qbice 0.6.5

use std::{collections::VecDeque, sync::Arc};

use qbice_stable_hash::{Compact128, StableHash};
use tokio::{sync::RwLock, task::JoinSet};

use crate::{
    Engine, Query, TrackedEngine,
    config::{Config, WriteTransaction},
    engine::{
        computation_graph::{
            ActiveInputSessionGuard,
            caller::{CallerInformation, CallerKind, QueryCaller},
        },
        guard::GuardExt,
    },
    query::QueryID,
};

/// A transactional session for setting and updating input query values.
///
/// `InputSession` provides a batch interface for modifying input queries with
/// automatic dirty propagation. All changes are buffered during the session
/// lifetime and committed atomically when the session is dropped.
///
/// # Lifecycle
///
/// 1. **Creation**: Obtained via [`Engine::input_session()`]
///    - Acquires a write lock on the computation graph
///    - Increments the global timestamp to invalidate running queries
/// 2. **Modification**: Use [`set_input`](Self::set_input) or
///    [`refresh`](Self::refresh)
///    - Changes are buffered in memory
///    - Dirty queries are tracked for later propagation
/// 3. **Commit**: Automatically triggered on drop
///    - Dirty propagation executes in parallel via the Rayon thread pool
///    - Write buffer is submitted to persistent storage
///    - Statistics are reset and internal state is cleaned up
///
/// # Dirty Propagation
///
/// When the session ends, the engine:
/// - Compares new values with stored values via fingerprint hashing
/// - Marks queries whose values changed as dirty
/// - Recursively marks all dependent queries as dirty
/// - Ensures downstream queries will recompute on next access
///
/// # Concurrency
///
/// **Only one input session can be active at a time.** Creating a new session
/// while another exists will deadlock, as the write lock cannot be acquired.
///
/// The session increments the timestamp on creation, causing any in-flight
/// queries to become stale and await cancellation. This ensures consistency
/// between the input changes and query results.
pub struct InputSession<C: Config> {
    engine: Arc<Engine<C>>,
    dirty_batch: Arc<RwLock<VecDeque<QueryID>>>,
    comitted: bool,

    #[allow(clippy::type_complexity)]
    transaction:
        Arc<RwLock<Option<(WriteTransaction<C>, ActiveInputSessionGuard)>>>,
}

impl<C: Config> Drop for InputSession<C> {
    // NOTE: the commit needs to always happen, even if the user forgets to call
    // `commit()`. Thus, we perform the commit in the `Drop` impl.
    fn drop(&mut self) {
        if self.comitted {
            return;
        }

        let transaction = self.transaction.clone();
        let dirty_batch = self.dirty_batch.clone();
        let engine = self.engine.clone();

        tokio::spawn(async move {
            let Some((transaction, guard)) = transaction.write().await.take()
            else {
                // the transaction has already been committed
                return;
            };

            let dirty_batch = std::mem::take(&mut *dirty_batch.write().await);

            Self::commit_internal(engine, dirty_batch, transaction).await;

            drop(guard);
        });
    }
}

impl<C: Config> InputSession<C> {
    /// Commits the input session, performing dirty propagation and submitting
    /// the write buffer.
    pub async fn commit(mut self) {
        let engine = self.engine.clone();

        async move {
            self.comitted = true;

            let dirty_batch =
                std::mem::take(&mut *self.dirty_batch.write().await);
            let (transaction, guard) =
                self.transaction.write().await.take().unwrap();

            Self::commit_internal(engine, dirty_batch, transaction).await;

            drop(guard);
        }
        .guarded()
        .await;
    }

    async fn commit_internal(
        engine: Arc<Engine<C>>,
        dirty_batch: VecDeque<QueryID>,
        mut transaction: WriteTransaction<C>,
    ) {
        engine.computation_graph.reset_statistic();
        engine.clear_dirtied_queries();

        transaction = engine
            .dirty_propagate_from_batch(dirty_batch.into_iter(), transaction)
            .await;

        engine.submit_write_buffer(transaction);
    }
}

impl<C: Config> std::fmt::Debug for InputSession<C> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("InputSession")
            .field("engine", &self.engine)
            .field("dirty_batch", &self.dirty_batch)
            .finish_non_exhaustive()
    }
}

