surrealdb 3.3.1

use std::borrow::Cow;
use std::future::IntoFuture;
use std::marker::PhantomData;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use async_channel::Receiver;
use futures::StreamExt;
use surrealdb_types::{ConfigurationError, QueryError, SerializationError, ValidationError};
#[cfg(not(target_family = "wasm"))]
pub(crate) use tokio::spawn;
use uuid::Uuid;
#[cfg(target_family = "wasm")]
pub(crate) use wasm_bindgen_futures::spawn_local as spawn;

use crate::conn::{Router, ctx, ctx_txn};
use crate::engine::any::Any;
use crate::method::{BoxFuture, Live, OnceLockExt, Select};
use crate::notification::Notification;
use crate::opt::Resource;
use crate::types::{
	Action, Notification as CoreNotification, RecordId, SurrealValue, Value, Variables,
};
use crate::{Connection, Error, ExtraFeatures, Result, Surreal};

fn into_future<C, O>(this: Select<C, O, Live>) -> BoxFuture<Result<Stream<O>>>
where
	C: Connection,
{
	let Select {
		client,
		resource,
		..
	} = this;
	Box::pin(async move {
		let router = client.inner.router.extract()?;
		if !router.features.contains(&ExtraFeatures::LiveQueries) {
			return Err(Error::configuration(
				"The protocol or storage engine does not support live queries on this architecture"
					.to_string(),
				ConfigurationError::LiveQueryNotSupported,
			));
		}

		let what_resource = resource?;

		let mut variables = Variables::new();
		let what = what_resource.for_sql_query(&mut variables)?;

		// Generate the LIVE SELECT SQL based on resource type
		let query = match what_resource {
			Resource::Table(table) => {
				variables.insert("_table".to_string(), Value::Table(table));
				format!("LIVE SELECT * FROM {what}")
			}
			Resource::RecordId(record) => {
				// For a specific record, we need to query the table with a WHERE clause
				// because LIVE queries don't support record IDs directly
				variables.insert("_table".to_string(), Value::Table(record.table.clone()));
				variables.insert("_record_id".to_string(), Value::RecordId(record));
				"LIVE SELECT * FROM $_table WHERE id = $_record_id".to_string()
			}
			Resource::Object(_) => {
				return Err(Error::validation(
					"Live queries on objects not supported".to_string(),
					ValidationError::InvalidParams,
				));
			}
			Resource::Array(_) => {
				return Err(Error::validation(
					"Live queries on arrays not supported".to_string(),
					ValidationError::InvalidParams,
				));
			}
			Resource::Range(query_range) => {
				// For live queries with ranges, we can't use the range in FROM clause
				// We need to use the table and add WHERE conditions
				variables.insert("_table".to_string(), Value::Table(query_range.table.clone()));
				let table_expr = "$_table";

				// Build WHERE clause for range queries
				let mut conditions = Vec::new();

				// Handle start bound
				match &query_range.range.start {
					std::ops::Bound::Included(key) => {
						variables.insert(
							"_start".to_string(),
							Value::RecordId(RecordId::new(query_range.table.clone(), key.clone())),
						);
						conditions.push("id >= $_start");
					}
					std::ops::Bound::Excluded(key) => {
						variables.insert(
							"_start".to_string(),
							Value::RecordId(RecordId::new(query_range.table.clone(), key.clone())),
						);
						conditions.push("id > $_start");
					}
					std::ops::Bound::Unbounded => {}
				}

				// Handle end bound
				match &query_range.range.end {
					std::ops::Bound::Included(key) => {
						variables.insert(
							"_end".to_string(),
							Value::RecordId(RecordId::new(query_range.table.clone(), key.clone())),
						);
						conditions.push("id <= $_end");
					}
					std::ops::Bound::Excluded(key) => {
						variables.insert(
							"_end".to_string(),
							Value::RecordId(RecordId::new(query_range.table.clone(), key.clone())),
						);
						conditions.push("id < $_end");
					}
					std::ops::Bound::Unbounded => {}
				}

				// Build final query
				if conditions.is_empty() {
					format!("LIVE SELECT * FROM {table_expr}")
				} else {
					format!("LIVE SELECT * FROM {table_expr} WHERE {}", conditions.join(" AND "))
				}
			}
		};

		// Execute the LIVE SELECT query directly to get the UUID
		let results = router
			.query_results(ctx_txn(client.session_id, None), Cow::Owned(query), variables)
			.await?;

		// Get the first result which should be the UUID
		let result = results.into_iter().next().ok_or_else(|| {
			Error::query("LIVE query returned no results".to_string(), QueryError::NotExecuted)
		})?;

		let id = match result.result? {
			Value::Uuid(id) => *id,
			Value::Array(mut arr) if arr.len() == 1 => match arr.pop() {
				Some(Value::Uuid(id)) => *id,
				_ => {
					return Err(Error::internal(
						"successful live query didn't return a uuid".to_string(),
					));
				}
			},
			other => {
				return Err(Error::internal(format!(
					"successful live query didn't return a uuid, got: {:?}",
					other
				)));
			}
		};

		let rx = register(router, id, client.session_id).await?;
		Ok(Stream::new(Arc::clone(&client.inner).into(), id, Some(rx)))
	})
}

