jingle 0.5.2

use crate::analysis::pcode_store::{PcodeOpRef, PcodeStore};
use crate::modeling::machine::cpu::concrete::ConcretePcodeAddress;
use jingle_sleigh::{PcodeOperation, SleighArchInfo};
pub use model::{CfgState, CfgStateModel, ModelTransition};
use petgraph::Direction;
use petgraph::graph::NodeIndex;
use petgraph::prelude::StableDiGraph;
use petgraph::visit::EdgeRef;
use std::borrow::Borrow;
use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::fmt::{Display, Formatter, LowerHex};
use std::rc::Rc;

pub(crate) mod model;

#[derive(Debug, Default, Copy, Clone, Hash)]
pub struct EmptyEdge;

impl Display for EmptyEdge {
    fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result {
        Ok(())
    }
}
impl LowerHex for EmptyEdge {
    fn fmt(&self, _: &mut Formatter<'_>) -> std::fmt::Result {
        Ok(())
    }
}

#[derive(Debug)]
pub struct PcodeCfg<N, D> {
    pub(crate) graph: StableDiGraph<N, EmptyEdge>,
    // Key pcode ops by concrete p-code address so lookups by address are efficient.
    // For CFG states `N` that can produce a `ConcretePcodeAddress` via `location()`,
    // we provide lookup helpers that map a `N` to its address and query this map.
    pub(crate) ops: HashMap<ConcretePcodeAddress, D>,
    pub(crate) indices: HashMap<N, NodeIndex>,
}

#[derive(Debug)]
pub struct ModeledPcodeCfg<N: CfgState = ConcretePcodeAddress, D = PcodeOperation> {
    pub(crate) cfg: PcodeCfg<N, D>,
    #[allow(unused)]
    pub(crate) info: SleighArchInfo,
    pub(crate) models: HashMap<N, N::Model>,
}

#[derive(Clone)]
pub struct PcodeCfgVisitor<'a, N: CfgState = ConcretePcodeAddress, D = PcodeOperation> {
    cfg: &'a ModeledPcodeCfg<N, D>,
    location: N,
    pub(crate) visited_locations: Rc<RefCell<HashSet<N>>>,
}

impl<'a, N: CfgState, D: ModelTransition<N::Model>> PcodeCfgVisitor<'a, N, D> {
    pub(crate) fn successors(&mut self) -> impl Iterator<Item = Self> {
        self.cfg
            .cfg
            .successors(&self.location)
            .into_iter()
            .flatten()
            .flat_map(|n| {
                // Use a short-lived borrow so the RefMut is released before we construct the new visitor
                let is_repeat = {
                    let mut set = self.visited_locations.borrow_mut();
                    if set.contains(n) {
                        true
                    } else {
                        set.insert(n.clone());
                        false
                    }
                };

                if is_repeat {
                    None
                } else {
                    Some(Self {
                        cfg: self.cfg,
                        location: n.clone(),
                        visited_locations: self.visited_locations.clone(),
                    })
                }
            })
    }

    pub(crate) fn transition(&self) -> Option<&D> {
        // Lookup the transition by the concrete address of the current location (if any)
        self.cfg.cfg.get_op_at(&self.location)
    }

    pub fn location(&self) -> &N {
        &self.location
    }

    pub fn state(&self) -> Option<&N::Model> {
        self.cfg.models.get(&self.location)
    }
}

impl<N: CfgState, D: ModelTransition<N::Model>> Default for PcodeCfg<N, D> {
    fn default() -> Self {
        Self::new()
    }
}

impl<N: CfgState, D: ModelTransition<N::Model>> PcodeCfg<N, D> {
    pub fn new() -> Self {
        Self {
            graph: Default::default(),
            ops: Default::default(),
            indices: Default::default(),
        }
    }

    pub fn graph(&self) -> &StableDiGraph<N, EmptyEdge> {
        &self.graph
    }

    pub fn nodes(&self) -> impl Iterator<Item = &N> {
        self.indices.keys()
    }

    /// Check whether the CFG contains a node (by value)
    pub fn has_node<T: Borrow<N>>(&self, node: T) -> bool {
        self.indices.contains_key(node.borrow())
    }

    /// Lookup an op by a CFG state `N` by first resolving its concrete p-code address via
    /// `location()` and then performing the address-keyed lookup.
    pub fn get_op_at<T: Borrow<N>>(&self, addr: T) -> Option<&D> {
        let n = addr.borrow();
        // If the state has an associated concrete address, look up by that address.
        n.location().and_then(|a| self.ops.get(&a))
    }

