//! Context registry.
//!
//! Mirrors the `contexts: Map<string, Map<ParamKey, ContextEntry>>`
//! shape in `frontends/mizan-base/src/index.ts`. Each entry holds the
//! latest `ContextState`, a `tokio::sync::watch::Sender` for notifying
//! subscribers, and a fetch function the registry invokes on demand.
//!
//! Subscribers receive a `ContextHandle` whose `rx: watch::Receiver`
//! they read from in their own loop. Watch channels overwrite the
//! previous value if the receiver hasn't consumed it yet — the render
//! loop sees only the latest state on each tick, never an intermediate
//! one. The TS kernel achieves the same effect via React's external
//! store re-render coalescing.

use std::collections::BTreeMap;
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;

use serde_json::Value;
use tokio::sync::{Mutex, RwLock, mpsc, watch};
use tokio_util::sync::CancellationToken;

use crate::error::MizanError;


#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ContextStatus {
    Idle,
    Loading,
    Success,
    Error,
}


#[derive(Debug, Clone)]
pub struct ContextState<T> {
    pub data: Option<T>,
    pub status: ContextStatus,
    pub error: Option<Arc<MizanError>>,
}


pub type ContextStateRaw = ContextState<Value>;


impl ContextStateRaw {
    pub fn idle() -> Self {
        Self { data: None, status: ContextStatus::Idle, error: None }
    }
}


pub type FetchFn = Arc<
    dyn Fn() -> Pin<Box<dyn Future<Output = Result<Value, MizanError>> + Send + 'static>>
        + Send
        + Sync,
>;


struct ContextEntry {
    params: Value,
    tx: watch::Sender<ContextStateRaw>,
    fetch_fn: FetchFn,
    refetch_tx: mpsc::UnboundedSender<()>,
    /// Cancel signal for the entry's spawned refetch loop. Set when the
    /// last handle on the entry unregisters.
    cancel: CancellationToken,
}


pub struct ContextRegistry {
    /// Outer key: context name. Inner key: `stable_key(params)`.
    entries: RwLock<HashMap<String, HashMap<String, Arc<Mutex<ContextEntry>>>>>,
}


impl Default for ContextRegistry {
    fn default() -> Self {
        Self::new()
    }
}


impl ContextRegistry {
    pub fn new() -> Self {
        Self { entries: RwLock::new(HashMap::new()) }
    }

    /// Register an instance of `(context_name, params)`. Idempotent —
    /// re-registering the same key returns a handle on the existing
    /// entry (the fetch_fn closure is replaced so the latest binding
    /// wins).
    pub async fn register(
        self: &Arc<Self>,
        name: impl Into<String>,
        params: Value,
        fetch_fn: FetchFn,
        initial_data: Option<Value>,
    ) -> ContextHandle {
        let name = name.into();
        let key = stable_key(&params);

        let mut outer = self.entries.write().await;
        let inner = outer.entry(name.clone()).or_default();

        if let Some(existing) = inner.get(&key).cloned() {
            // Update the fetch closure so the latest registration's
            // closure wins (matches the TS Strict-Mode behavior).
            {
                let mut entry = existing.lock().await;
                entry.fetch_fn = fetch_fn;
            }
            let entry = existing.lock().await;
            return ContextHandle {
                rx: entry.tx.subscribe(),
                refetch_tx: entry.refetch_tx.clone(),
                cancel: entry.cancel.clone(),
                registry: Arc::clone(self),
                name,
                key,
            };
        }

        let initial = match initial_data {
            Some(data) => ContextState { data: Some(data), status: ContextStatus::Success, error: None },
            None => ContextStateRaw::idle(),
        };
        let (tx, _rx) = watch::channel(initial);
        let (refetch_tx, mut refetch_rx) = mpsc::unbounded_channel::<()>();
        let cancel = CancellationToken::new();

        let entry = Arc::new(Mutex::new(ContextEntry {
            params: params.clone(),
            tx: tx.clone(),
            fetch_fn: fetch_fn.clone(),
            refetch_tx: refetch_tx.clone(),
            cancel: cancel.clone(),
        }));
        inner.insert(key.clone(), Arc::clone(&entry));
        drop(outer);

        // Spawn the entry's refetch loop. The loop owns its own fetch
        // closure handle resolution via the entry mutex — each tick
        // reads the latest closure, so updates via re-register apply.
        let entry_for_task = Arc::clone(&entry);
        let cancel_for_task = cancel.clone();
        tokio::spawn(async move {
            loop {
                tokio::select! {
                    _ = cancel_for_task.cancelled() => break,
                    msg = refetch_rx.recv() => {
                        if msg.is_none() { break; }
                        let (fetch_fn, tx) = {
                            let entry = entry_for_task.lock().await;
                            (entry.fetch_fn.clone(), entry.tx.clone())
                        };
                        // Loading state
                        let cur = tx.borrow().clone();
                        let loading = ContextState { data: cur.data, status: ContextStatus::Loading, error: None };
                        let _ = tx.send(loading);
                        // Drive the fetch
                        match fetch_fn().await {
                            Ok(data) => {
                                let _ = tx.send(ContextState { data: Some(data), status: ContextStatus::Success, error: None });
                            }
                            Err(err) => {
                                let cur = tx.borrow().clone();
                                let _ = tx.send(ContextState {
                                    data: cur.data,
                                    status: ContextStatus::Error,
                                    error: Some(Arc::new(err)),
                                });
                            }
                        }
                    }
                }
            }
        });

        ContextHandle {
            rx: tx.subscribe(),
            refetch_tx,
            cancel,
            registry: Arc::clone(self),
            name,
            key,
        }
    }