pub(crate) async fn register(
	router: &Router,
	id: Uuid,
	session_id: Uuid,
) -> Result<Receiver<Result<CoreNotification>>> {
	let (tx, rx) = async_channel::unbounded();
	router.engine.subscribe_live(ctx(session_id), id, tx).await?;
	Ok(rx)
}

impl<'r, Client> IntoFuture for Select<'r, Client, Value, Live>
where
	Client: Connection,
{
	type Output = Result<Stream<Value>>;
	type IntoFuture = BoxFuture<'r, Self::Output>;

	fn into_future(self) -> Self::IntoFuture {
		into_future(self)
	}
}

impl<'r, Client, R> IntoFuture for Select<'r, Client, Option<R>, Live>
where
	Client: Connection,
	R: SurrealValue,
{
	type Output = Result<Stream<Option<R>>>;
	type IntoFuture = BoxFuture<'r, Self::Output>;

	fn into_future(self) -> Self::IntoFuture {
		into_future(self)
	}
}

impl<'r, Client, R> IntoFuture for Select<'r, Client, Vec<R>, Live>
where
	Client: Connection,
	R: SurrealValue,
{
	type Output = Result<Stream<Vec<R>>>;
	type IntoFuture = BoxFuture<'r, Self::Output>;

	fn into_future(self) -> Self::IntoFuture {
		into_future(self)
	}
}

/// Notifications from [`Select::live`](crate::method::Select::live) (built via
/// [`Surreal::select`](crate::Surreal::select)) or from live handles inside
/// [`Surreal::query`](crate::Surreal::query).
#[derive(Debug)]
#[must_use = "streams do nothing unless you poll them"]
pub struct Stream<R> {
	pub(crate) client: Surreal<Any>,
	// We no longer need the lifetime and the type parameter
	// Leaving them in for backwards compatibility
	pub(crate) id: Uuid,
	pub(crate) rx: Option<Pin<Box<Receiver<Result<CoreNotification>>>>>,
	pub(crate) response_type: PhantomData<R>,
}

impl<R> Stream<R> {
	pub(crate) fn new(
		client: Surreal<Any>,
		id: Uuid,
		rx: Option<Receiver<Result<CoreNotification>>>,
	) -> Self {
		Self {
			id,
			rx: rx.map(Box::pin),
			client,
			response_type: PhantomData,
		}
	}
}

macro_rules! poll_next {
	($result:ident => $body:expr_2021) => {
		fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
			let Some(ref mut rx) = self.as_mut().rx else {
				return Poll::Ready(None);
			};
			let polled = rx.poll_next_unpin(cx);
			match polled {
				Poll::Ready(Some($result)) => {
					// `Killed` is the subscription's own end: the datastore has
					// removed it, so the stream has nothing left to receive and
					// nothing left to tear down. Dropping the receiver here is
					// what tells `Drop` that, which would otherwise ask for a
					// subscription that is already gone to be killed.
					if let Ok(notification) = &$result
						&& matches!(notification.action, Action::Killed)
					{
						self.as_mut().rx = None;
					}
					$body
				}
				Poll::Ready(None) => Poll::Ready(None),
				Poll::Pending => Poll::Pending,
			}
		}
	};
}

impl futures::Stream for Stream<Value> {
	type Item = Result<Notification<Value>>;