impl<C: Config> Engine<C> {
    /// Creates an input session for setting or updating input query values.
    ///
    /// An input session provides a transactional interface for modifying
    /// input values. All changes are batched and committed atomically when
    /// the session is dropped.
    ///
    /// # Lifecycle
    ///
    /// 1. **Create**: Call this method to begin a session
    /// 2. **Modify**: Use [`set_input`](InputSession::set_input) to update
    ///    values
    /// 3. **Commit**: Drop the session to apply changes and propagate dirtiness
    ///
    /// # Dirty Propagation
    ///
    /// When the session is dropped, the engine:
    /// - Compares new values with existing ones via fingerprints
    /// - Marks changed queries and their dependents as dirty
    /// - Increments the global timestamp (if any changes occurred)
    ///
    /// # Cancellation and Timestamp Management
    ///
    /// The input session increments the global timestamp when it's created. If
    /// there're any other running queries, they will **stuck** in the pending
    /// loop forever. This signifies that the old running queries are no longer
    /// valid and must be dropped.
    ///
    /// # Deadlocks
    ///
    /// There can only be one active input session at a time. Attempting to
    /// create a new session while another is active will result in a
    /// deadlock.
    #[must_use]
    #[allow(clippy::unused_async)]
    pub async fn input_session(self: &Arc<Self>) -> InputSession<C> {
        // acquire a write lock on the write buffer and increment timestamp
        let (write_buffer_with_lock, active_input_session_guard) =
            self.acquire_active_input_session_guard().await;

        InputSession {
            dirty_batch: Arc::new(RwLock::new(VecDeque::new())),
            engine: self.clone(),
            comitted: false,
            transaction: Arc::new(RwLock::new(Some((
                write_buffer_with_lock,
                active_input_session_guard,
            )))),
        }
    }
}

/// The result of setting an input query value.
///
/// This enum indicates whether an input query was newly created, had its value
/// updated, or remained unchanged after a
/// [`set_input`](InputSession::set_input) operation.
///
/// The variant determines whether dirty propagation will occur when the input
/// session is dropped.
#[derive(
    Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, StableHash,
)]
pub enum SetInputResult {
    /// The input query was set for the first time.
    ///
    /// This indicates that no prior value existed for this query. A new entry
    /// was created in the computation graph. No dirty propagation is needed
    /// since there are no dependents yet.
    Fresh,

    /// The input query's value was updated.
    ///
    /// The new value's fingerprint differs from the previously stored value,
    /// indicating a semantic change. This query and all its dependents will be
    /// marked dirty for recomputation when the input session is dropped.
    Updated,

    /// The input query's value remained unchanged.
    ///
    /// The new value's fingerprint matches the previously stored value,
    /// indicating no semantic change occurred. No dirty propagation will be
    /// performed, and dependents remain valid.
    Unchanged,
}

impl<C: Config> InputSession<C> {
    /// Sets the value for an input query.
    ///
    /// This method updates the value associated with the given query. If the
    /// new value differs from the existing value (based on fingerprint
    /// comparison), the query and its dependents will be marked as dirty for
    /// recomputation.
    ///
    /// # Behavior
    ///
    /// - **New queries**: If the query has never been set before, a new entry
    ///   is created
    /// - **Changed values**: If the value fingerprint differs from the stored
    ///   value, dirty propagation is scheduled
    /// - **Timestamp management**: The first change in a session increments the
    ///   global timestamp
    ///
    /// All dirty propagation happens when the `InputSession` is dropped, not
    /// when this method is called.
    ///
    /// # Type Parameters
    ///
    /// - `Q`: The query type, must implement [`Query`]
    ///
    /// # Arguments
    ///
    /// - `query`: The input query key
    /// - `new_value`: The new value to associate with this query
    pub async fn set_input<Q: Query>(
        &mut self,
        query: Q,
        new_value: Q::Value,
    ) -> SetInputResult {
        let query_hash = self.engine.hash(&query);
        let query_id = QueryID::new::<Q>(query_hash);

        let mut snapshot = self.engine.get_exclusive_snapshot(query_hash).await;

        let query_value_fingerprint = self.engine.hash(&new_value);

        // has prior node infos, check for fingerprint diff
        // also, unwire the backward edges (if any)
        let set_input_result = (snapshot.node_info().await).map_or(
            SetInputResult::Fresh,
            |node_info| {
                let fingerprint_diff =
                    node_info.value_fingerprint() != query_value_fingerprint;

                // only dirty propagate if the fingerprint differs
                if fingerprint_diff {
                    SetInputResult::Updated
                } else {
                    SetInputResult::Unchanged
                }
            },
        );

        // should be safe since we're holding input session phase guard
        let ts = unsafe { self.engine.get_current_timestamp_unchecked() };

        let transaction = self.transaction.clone();
        let dirty_batch = self.dirty_batch.clone();

        async move {
            if set_input_result == SetInputResult::Updated {
                dirty_batch.write().await.push_back(query_id);
            }

            let mut transaction = transaction.write().await;

            let Some((write_buffer, _guard)) = transaction.as_mut() else {
                panic!("InputSession transaction has already been committed");
            };

            snapshot
                .set_computed_input(
                    query,
                    query_hash,
                    new_value,
                    query_value_fingerprint,
                    write_buffer,
                    true,
                    ts,
                )
                .await;

            set_input_result
        }
        .guarded()
        .await
    }