    /// Merge a value into a context entry's bundle slot. Mirrors the
    /// TS kernel `merge(context, params, slot, value)` call.
    pub async fn merge(
        &self,
        name: &str,
        params: Option<&Value>,
        slot: &str,
        value: &Value,
    ) {
        let key = match params {
            Some(p) => stable_key(p),
            None => stable_key(&Value::Object(Default::default())),
        };
        let entry_handle = {
            let outer = self.entries.read().await;
            outer.get(name).and_then(|inner| inner.get(&key)).cloned()
        };
        let Some(entry_arc) = entry_handle else { return };
        let entry = entry_arc.lock().await;
        let cur = entry.tx.borrow().clone();
        let Some(bundle) = cur.data.as_ref() else { return };
        let Some(merged) = crate::merge::merge_into_bundle(bundle, slot, value) else { return };
        let _ = entry.tx.send(ContextState {
            data: Some(merged),
            status: ContextStatus::Success,
            error: None,
        });
    }

    /// Trigger refetch on every entry of `name`.
    pub async fn invalidate_broad(&self, name: &str) {
        let entries = {
            let outer = self.entries.read().await;
            outer.get(name).map(|inner| inner.values().cloned().collect::<Vec<_>>())
        };
        let Some(entries) = entries else { return };
        for entry in entries {
            let tx = {
                let e = entry.lock().await;
                e.refetch_tx.clone()
            };
            let _ = tx.send(());
        }
    }

    /// Trigger refetch on the single entry matching `(name, params)`.
    pub async fn invalidate_scoped(&self, name: &str, params: &Value) {
        let key = stable_key(params);
        let entry_arc = {
            let outer = self.entries.read().await;
            outer.get(name).and_then(|inner| inner.get(&key)).cloned()
        };
        let Some(entry_arc) = entry_arc else { return };
        let tx = {
            let entry = entry_arc.lock().await;
            entry.refetch_tx.clone()
        };
        let _ = tx.send(());
    }

    async fn unregister(&self, name: &str, key: &str) {
        let mut outer = self.entries.write().await;
        if let Some(inner) = outer.get_mut(name) {
            if let Some(entry) = inner.remove(key) {
                let entry = entry.lock().await;
                entry.cancel.cancel();
            }
            if inner.is_empty() {
                outer.remove(name);
            }
        }
    }
}


pub struct ContextHandle {
    pub rx: watch::Receiver<ContextStateRaw>,
    refetch_tx: mpsc::UnboundedSender<()>,
    cancel: CancellationToken,
    registry: Arc<ContextRegistry>,
    name: String,
    key: String,
}


impl ContextHandle {
    /// Drive a refetch. Returns immediately; the new state lands on
    /// `rx` once the kernel's refetch task finishes the fetch.
    pub fn refetch(&self) {
        let _ = self.refetch_tx.send(());
    }

    pub fn state(&self) -> ContextStateRaw {
        self.rx.borrow().clone()
    }

    pub fn cancel_token(&self) -> CancellationToken {
        self.cancel.clone()
    }

    pub async fn unregister(self) {
        self.registry.unregister(&self.name, &self.key).await;
    }
}


/// Byte-identical to TS `JSON.stringify(params, Object.keys(params).sort())`.
///
/// Uses `BTreeMap` for deterministic key ordering and serializes via
/// `serde_json::to_string` (compact, no whitespace) — matches the TS
/// default. Non-object / non-string params (numbers, booleans) pass
/// through serde_json's standard JSON representation.
pub fn stable_key(params: &Value) -> String {
    match params {
        Value::Object(map) => {
            let sorted: BTreeMap<&String, &Value> = map.iter().collect();
            serde_json::to_string(&sorted).unwrap_or_default()
        }
        other => serde_json::to_string(other).unwrap_or_default(),
    }
}


#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    #[test]
    fn stable_key_sorts_object_keys() {
        let a = stable_key(&json!({"b": 1, "a": 2}));
        let b = stable_key(&json!({"a": 2, "b": 1}));
        assert_eq!(a, b);
        assert_eq!(a, r#"{"a":2,"b":1}"#);
    }

    #[test]
    fn stable_key_handles_empty_object() {
        assert_eq!(stable_key(&json!({})), "{}");
    }

    #[tokio::test]
    async fn register_and_refetch() {
        let registry = Arc::new(ContextRegistry::new());
        let counter = Arc::new(std::sync::atomic::AtomicU32::new(0));
        let counter_clone = Arc::clone(&counter);
        let fetch_fn: FetchFn = Arc::new(move || {
            let counter = Arc::clone(&counter_clone);
            Box::pin(async move {
                let n = counter.fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1;
                Ok(json!({ "count": n }))
            })
        });

        let mut handle = registry.register("test", json!({}), fetch_fn, None).await;
        handle.refetch();
        // Poll until success — watch::Receiver::changed() returns once
        // per "newest value seen" advance, so back-to-back sends from the
        // refetch task can coalesce into a single notification. The loop
        // ignores intermediate Loading states and waits for Success.
        loop {
            tokio::time::timeout(std::time::Duration::from_secs(2), handle.rx.changed())
                .await
                .expect("changed timed out")
                .unwrap();
            if handle.state().status == ContextStatus::Success {
                break;
            }
        }
        let state = handle.state();
        assert_eq!(state.data.unwrap()["count"], 1);
    }
}