    pub fn get_op_at_address<C: Borrow<ConcretePcodeAddress>>(&self, addr: C) -> Option<&D> {
        self.ops.get(addr.borrow())
    }

    pub fn add_node<T: Borrow<N>>(&mut self, node: T) {
        let node = node.borrow();
        if !self.indices.contains_key(node) {
            let idx = self.graph.add_node(node.clone());
            self.indices.insert(node.clone(), idx);
        }
    }

    pub fn replace_and_combine_nodes<T: Borrow<N>, S: Borrow<N>>(
        &mut self,
        old_weight: T,
        new_weight: S,
    ) {
        // Copy the indices first to avoid borrow issues
        let old_idx = self.indices.get(old_weight.borrow()).copied();
        let new_idx = self.indices.get(new_weight.borrow()).copied();

        tracing::debug!(
            "replace_and_combine_nodes called: old_idx={:?}, new_idx={:?}",
            old_idx,
            new_idx
        );

        if let (Some(old_idx), Some(new_idx)) = (old_idx, new_idx) {
            // If the indices are the same, the nodes are already merged - nothing to do
            if old_idx == new_idx {
                tracing::debug!("Indices are identical, skipping merge");
                return;
            }

            tracing::debug!("Both nodes found with different indices, proceeding with merge");
            // We are going to keep the old index, but replace its weight with new_weight
            // All edges from new_idx will be redirected to old_idx, then new_idx is removed

            // Redirect all incoming edges from new_idx to old_idx
            let incoming: Vec<_> = self
                .graph
                .edges_directed(new_idx, Direction::Incoming)
                .map(|edge| edge.source())
                .collect();
            for source in incoming {
                if !self.graph.contains_edge(source, old_idx) {
                    self.graph.add_edge(source, old_idx, EmptyEdge);
                }
            }

            // Redirect all outgoing edges from new_idx to old_idx
            let outgoing: Vec<_> = self
                .graph
                .edges_directed(new_idx, Direction::Outgoing)
                .map(|edge| edge.target())
                .collect();
            for target in &outgoing {
                if !self.graph.contains_edge(old_idx, *target) {
                    self.graph.add_edge(old_idx, *target, EmptyEdge);
                }
            }

            // Remove the new node from the graph (using StableGraph so indices remain valid)
            self.graph.remove_node(new_idx);

            // Update the weight at old_idx to be new_weight
            if let Some(node_weight) = self.graph.node_weight_mut(old_idx) {
                *node_weight = new_weight.borrow().clone();
            }

            // Update the indices map: new_weight should now map to old_idx
            self.indices.insert(new_weight.borrow().clone(), old_idx);
            self.indices.remove(old_weight.borrow());

            // Update the ops map: prefer the op from new_weight if it exists, otherwise use old_weight's op
            // Prefer the op from `new_weight` (by its concrete address), otherwise fall back to `old_weight`.
            let op_to_keep = {
                let new_loc = new_weight.borrow().location();
                let old_loc = old_weight.borrow().location();
                new_loc
                    .and_then(|a| self.ops.get(&a).cloned())
                    .or_else(|| old_loc.and_then(|a| self.ops.get(&a).cloned()))
            };

            // Remove any existing ops keyed by the concrete addresses of the old/new weights.
            if let Some(loc) = old_weight.borrow().location() {
                self.ops.remove(&loc);
            }
            if let Some(loc) = new_weight.borrow().location() {
                self.ops.remove(&loc);
            }

            if let Some(op) = op_to_keep {
                if let Some(loc) = new_weight.borrow().location() {
                    self.ops.insert(loc, op);
                } else {
                    // If the new_weight has no concrete location, we cannot re-insert under an address.
                    dbg!("Missing new_weight location for op insertion!");
                }
            } else {
                dbg!("Missing op!");
            }
        }
    }

    pub fn add_edge<A, B, C>(&mut self, from: A, to: B, op: C)
    where
        A: Borrow<N>,
        B: Borrow<N>,
        C: Borrow<D>,
    {
        let from = from.borrow();
        let to = to.borrow();
        let op = op.borrow();
        self.add_node(from);
        self.add_node(to);
        // Insert op keyed by the concrete address of `from` (if available).
        if let Some(loc) = from.location() {
            self.ops.insert(loc, op.clone());
        }
        let from_idx = *self.indices.get(from).unwrap();
        let to_idx = *self.indices.get(to).unwrap();
        // Return early if the edge already exists
        if self.graph.find_edge(from_idx, to_idx).is_some() {
            return;
        }
        self.graph.add_edge(from_idx, to_idx, EmptyEdge);
    }

