//! The reactive graph: a thread-local arena of nodes (signals, memos,
//! effects) with auto-tracked dependencies and mark-dirty /
//! evaluate-on-demand invalidation.
//!
//! Invalidation uses the three-state algorithm (Clean / Check / Dirty):
//! writing a signal marks direct subscribers Dirty and transitive ones Check;
//! when an effect or memo is pulled, a Check node first asks its sources to
//! update and only recomputes if one of them actually changed value. This
//! gives glitch-freedom on diamond dependencies without running anything
//! twice, and equality gating stops propagation where recomputed values are
//! unchanged.
//!
//! User closures (memo bodies, effect bodies) are moved out of the arena
//! while they run, so they can freely re-enter the runtime (read signals,
//! create nodes, set signals).

use std::any::Any;
use std::cell::RefCell;

use slotmap::{SlotMap, new_key_type};

new_key_type! {
    pub(crate) struct NodeKey;
    pub(crate) struct ScopeKey;
}

#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Debug)]
enum State {
    Clean,
    /// Possibly stale: a transitive source changed; sources must be checked
    /// before deciding whether to recompute.
    Check,
    /// Definitely stale: a direct source changed value.
    Dirty,
}

#[derive(Copy, Clone, PartialEq, Eq, Debug)]
pub(crate) enum Kind {
    Signal,
    Memo,
    Effect,
}

/// Body of a memo or effect. Memos compute their new value, write it into
/// the slot (`None` on the very first run), and report whether it changed.
/// Effects ignore the slot and the return value is meaningless.
pub(crate) type RunFn = Box<dyn FnMut(&mut Option<Box<dyn Any>>) -> bool>;

struct Node {
    kind: Kind,
    state: State,
    /// Present for signals and memos.
    value: Option<Box<dyn Any>>,
    /// Present for memos and effects; `None` while the body is executing.
    run: Option<RunFn>,
    sources: Vec<NodeKey>,
    subscribers: Vec<NodeKey>,
    /// Effects only: already in the flush queue.
    queued: bool,
    scope: ScopeKey,
}

#[derive(Default)]
struct Scope {
    nodes: Vec<NodeKey>,
    children: Vec<ScopeKey>,
    parent: Option<ScopeKey>,
}

pub(crate) struct Runtime {
    nodes: SlotMap<NodeKey, Node>,
    scopes: SlotMap<ScopeKey, Scope>,
    root_scope: ScopeKey,
    current_scope: ScopeKey,
    /// The node currently executing, for dependency auto-tracking.
    observer: Option<NodeKey>,
    /// Effects waiting for the next flush.
    queue: Vec<NodeKey>,
    flushing: bool,
}

thread_local! {
    static RUNTIME: RefCell<Runtime> = RefCell::new(Runtime::new());
}

/// Rounds of effect execution per flush before assuming effects are feeding
/// each other in a cycle.
const MAX_FLUSH_ROUNDS: usize = 1000;

impl Runtime {
    fn new() -> Self {
        let mut scopes = SlotMap::with_key();
        let root_scope = scopes.insert(Scope::default());
        Self {
            nodes: SlotMap::with_key(),
            scopes,
            root_scope,
            current_scope: root_scope,
            observer: None,
            queue: Vec::new(),
            flushing: false,
        }
    }

    fn insert_node(
        &mut self,
        kind: Kind,
        value: Option<Box<dyn Any>>,
        run: Option<RunFn>,
    ) -> NodeKey {
        let scope = self.current_scope;
        let key = self.nodes.insert(Node {
            kind,
            state: State::Clean,
            value,
            run,
            sources: Vec::new(),
            subscribers: Vec::new(),
            queued: false,
            scope,
        });
        self.scopes[scope].nodes.push(key);
        key
    }

    /// Record that the running observer depends on `key`.
    fn track(&mut self, key: NodeKey) {
        let Some(observer) = self.observer else {
            return;
        };
        if observer == key {
            return;
        }
        let Some(node) = self.nodes.get_mut(key) else {
            return;
        };
        if !node.subscribers.contains(&observer) {
            node.subscribers.push(observer);
            if let Some(obs) = self.nodes.get_mut(observer) {
                obs.sources.push(key);
            }
        }
    }

    /// Mark `key` at least `level` stale, propagating Check to transitive
    /// subscribers and queueing any effects encountered.
    fn stale(&mut self, key: NodeKey, level: State) {
        let Some(node) = self.nodes.get_mut(key) else {
            return;
        };
        if node.state >= level {
            return;
        }
        let was_clean = node.state == State::Clean;
        node.state = level;
        if node.kind == Kind::Effect && !node.queued {
            node.queued = true;
            self.queue.push(key);
        }
        if was_clean {
            let subscribers = node.subscribers.clone();
            for sub in subscribers {
                self.stale(sub, State::Check);
            }
        }
    }