    /// Inspects the current value and performs an update based on the provided
    /// function.
    ///
    /// This method allows you to read the existing value of an input query,
    /// compute a new value based on it, and update the query atomically. The
    /// update function receives an `Option<Q::Value>`, which is `None` if the
    /// query has never been set before, or `Some(value)` if it has an existing
    /// value.
    pub async fn update<Q: Query>(
        &mut self,
        query: Q,
        update_fn: impl FnOnce(Option<Q::Value>) -> Q::Value,
    ) -> SetInputResult {
        let query_hash = self.engine.hash(&query);
        let query_id = QueryID::new::<Q>(query_hash);

        let mut snapshot = self.engine.get_exclusive_snapshot(query_hash).await;

        let current_value = snapshot.query_result().await;
        let new_value = update_fn(current_value);

        let query_value_fingerprint = self.engine.hash(&new_value);

        // has prior node infos, check for fingerprint diff
        // also, unwire the backward edges (if any)
        let set_input_result = (snapshot.node_info().await).map_or(
            SetInputResult::Fresh,
            |node_info| {
                let fingerprint_diff =
                    node_info.value_fingerprint() != query_value_fingerprint;

                // only dirty propagate if the fingerprint differs
                if fingerprint_diff {
                    SetInputResult::Updated
                } else {
                    SetInputResult::Unchanged
                }
            },
        );

        // should be safe since we're holding input session phase guard
        let ts = unsafe { self.engine.get_current_timestamp_unchecked() };

        let transaction = self.transaction.clone();
        let dirty_batch = self.dirty_batch.clone();

        async move {
            if set_input_result == SetInputResult::Updated {
                dirty_batch.write().await.push_back(query_id);
            }

            let mut transaction = transaction.write().await;

            let Some((write_buffer, _guard)) = transaction.as_mut() else {
                panic!("InputSession transaction has already been committed");
            };

            snapshot
                .set_computed_input(
                    query,
                    query_hash,
                    new_value,
                    query_value_fingerprint,
                    write_buffer,
                    true,
                    ts,
                )
                .await;

            set_input_result
        }
        .guarded()
        .await
    }

