//! Phase 5: Finalization
//!
//! Removes unused constants and builds the OnnxGraphBuilder.

use std::{
    cell::RefCell,
    collections::{HashMap, HashSet},
    rc::Rc,
};

use crate::{
    graph_state::GraphState,
    ir::{Argument, NodeType, OnnxGraphBuilder, RawNode},
};

/// Finalize the graph by removing unused constants and building OnnxGraphBuilder
///
/// Returns OnnxGraphBuilder which still contains RawNode instances.
/// The caller should convert this to OnnxGraph using convert_to_graph().
pub(crate) fn finalize(
    nodes: &mut Vec<RawNode>,
    inputs: Vec<Argument>,
    outputs: &mut Vec<Argument>,
    state_rc: Rc<RefCell<GraphState>>,
) -> OnnxGraphBuilder {
    remove_unreferenced_constants(nodes, outputs);

    OnnxGraphBuilder {
        nodes: std::mem::take(nodes),
        inputs,
        outputs: std::mem::take(outputs),
        graph_state: Some(state_rc),
    }
}

/// Remove constant nodes that have zero runtime references
fn remove_unreferenced_constants(nodes: &mut Vec<RawNode>, outputs: &[Argument]) {
    // Build map of constant output names to node indices
    let mut constant_output_to_idx: HashMap<String, usize> = HashMap::new();
    for (idx, node) in nodes.iter().enumerate() {
        if node.node_type == NodeType::Constant {
            for output in &node.outputs {
                constant_output_to_idx.insert(output.name.clone(), idx);
            }
        }
    }

    // Count references (only Constant/Dynamic, not Static)
    let mut constant_references: HashMap<String, usize> = HashMap::new();

    for node in nodes.iter() {
        for input in &node.inputs {
            if (input.is_constant() || input.is_dynamic())
                && constant_output_to_idx.contains_key(&input.name)
            {
                *constant_references.entry(input.name.clone()).or_insert(0) += 1;
            }
        }
    }

    for output in outputs {
        if (output.is_constant() || output.is_dynamic())
            && constant_output_to_idx.contains_key(&output.name)
        {
            *constant_references.entry(output.name.clone()).or_insert(0) += 1;
        }
    }

    // Mark constants with zero references for removal
    let mut constants_to_remove = HashSet::new();
    for (output_name, &node_idx) in &constant_output_to_idx {
        let ref_count = constant_references.get(output_name).unwrap_or(&0);
        if *ref_count == 0 {
            constants_to_remove.insert(node_idx);
        }
    }

    // Filter out unreferenced constants
    let mut i = 0;
    nodes.retain(|_node| {
        let keep = !constants_to_remove.contains(&i);
        i += 1;
        keep
    });
}