    /// Return successors of a node (by value) as references into the backing CFG.
    /// Returns `None` if the node is not present in the CFG.
    pub fn successors<T: Borrow<N>>(&self, node: T) -> Option<Vec<&N>> {
        let n = node.borrow();
        let idx = *self.indices.get(n)?;
        let succs: Vec<&N> = self
            .graph
            .edges_directed(idx, Direction::Outgoing)
            .map(|e| self.graph.node_weight(e.target()).unwrap())
            .collect();
        Some(succs)
    }

    /// Return predecessors of a node (by value) as references into the backing CFG.
    /// Returns `None` if the node is not present in the CFG.
    pub fn predecessors<T: Borrow<N>>(&self, node: T) -> Option<Vec<&N>> {
        let n = node.borrow();
        let idx = *self.indices.get(n)?;
        let preds: Vec<&N> = self
            .graph
            .edges_directed(idx, Direction::Incoming)
            .map(|e| self.graph.node_weight(e.source()).unwrap())
            .collect();
        Some(preds)
    }

    pub fn leaf_nodes(&self) -> impl Iterator<Item = &N> {
        self.graph
            .externals(Direction::Outgoing)
            .map(move |idx| self.graph.node_weight(idx).unwrap())
    }

    pub fn entry_nodes(&self) -> impl Iterator<Item = &N> {
        self.graph
            .externals(Direction::Incoming)
            .map(move |idx| self.graph.node_weight(idx).unwrap())
    }

    pub fn edge_weights(&self) -> impl Iterator<Item = &D> {
        self.ops.values()
    }

    pub fn nodes_for_location<S: PartialEq<N>>(&self, location: S) -> impl Iterator<Item = &N> {
        self.nodes().filter(move |a| location == **a)
    }

    /// Create a `ModeledPcodeCfg` by generating SMT models for all nodes in the CFG.
    pub fn smt_model(self, info: SleighArchInfo) -> ModeledPcodeCfg<N, D> {
        ModeledPcodeCfg::new(self, info)
    }

    /// Build a sub-CFG from a set of `NodeIndex` values.
    ///
    /// This will include:
    /// - all nodes whose indices are in `selected`
    /// - all edges between those nodes
    ///
    /// If the resulting subgraph has more than one connected component (treating edges as undirected
    /// for connectivity), this function will panic.
    pub fn sub_cfg_from_indices(&self, selected: &HashSet<NodeIndex>) -> PcodeCfg<N, D> {
        // Build the set of NodeIndex that we will include: exactly the selected nodes.
        // Do NOT include outgoing neighbors or nodes that only have incoming edges into the selected set.
        let mut include: HashSet<NodeIndex> = HashSet::new();
        for &idx in selected {
            include.insert(idx);
        }

        // Collect the nodes to work with
        let nodes: Vec<NodeIndex> = include.iter().copied().collect();

        // If there are nodes, verify the induced subgraph is a single (undirected) component.
        if !nodes.is_empty() {
            let mut visited: HashSet<NodeIndex> = HashSet::new();
            let mut stack: Vec<NodeIndex> = Vec::new();

            // Start from the first node
            stack.push(nodes[0]);
            visited.insert(nodes[0]);

            while let Some(u) = stack.pop() {
                // Outgoing neighbors
                for e in self.graph.edges_directed(u, Direction::Outgoing) {
                    let v = e.target();
                    if include.contains(&v) && visited.insert(v) {
                        stack.push(v);
                    }
                }
                // Incoming neighbors (treat graph as undirected for connectivity)
                for e in self.graph.edges_directed(u, Direction::Incoming) {
                    let v = e.source();
                    if include.contains(&v) && visited.insert(v) {
                        stack.push(v);
                    }
                }
            }

            if visited.len() != include.len() {
                panic!("sub_cfg_from_indices produced more than one component");
            }
        }

        // Build the new graph and maps containing only the included nodes and edges between them
        let mut new_graph = StableDiGraph::<N, EmptyEdge>::default();
        let mut new_indices: HashMap<N, NodeIndex> = HashMap::new();
        // New ops map keyed by concrete pcode addresses (consistent with the parent CFG)
        let mut new_ops: HashMap<ConcretePcodeAddress, D> = HashMap::new();

        // Add nodes and their ops (copy ops by address if the node exposes a concrete address)
        for &old_idx in &nodes {
            let n = self.graph.node_weight(old_idx).unwrap().clone();
            let new_idx = new_graph.add_node(n.clone());
            new_indices.insert(n.clone(), new_idx);
            if let Some(loc) = n.location() {
                if let Some(op) = self.ops.get(&loc) {
                    new_ops.insert(loc, op.clone());
                }
            }
        }

        // Add edges where both endpoints are included
        for edge in self.graph.edge_indices() {
            let (from, to) = self.graph.edge_endpoints(edge).unwrap();
            if include.contains(&from) && include.contains(&to) {
                let from_w = self.graph.node_weight(from).unwrap();
                let to_w = self.graph.node_weight(to).unwrap();
                let from_new = *new_indices.get(from_w).unwrap();
                let to_new = *new_indices.get(to_w).unwrap();
                new_graph.add_edge(from_new, to_new, EmptyEdge);
            }
        }

        PcodeCfg {
            graph: new_graph,
            ops: new_ops,
            indices: new_indices,
        }
    }
}