    /// Refreshes all external input queries of type `Q`.
    ///
    /// This method re-executes all queries of type `Q` that were previously
    /// computed with [`ExecutionStyle::ExternalInput`]. For each query:
    ///
    /// 1. The executor is re-invoked to fetch the latest external data
    /// 2. The new result is compared with the old result via fingerprints
    /// 3. Only if the result changed, the query is marked dirty for propagation
    ///
    /// # Use Case
    ///
    /// External input queries represent data from the outside world (files,
    /// network, databases, etc.). When you know the external data has changed,
    /// call this method to refresh and update all queries of that type.
    ///
    /// # Type Parameters
    ///
    /// - `Q`: The query type to refresh. Must implement [`Query`].
    ///
    /// # Cancellation
    ///
    /// This method can be cancelled. However, if cancelled midway, not all
    /// queries of type `Q` may be refreshed, potentially leaving some stale
    /// data in the computation graph. Nevertheless, it's guaranteed that the
    /// engine's internal state will remain consistent.
    #[allow(clippy::too_many_lines)]
    pub async fn refresh<Q: Query>(&mut self) {
        struct RefreshResult<K, V> {
            query_input_hash: Compact128,
            query_result_hash: Compact128,
            query_input: K,
            new_value: V,
        }

        let type_id = Q::STABLE_TYPE_ID;

        // Get all query hashes for this type that were computed as
        // ExternalInput
        let external_input_set =
            self.engine.get_external_input_queries(&type_id).await;

        // should be safe since we're holding input session phase guard
        let timestamp =
            unsafe { self.engine.get_current_timestamp_unchecked() };

        let hashes = external_input_set.collect::<Vec<_>>();

        let expected_parallelism = std::thread::available_parallelism()
            .map_or_else(|_| 1, std::num::NonZero::get)
            * 4;
        let chunk_size =
            std::cmp::max((hashes.len()) / expected_parallelism, 1);

        let mut join_set = JoinSet::new();

        for chunk in hashes.chunks(chunk_size) {
            let engine = self.engine.clone();
            let chunk = chunk.to_owned();

            join_set.spawn(async move {
                let mut results = Vec::with_capacity(chunk.len());

                for query_hash in chunk.iter().copied() {
                    let query_id = QueryID::new::<Q>(query_hash);

                    let query = engine.get_query_input::<Q>(&query_hash).await;

                    // Create a tracked engine to execute the query
                    let wait_group = waitgroup::WaitGroup::new();

                    let tracked_engine = TrackedEngine::new(
                        engine.clone(),
                        CallerInformation::new(
                            CallerKind::Query(QueryCaller::new_external_input(
                                query_id,
                                wait_group.worker(),
                            )),
                            timestamp,
                            None,
                        ),
                    );

                    // Invoke the executor to get the new value
                    let entry =
                        engine.executor_registry.get_executor_entry::<Q>();
                    let result = entry
                        .invoke_executor::<Q>(&query, &tracked_engine)
                        .await;

                    // Wait for any spawned tasks to complete
                    drop(tracked_engine);
                    wait_group.wait().await;

                    let new_value = match result {
                        Ok(value) => value,
                        Err(panic) => panic.resume_unwind(),
                    };

                    let new_fingerprint = engine.hash(&new_value);

                    results.push(RefreshResult {
                        query_input: query,
                        query_input_hash: query_hash,
                        query_result_hash: new_fingerprint,
                        new_value,
                    });
                }

                results
            });
        }

        let dirty_batch = self.dirty_batch.clone();
        let engine = self.engine.clone();
        let transaction = self.transaction.clone();

        async move {
            let mut dirty_batch = dirty_batch.write().await;
            let mut transaction = transaction.write().await;

            let Some((transaction, _guard)) = transaction.as_mut() else {
                panic!("InputSession transaction has already been committed");
            };

            while let Some(res) = join_set.join_next().await {
                match res {
                    Ok(results) => {
                        for refresh_result in results {
                            let mut snapshot = engine
                                .get_exclusive_snapshot::<Q>(
                                    refresh_result.query_input_hash,
                                )
                                .await;

                            let fingerprint_diff =
                                snapshot.value_fingerprint().await.unwrap()
                                    != refresh_result.query_result_hash;

                            if fingerprint_diff {
                                let query_id = QueryID::new::<Q>(
                                    refresh_result.query_input_hash,
                                );

                                dirty_batch.push_back(query_id);
                            }

                            snapshot
                                .set_computed_input(
                                    refresh_result.query_input,
                                    refresh_result.query_input_hash,
                                    refresh_result.new_value,
                                    refresh_result.query_result_hash,
                                    transaction,
                                    false,
                                    timestamp,
                                )
                                .await;
                        }
                    }

                    Err(er) => match er.try_into_panic() {
                        Ok(panic_reason) => {
                            std::panic::resume_unwind(panic_reason);
                        }
                        Err(er) => {
                            panic!(
                                "Failed to refresh external input query: {er}"
                            );
                        }
                    },
                }
            }
        }
        .guarded()
        .await;
    }

    /// Interns a value, returning a reference-counted handle to the shared
    /// allocation.
    ///
    /// This is a delegation to [`Engine::intern`]. See its documentation for
    /// more details.
    pub fn intern<
        T: StableHash + crate::Identifiable + Send + Sync + 'static,
    >(
        &self,
        value: T,
    ) -> qbice_storage::intern::Interned<T> {
        self.engine.intern(value)
    }

    /// Interns an unsized value, returning a reference-counted handle to the
    /// shared allocation.
    ///
    /// This is a delegation to [`Engine::intern_unsized`]. See its
    /// documentation for more details.
    pub fn intern_unsized<
        T: StableHash + crate::Identifiable + Send + Sync + 'static + ?Sized,
        Q: std::borrow::Borrow<T> + Send + Sync + 'static,
    >(
        &self,
        value: Q,
    ) -> qbice_storage::intern::Interned<T>
    where
        Arc<T>: From<Q>,
    {
        self.engine.intern_unsized(value)
    }
}