    /// Detach `key` from everything it reads, ahead of a re-run.
    fn clear_sources(&mut self, key: NodeKey) {
        let sources = match self.nodes.get_mut(key) {
            Some(node) => std::mem::take(&mut node.sources),
            None => return,
        };
        for source in sources {
            if let Some(node) = self.nodes.get_mut(source) {
                node.subscribers.retain(|s| *s != key);
            }
        }
    }

    fn remove_node(&mut self, key: NodeKey) {
        self.clear_sources(key);
        let Some(node) = self.nodes.remove(key) else {
            return;
        };
        // Dependents lose this source but stay functional; their next re-run
        // simply won't re-establish the edge.
        for sub in node.subscribers {
            if let Some(sub_node) = self.nodes.get_mut(sub) {
                sub_node.sources.retain(|s| *s != key);
            }
        }
    }

    fn dispose_scope(&mut self, scope: ScopeKey) {
        let Some(data) = self.scopes.remove(scope) else {
            return;
        };
        for child in data.children {
            self.dispose_scope(child);
        }
        for node in data.nodes {
            self.remove_node(node);
        }
        if let Some(parent) = data.parent
            && let Some(parent_scope) = self.scopes.get_mut(parent)
        {
            parent_scope.children.retain(|c| *c != scope);
        }
        if self.current_scope == scope {
            self.current_scope = self.root_scope;
        }
    }
}

fn with_runtime<R>(f: impl FnOnce(&mut Runtime) -> R) -> R {
    RUNTIME.with(|rt| f(&mut rt.borrow_mut()))
}

/// Bring `key` up to date, recomputing it (and any stale sources) as needed.
/// May run user code; never holds the runtime borrow across it.
fn update_if_necessary(key: NodeKey) {
    let state = match with_runtime(|rt| rt.nodes.get(key).map(|n| n.state)) {
        Some(state) => state,
        None => return,
    };
    if state == State::Check {
        let sources = with_runtime(|rt| {
            rt.nodes
                .get(key)
                .map(|n| n.sources.clone())
                .unwrap_or_default()
        });
        for source in sources {
            update_if_necessary(source);
            // A recomputed source that changed value has promoted us to
            // Dirty; no need to check the rest.
            let now = with_runtime(|rt| rt.nodes.get(key).map(|n| n.state));
            match now {
                Some(State::Dirty) => break,
                Some(_) => {}
                None => return,
            }
        }
    }
    let state = match with_runtime(|rt| rt.nodes.get(key).map(|n| n.state)) {
        Some(state) => state,
        None => return,
    };
    if state == State::Dirty {
        recompute(key);
    }
    with_runtime(|rt| {
        if let Some(node) = rt.nodes.get_mut(key) {
            node.state = State::Clean;
        }
    });
}

/// Re-run a memo or effect body with dependency tracking.
fn recompute(key: NodeKey) {
    let taken = with_runtime(|rt| {
        rt.clear_sources(key);
        let node = rt.nodes.get_mut(key)?;
        let run = node.run.take()?;
        let value = node.value.take();
        let prev_observer = rt.observer.replace(key);
        Some((run, value, node.kind, prev_observer))
    });
    let Some((mut run, mut value, kind, prev_observer)) = taken else {
        return;
    };

    let changed = run(&mut value);

    with_runtime(|rt| {
        rt.observer = prev_observer;
        if let Some(node) = rt.nodes.get_mut(key) {
            node.run = Some(run);
            node.value = value;
            if changed && kind == Kind::Memo {
                let subscribers = node.subscribers.clone();
                for sub in subscribers {
                    rt.stale(sub, State::Dirty);
                }
            }
        }
        // If the node was disposed while running, its body and value are
        // simply dropped here.
    });
}

// --- crate-internal API used by the typed wrappers in lib.rs ---

pub(crate) fn create_node(kind: Kind, value: Option<Box<dyn Any>>, run: Option<RunFn>) -> NodeKey {
    with_runtime(|rt| rt.insert_node(kind, value, run))
}

/// Create and immediately execute a memo or effect, establishing its initial
/// dependencies.
pub(crate) fn create_and_run(kind: Kind, run: RunFn) -> NodeKey {
    let key = create_node(kind, None, Some(run));
    with_runtime(|rt| {
        if let Some(node) = rt.nodes.get_mut(key) {
            node.state = State::Dirty;
        }
    });
    recompute(key);
    with_runtime(|rt| {
        if let Some(node) = rt.nodes.get_mut(key) {
            node.state = State::Clean;
        }
    });
    key
}

