//! The main code used for tracking dependencies (mostly) invisibly to the user.
//!
//! Each mapping that could be memoised is represented by a Node. Each Node can have one or more dependencies, forming a directed acyclic graph.
//!
//! [`self::add_dep`] should be called whenever any mapping is retrieved, to record the dependency.
//!
//! [`self::next_node_id`] and [`self::with_node_id`] should be used to create new nodes in the graph when a new mapping is being evaluated.
//!
//! [`self::mark_dirty`] should be called from outside the evaluation of any mapping. In many cases, this requires the dependency ID to be stored externally alongside the key, which can be accomplished with [`self::current_node_id`]
//!
//! [`self::is_dirty`] should be used when a mapping is retrieved to determine if it needs to be re-evaluated. If it does, [`self::clear`] should be used to reset the mapping's dependencies and dirty state.
use std::{
    cell::RefCell,
    collections::{HashMap, VecDeque},
    fmt::Write,
    future::Future,
    sync::{Mutex, OnceLock},
    thread,
};

use tracing::instrument;

/// Identifier of a node, unique across a program run.
pub type NodeId = u64;

tokio::task_local! {
    /// The ID of the node in the dependency tree for the currently running computation
    static NODE_ID: NodeId;
}

thread_local! {
    static NEXT_NODE_ID_LSB: RefCell<u64> = RefCell::new(5);
}

/// Responsible for tracking dependencies.
/// There is one global instance of this, [`DEP_TRACKER`], with convenience functions for most of its functionality.
#[derive(Default)]
struct DepTracker {
    /// Neighbour-List representation of dependency tree
    deps: Mutex<HashMap<NodeId, NodeInfo>>,
}

impl DepTracker {
    /// See [`self::add_dep`]
    #[instrument(level = "trace", skip(self))]
    fn add_dep(&self, dep: NodeId) {
        self.deps
            .lock()
            .unwrap()
            .entry(dep)
            .and_modify(|v| v.dependents.push(NODE_ID.get()))
            .or_insert(NodeInfo {
                dependents: vec![NODE_ID.get()],
                dirty: false,
            });
    }

    /// See [`self::next_node_id`]
    fn next_node_id(&self) -> NodeId {
        let tid = thread::current().id().as_u64();
        let lsb = NEXT_NODE_ID_LSB.with(|lsb| {
            let old = *lsb.borrow();
            *lsb.borrow_mut() = old + 1;
            old
        });

        (tid.get() << 32) | (lsb & 0xffffffff)
    }

    /// See [`self::mark_dirty`]
    #[instrument(level = "trace", skip(self))]
    fn mark_dirty(&self, node: NodeId) {
        let mut lock = self.deps.lock().unwrap();
        let mut frontier = VecDeque::new();
        frontier.push_back(node);
        while let Some(node_id) = frontier.pop_front() {
            let node = match lock.get_mut(&node_id) {
                Some(x) => x,
                None => continue,
            };

            if node.dirty {
                continue;
            }
            node.dirty = true;

            for dependent in node.dependents.iter() {
                frontier.push_back(*dependent);
            }
        }
    }

    /// See [`self::is_dirty`]
    #[instrument(level = "trace", skip(self))]
    fn is_dirty(&self, node: NodeId) -> bool {
        self.deps
            .lock()
            .unwrap()
            .get(&node)
            .map(|v| v.dirty)
            .unwrap_or(false)
    }

    /// See [`self::clear`]
    #[instrument(level = "trace", skip(self))]
    fn clear(&self, node: NodeId) {
        let mut lock = self.deps.lock().unwrap();
        let node = match lock.get_mut(&node) {
            Some(x) => x,
            None => return,
        };

        node.dirty = false;
        node.dependents = vec![];
    }
}

/// Info about a single dependency node
struct NodeInfo {
    dependents: Vec<NodeId>,
    dirty: bool,
}

/// The global instance of the dependency tracker
static DEP_TRACKER: OnceLock<DepTracker> = OnceLock::new();
fn dep_tracker() -> &'static DepTracker {
    DEP_TRACKER.get_or_init(DepTracker::default)
}

/// Register a dependency for the current node
pub fn add_dep(dep: NodeId) {
    dep_tracker().add_dep(dep);
}

/// Get the next node id to use
pub fn next_node_id() -> NodeId {
    dep_tracker().next_node_id()
}

/// Get the ID of the current node
pub fn current_node_id() -> NodeId {
    NODE_ID.get()
}

/// Mark the given node and all nodes that depend on it 'dirty'.
///
/// This should be used in conjunction with [`self::is_dirty`] to recompute values when necessary.
pub fn mark_dirty(dep: NodeId) {
    dep_tracker().mark_dirty(dep)
}

/// Check if the given node is dirty, indicating it should be recomputed.
pub fn is_dirty(dep: NodeId) -> bool {
    dep_tracker().is_dirty(dep)
}

/// Clear all information about a given node.
///
/// This should be used when a node is recomputed, to ensure stale dependencies are removed.
pub fn clear(dep: NodeId) {
    dep_tracker().clear(dep)
}

/// Generate a graphviz representation of the current dependency graph, for debugging.
///
/// Dirty nodes will be coloured grey.
pub fn dep_graphviz() -> String {
    let mut out = String::new();
    let deps = dep_tracker().deps.lock().unwrap();
    writeln!(&mut out, "digraph G {{").unwrap();
    for (id, info) in deps.iter() {
        if *id == 0 {
            writeln!(&mut out, "\t{} [label=Root]", id).unwrap();
        } else if info.dirty {
            writeln!(&mut out, "\t{} [style=filled,color=lightgrey]", id).unwrap();
        }

        for dep in info.dependents.iter() {
            writeln!(&mut out, "\t{} -> {}", dep, id).unwrap();
        }
    }

    writeln!(&mut out, "}}").unwrap();

    out
}

/// Decorate the given future to have the given dependency ID.
pub async fn with_node_id<F>(dep: NodeId, f: F) -> F::Output
where
    F: Future + Send,
    F::Output: Send,
{
    NODE_ID.scope(dep, f).await
}