impl<'op, L: CfgState> PcodeStore<'op> for PcodeCfg<L, PcodeOperation> {
    fn get_pcode_op_at<T: Borrow<ConcretePcodeAddress>>(
        &'op self,
        addr: T,
    ) -> Option<crate::analysis::pcode_store::PcodeOpRef<'op>> {
        let addr = *addr.borrow();
        self.get_op_at_address(addr)
            .map(crate::analysis::pcode_store::PcodeOpRef::from)
    }
}

impl<'op, L: CfgState> PcodeStore<'op> for PcodeCfg<L, PcodeOpRef<'op>> {
    fn get_pcode_op_at<T: Borrow<ConcretePcodeAddress>>(
        &'op self,
        addr: T,
    ) -> Option<crate::analysis::pcode_store::PcodeOpRef<'op>> {
        let addr = *addr.borrow();
        // The CFG stores `PcodeOpRef<'static>` values. We can return a borrowed
        // `PcodeOpRef<'op>` that references the inner operation without cloning.
        self.get_op_at_address(addr).map(|op_ref_static| {
            crate::analysis::pcode_store::PcodeOpRef::from(op_ref_static.as_ref())
        })
    }
}

/// Trait providing a `basic_blocks` transformation for CFGs whose op storage type
/// can yield a reference to a `PcodeOperation` via `AsRef<PcodeOperation>`.
pub trait BasicBlocks<N: CfgState> {
    fn basic_blocks(&self) -> PcodeCfg<N, Vec<PcodeOperation>>;
}