/// Read a node's value. Tracks the read if `track` and the node is up to
/// date first if it is a memo.
pub(crate) fn read<R>(key: NodeKey, track: bool, f: impl FnOnce(Option<&dyn Any>) -> R) -> R {
    let is_memo = with_runtime(|rt| rt.nodes.get(key).map(|n| n.kind == Kind::Memo));
    if is_memo == Some(true) {
        update_if_necessary(key);
    }
    with_runtime(|rt| {
        if track {
            rt.track(key);
        }
        f(rt.nodes.get(key).and_then(|n| n.value.as_deref()))
    })
}

/// Write a signal's value (already compared by the typed wrapper) and mark
/// dependents stale.
pub(crate) fn write(key: NodeKey, value: Box<dyn Any>) {
    with_runtime(|rt| {
        let Some(node) = rt.nodes.get_mut(key) else {
            return;
        };
        debug_assert_eq!(node.kind, Kind::Signal, "write is only for plain signals");
        node.value = Some(value);
        let subscribers = node.subscribers.clone();
        for sub in subscribers {
            rt.stale(sub, State::Dirty);
        }
    });
}

/// Mutate a signal's value in place and mark dependents stale.
pub(crate) fn write_with(key: NodeKey, f: impl FnOnce(&mut dyn Any)) {
    let changed = with_runtime(|rt| {
        let Some(node) = rt.nodes.get_mut(key) else {
            return false;
        };
        if let Some(value) = node.value.as_deref_mut() {
            f(value);
        }
        true
    });
    if changed {
        with_runtime(|rt| {
            let subscribers = rt
                .nodes
                .get(key)
                .map(|n| n.subscribers.clone())
                .unwrap_or_default();
            for sub in subscribers {
                rt.stale(sub, State::Dirty);
            }
        });
    }
}

pub(crate) fn node_exists(key: NodeKey) -> bool {
    with_runtime(|rt| rt.nodes.contains_key(key))
}

pub(crate) fn dispose_node(key: NodeKey) {
    with_runtime(|rt| {
        if let Some(node) = rt.nodes.get(key) {
            let scope = node.scope;
            if let Some(scope) = rt.scopes.get_mut(scope) {
                scope.nodes.retain(|n| *n != key);
            }
        }
        rt.remove_node(key);
    });
}

/// Run every queued effect (and any they transitively queue) to quiescence.
pub fn flush_effects() {
    let already_flushing = with_runtime(|rt| {
        if rt.flushing {
            true
        } else {
            rt.flushing = true;
            false
        }
    });
    if already_flushing {
        return;
    }
    let mut rounds = 0;
    loop {
        let batch = with_runtime(|rt| std::mem::take(&mut rt.queue));
        if batch.is_empty() {
            break;
        }
        rounds += 1;
        assert!(
            rounds <= MAX_FLUSH_ROUNDS,
            "effect flush did not converge after {MAX_FLUSH_ROUNDS} rounds; \
             effects are likely writing signals they depend on"
        );
        for key in batch {
            let runnable = with_runtime(|rt| match rt.nodes.get_mut(key) {
                Some(node) if node.queued => {
                    node.queued = false;
                    true
                }
                _ => false,
            });
            if runnable {
                update_if_necessary(key);
            }
        }
    }
    with_runtime(|rt| rt.flushing = false);
}

/// Whether any effects are waiting for [`flush_effects`].
pub fn has_pending_effects() -> bool {
    with_runtime(|rt| !rt.queue.is_empty())
}

// --- scopes ---

pub(crate) fn create_scope() -> ScopeKey {
    with_runtime(|rt| {
        let parent = rt.current_scope;
        let key = rt.scopes.insert(Scope {
            parent: Some(parent),
            ..Default::default()
        });
        rt.scopes[parent].children.push(key);
        key
    })
}

pub(crate) fn enter_scope(scope: ScopeKey) -> ScopeKey {
    with_runtime(|rt| {
        let prev = rt.current_scope;
        if rt.scopes.contains_key(scope) {
            rt.current_scope = scope;
        }
        prev
    })
}

pub(crate) fn restore_scope(scope: ScopeKey) {
    with_runtime(|rt| {
        if rt.scopes.contains_key(scope) {
            rt.current_scope = scope;
        } else {
            rt.current_scope = rt.root_scope;
        }
    });
}

pub(crate) fn dispose_scope(scope: ScopeKey) {
    with_runtime(|rt| rt.dispose_scope(scope));
}