	poll_next! {
		result => match result {
			Ok(notification) => {
				match notification.action {
					Action::Killed => Poll::Ready(None),
					action => Poll::Ready(Some(Ok(Notification {
						query_id: notification.id,
						action,
						data: notification.result,
					}))),
				}
			}
			Err(error) => {
				Poll::Ready(Some(Err(error)))
			}
		}
	}

	fn size_hint(&self) -> (usize, Option<usize>) {
		(0, None)
	}
}

macro_rules! poll_next_and_convert {
	() => {
		poll_next! {
			result => match result {
				Ok(notification) => {
					Poll::Ready(deserialize(notification))
				}
				Err(error) => {
					Poll::Ready(Some(Err(error)))
				}
			}
		}
	};
}

impl<R> futures::Stream for Stream<Option<R>>
where
	R: SurrealValue + Unpin,
{
	type Item = Result<Notification<R>>;

	poll_next_and_convert! {}
}

impl<R> futures::Stream for Stream<Vec<R>>
where
	R: SurrealValue + Unpin,
{
	type Item = Result<Notification<R>>;

	poll_next_and_convert! {}
}

impl<R> futures::Stream for Stream<Notification<R>>
where
	R: SurrealValue + Unpin,
{
	type Item = Result<Notification<R>>;

	poll_next_and_convert! {}
}

/// Ends the live query behind a [`Stream`], on the way out of its `Drop`.
///
/// A dropped stream is gone whether or not the engine could end its
/// subscription, so the failure is reported rather than returned: there is no
/// caller left to return it to, and a subscription that outlives its stream
/// goes on being registered -- and, on the engines that capture events per
/// table, goes on costing every later write.
pub(crate) fn kill<Client>(client: &Surreal<Client>, uuid: Uuid)
where
	Client: Connection,
{
	let client = client.clone();
	spawn(async move {
		if let Ok(router) = client.inner.router.extract()
			&& let Err(error) = router.engine.kill(ctx(client.session_id), uuid).await
		{
			tracing::warn!("Failed to kill live query '{uuid}'; {error}");
		}
	});
}

impl<R> Drop for Stream<R> {
	/// Close the live query stream
	///
	/// This kills the live query process responsible for this stream.
	fn drop(&mut self) {
		if self.rx.is_some() {
			kill(&self.client, self.id);
		}
	}
}

fn deserialize<R>(notification: CoreNotification) -> Option<Result<Notification<R>>>
where
	R: SurrealValue,
{
	let query_id = notification.id;
	let action = notification.action;
	match action {
		Action::Killed => None,
		action => match R::from_value(notification.result) {
			Ok(data) => Some(Ok(Notification {
				query_id,
				data,
				action,
			})),
			Err(error) => Some(Err(Error::serialization(
				error.to_string(),
				SerializationError::Deserialization,
			))),
		},
	}
}

#[cfg(test)]
mod tests {
	use futures::StreamExt;

	use super::*;

	/// A notification for `id`, as the engine delivers one.
	fn notification(id: Uuid, action: Action) -> CoreNotification {
		CoreNotification::new(id.into(), None, action, Value::None, Value::None)
	}

	/// A `Killed` notification ends the stream *and* everything behind it.
	///
	/// The subscription is gone from the datastore by the time this arrives --
	/// it is what the removal queues -- so a stream that has seen one must not
	/// ask for it to be killed again on the way out of `Drop`: there is nothing
	/// left to kill, and the attempt fails.
	#[tokio::test]
	async fn a_killed_notification_leaves_nothing_to_tear_down() {
		let id = Uuid::now_v7();
		let (sender, receiver) = async_channel::unbounded();
		let mut stream: Stream<Value> = Stream::new(Surreal::init(), id, Some(receiver));

		sender.send(Ok(notification(id, Action::Create))).await.unwrap();
		sender.send(Ok(notification(id, Action::Killed))).await.unwrap();

		assert!(stream.next().await.is_some(), "the change is delivered");
		assert!(stream.rx.is_some(), "the subscription is still this stream's to end");

		assert!(stream.next().await.is_none(), "a killed subscription ends the stream");
		assert!(
			stream.rx.is_none(),
			"a stream that has seen the subscription end has nothing left to tear down"
		);
	}
}