impl<N: CfgState, D: AsRef<PcodeOperation>> BasicBlocks<N> for PcodeCfg<N, D>
where
    PcodeOperation: Clone,
{
    fn basic_blocks(&self) -> PcodeCfg<N, Vec<PcodeOperation>> {
        use petgraph::visit::EdgeRef;
        // Step 1: Initialize new graph and maps
        let mut graph = StableDiGraph::<N, EmptyEdge>::default();
        // ops keyed by concrete pcode addresses
        let mut ops: HashMap<ConcretePcodeAddress, Vec<PcodeOperation>> = HashMap::new();
        let mut indices: HashMap<N, NodeIndex> = HashMap::new();
        // Step 2: Wrap each op in a Vec and add nodes
        for node in self.graph.node_indices() {
            let n = self.graph.node_weight(node).unwrap().clone();
            let op = if let Some(loc) = n.location() {
                self.ops
                    .get(&loc)
                    .map(|opd| vec![opd.as_ref().clone()])
                    .unwrap_or_default()
            } else {
                Vec::new()
            };
            let idx = graph.add_node(n.clone());
            graph.add_node(n.clone());
            indices.insert(n.clone(), idx);
            if let Some(loc) = n.location() {
                ops.insert(loc, op);
            }
        }
        // Step 3: Add edges
        for edge in self.graph.edge_indices() {
            let (fromidx, toidx) = self.graph.edge_endpoints(edge).unwrap();
            let from = self.graph.node_weight(fromidx).unwrap();
            let to = self.graph.node_weight(toidx).unwrap();
            let from_idx = *indices.get(from).unwrap();
            let to_idx = *indices.get(to).unwrap();
            graph.add_edge(from_idx, to_idx, EmptyEdge);
        }
        // Step 4: Wait-list algorithm for merging
        let mut changed = true;
        while changed {
            changed = false;
            for node in graph.node_indices() {
                // Only consider nodes still present
                let out_edges: Vec<_> = graph.edges_directed(node, Direction::Outgoing).collect();
                if out_edges.len() != 1 {
                    continue;
                }
                let edge = out_edges[0];
                let target = edge.target();
                let in_edges: Vec<_> = graph.edges_directed(target, Direction::Incoming).collect();
                if in_edges.len() != 1 {
                    continue;
                }

                // Merge target into node
                let src_n = graph.node_weight(node).unwrap().clone();
                let tgt_n = graph.node_weight(target).unwrap().clone();
                // Fix borrow: collect target ops first (by concrete address)
                let tgt_ops = if let Some(tgt_loc) = tgt_n.location() {
                    ops.get(&tgt_loc).cloned().unwrap_or_default()
                } else {
                    Vec::new()
                };
                if !tgt_ops.is_empty() {
                    if let Some(src_loc) = src_n.location() {
                        ops.entry(src_loc).or_default().extend(tgt_ops);
                    }
                }
                // Redirect outgoing edges of target to source
                let tgt_out_edges: Vec<_> = graph
                    .edges_directed(target, Direction::Outgoing)
                    .map(|e| e.target())
                    .collect();
                for tgt_out in tgt_out_edges {
                    graph.add_edge(node, tgt_out, EmptyEdge);
                }
                // Remove target node and its ops
                graph.remove_node(target);
                if let Some(tgt_loc) = tgt_n.location() {
                    ops.remove(&tgt_loc);
                }
                indices.remove(&tgt_n);
                changed = true;
                break; // Restart after each merge
            }
        }
        // Step 5: Build a new graph using only connected nodes
        let mut new_graph = StableDiGraph::<N, EmptyEdge>::default();
        let mut new_ops: HashMap<ConcretePcodeAddress, Vec<PcodeOperation>> = HashMap::new();
        let mut new_indices: HashMap<N, NodeIndex> = HashMap::new();
        // Collect connected nodes
        let connected_nodes: Vec<_> = graph
            .node_indices()
            .filter(|&node| {
                graph
                    .edges_directed(node, Direction::Incoming)
                    .next()
                    .is_some()
                    || graph
                        .edges_directed(node, Direction::Outgoing)
                        .next()
                        .is_some()
            })
            .collect();
        // Add connected nodes to new graph
        for node in connected_nodes.iter() {
            let n = graph.node_weight(*node).unwrap().clone();
            let idx = new_graph.add_node(n.clone());
            new_indices.insert(n.clone(), idx);
            if let Some(loc) = n.location() {
                if let Some(op) = ops.get(&loc) {
                    new_ops.insert(loc, op.clone());
                }
            }
        }
        // Add edges between connected nodes
        for edge in graph.edge_indices() {
            let (fromidx, toidx) = graph.edge_endpoints(edge).unwrap();
            if connected_nodes.contains(&fromidx) && connected_nodes.contains(&toidx) {
                let from = graph.node_weight(fromidx).unwrap();
                let to = graph.node_weight(toidx).unwrap();
                let from_idx = *new_indices.get(from).unwrap();
                let to_idx = *new_indices.get(to).unwrap();
                new_graph.add_edge(from_idx, to_idx, EmptyEdge);
            }
        }
        // Step 6: Build and return new PcodeCfg
        PcodeCfg {
            graph: new_graph,
            ops: new_ops,
            indices: new_indices,
        }
    }
}

impl<N: CfgState, D: ModelTransition<N::Model>> ModeledPcodeCfg<N, D> {
    pub fn new(cfg: PcodeCfg<N, D>, info: SleighArchInfo) -> Self {
        let mut models = HashMap::new();
        for node in cfg.nodes() {
            let model = node.new_const(&info);
            models.insert(node.clone(), model);
        }
        Self { cfg, models, info }
    }

    pub fn cfg(&self) -> &PcodeCfg<N, D> {
        &self.cfg
    }

    pub fn models(&self) -> &HashMap<N, N::Model> {
        &self.models
    }

    // Delegation methods to underlying PcodeCfg
    pub fn graph(&self) -> &StableDiGraph<N, EmptyEdge> {
        self.cfg.graph()
    }

    pub fn nodes(&self) -> impl Iterator<Item = &N> {
        self.cfg.nodes()
    }

    pub fn leaf_nodes(&self) -> impl Iterator<Item = &N> {
        self.cfg.leaf_nodes()
    }

    pub fn edge_weights(&self) -> impl Iterator<Item = &D> {
        self.cfg.edge_weights()
    }

    pub fn nodes_for_location<S: PartialEq<N>>(&self, location: S) -> impl Iterator<Item = &N> {
        self.cfg.nodes_for_location(location)
    }
}