rill-graph 0.6.0-M2

use rill_core::math::Transcendental;
use rill_core::queues::CommandEnum;
use rill_core::traits::Params;
use rill_core_actor::ActorRef;

use indexmap::IndexMap;
use rill_lang::builtin::SignatureSource;
use rill_lang::graph_ir::{EdgeKind, GraphEdge, GraphIr, GraphNode};
use std::collections::HashMap;

// ============================================================================
// Build Errors
// ============================================================================

/// Errors that can occur during graph construction.
#[derive(Debug, Clone)]
pub enum BuildError {
    /// A cycle was detected in the signal edge graph.
    CycleDetected,
    /// Backend creation failed.
    Backend(String),
    /// A node type is not registered in the built-in registry.
    UnknownNodeType(String),
    /// The graph topology is not supported for conversion to a flat chain.
    UnsupportedTopology(String),
    /// AST compilation failed.
    CompilationFailed(String),
}

impl std::fmt::Display for BuildError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::CycleDetected => write!(f, "graph cycle detected"),
            Self::Backend(msg) => write!(f, "backend error: {msg}"),
            Self::UnknownNodeType(msg) => write!(f, "unknown node type: {msg}"),
            Self::UnsupportedTopology(msg) => write!(f, "unsupported topology: {msg}"),
            Self::CompilationFailed(msg) => write!(f, "compilation failed: {msg}"),
        }
    }
}

// ============================================================================
// Node Storage
// ============================================================================

/// A deferred node recipe — constructed at build_ir time.
struct NodeRecipe<T: Transcendental, const BUF_SIZE: usize> {
    type_name: String,
    id: u32,
    name: String,
    params: Params,
    routing_entries: Vec<(usize, usize, f32)>,
    _phantom: std::marker::PhantomData<(T, [(); BUF_SIZE])>,
}

// ============================================================================
// GraphBuilder (Mutable Construction)
// ============================================================================

/// A named resource (tape loop) shared between nodes in the graph.
#[derive(Clone)]
pub struct GraphResource {
    /// Unique name referenced by node parameters.
    pub name: String,
    /// Resource kind string (`"tape"`).
    pub kind: String,
    /// Capacity in samples (for `"tape"` kind).
    pub capacity: usize,
}

/// Mutable builder for an immutable signal graph.
pub struct GraphBuilder<T: Transcendental, const BUF_SIZE: usize> {
    recipes: Vec<NodeRecipe<T, BUF_SIZE>>,
    signal_edges: Vec<(usize, usize, usize, usize)>,
    control_edges: Vec<(usize, usize, usize, usize)>,
    clock_edges: Vec<(usize, usize, usize, usize)>,
    feedback_edges: Vec<(usize, usize, usize, usize)>,
    resources: Vec<GraphResource>,
    sample_rate: Option<f32>,
    parent_ref: Option<ActorRef<CommandEnum>>,
}

impl<T: Transcendental, const BUF_SIZE: usize> Default for GraphBuilder<T, BUF_SIZE> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T: Transcendental, const BUF_SIZE: usize> GraphBuilder<T, BUF_SIZE> {
    /// Create a new empty graph builder.
    pub fn new() -> Self {
        Self {
            recipes: Vec::new(),
            signal_edges: Vec::new(),
            control_edges: Vec::new(),
            clock_edges: Vec::new(),
            feedback_edges: Vec::new(),
            resources: Vec::new(),
            sample_rate: None,
            parent_ref: None,
        }
    }

    /// Add a node by type name.
    ///
    /// Returns the index of the newly added node.
    pub fn add_node(&mut self, type_name: &str, params: &Params) -> usize {
        let id = self.recipes.len() as u32;
        self.add_node_with_id(type_name, params, id)
    }

    /// Add a node with an explicit `NodeId`.
    pub fn add_node_with_id(&mut self, type_name: &str, params: &Params, id: u32) -> usize {
        self.add_node_with_name(type_name, params, id, String::new())
    }

    /// Add a node with an explicit `NodeId` and a human-readable name
    /// (typically sourced from the JSON `name` field). The name becomes the
    /// program/anchor name in the compiled graph, used by `SetParameter` routing.
    pub fn add_node_with_name(
        &mut self,
        type_name: &str,
        params: &Params,
        id: u32,
        name: String,
    ) -> usize {
        let idx = self.recipes.len();
        self.recipes.push(NodeRecipe {
            type_name: type_name.to_string(),
            id,
            name,
            params: params.clone(),
            routing_entries: Vec::new(),
            _phantom: std::marker::PhantomData,
        });
        idx
    }

    /// Store a routing matrix entry to be applied at build time.
    pub fn add_routing_entry(&mut self, idx: usize, from: usize, to: usize, gain: f32) {
        if let Some(recipe) = self.recipes.get_mut(idx) {
            recipe.routing_entries.push((from, to, gain));
        }
    }

    /// Register a named resource (tape loop, buffer, etc.).
    pub fn add_resource(&mut self, resource: GraphResource) {
        self.resources.push(resource);
    }

    /// Number of nodes added to the builder so far.
    pub fn node_count(&self) -> usize {
        self.recipes.len()
    }

    /// Set the sample rate for this builder.
    pub fn set_sample_rate(&mut self, sr: f32) {
        self.sample_rate = Some(sr);
    }

    /// Set the parent RackCase actor reference (Graph → parent ClockTick).
    pub fn set_parent_ref(&mut self, parent: ActorRef<CommandEnum>) {
        self.parent_ref = Some(parent);
    }

    /// Connect signal ports.
    pub fn connect_signal(
        &mut self,
        from_node: usize,
        from_port: usize,
        to_node: usize,
        to_port: usize,
    ) {
        self.signal_edges
            .push((from_node, from_port, to_node, to_port));
    }

    /// Connect control ports (modulation values).
    pub fn connect_control(
        &mut self,
        from_node: usize,
        from_port: usize,
        to_node: usize,
        to_port: usize,
    ) {
        self.control_edges
            .push((from_node, from_port, to_node, to_port));
    }

    /// Connect clock ports (timing events).
    pub fn connect_clock(
        &mut self,
        from_node: usize,
        from_port: usize,
        to_node: usize,
        to_port: usize,
    ) {
        self.clock_edges
            .push((from_node, from_port, to_node, to_port));
    }

    /// Connect feedback ports (delay lines, state carryover).
    pub fn connect_feedback(
        &mut self,
        from_node: usize,
        from_port: usize,
        to_node: usize,
        to_port: usize,
    ) {
        self.feedback_edges
            .push((from_node, from_port, to_node, to_port));
    }

    /// Build a [`rill_lang::graph_ir::GraphIr`] using the built-in `Registry`.
    ///
    /// This is the new execution path. It looks up each node type in the registry,
    /// constructs placeholder IRs, and performs topological sort. Actual compilation
    /// to executable programs happens in a future phase.
    pub fn build_ir(
        self,
        registry: &rill_lang::builtin::Registry<T>,
    ) -> Result<GraphIr, BuildError> {
        // 1. Build index → name mapping
        let idx_to_name: HashMap<usize, String> = self
            .recipes
            .iter()
            .enumerate()
            .map(|(idx, recipe)| {
                let name = if recipe.name.is_empty() {
                    format!("node_{}", recipe.id)
                } else {
                    recipe.name.clone()
                };
                (idx, name)
            })
            .collect();

        // 2. Create GraphNodes from recipes
        let mut nodes: IndexMap<String, GraphNode> = IndexMap::new();
        let mut node_list: Vec<String> = Vec::new();

        for (idx, recipe) in self.recipes.iter().enumerate() {
            let name = idx_to_name[&idx].clone();
            node_list.push(name.clone());

            let sig = registry
                .builtin_sig(&recipe.type_name)
                .or_else(|| {
                    // Strip "rill/" prefix for graph node → lang builtin mapping
                    recipe
                        .type_name
                        .strip_prefix("rill/")
                        .and_then(|n| registry.builtin_sig(n))
                })
                .or_else(|| {
                    // Common suffix mappings
                    let mapped = match recipe.type_name.as_str() {
                        "rill/dry_wet_mix" => "dry_wet",
                        "rill/parametric_eq" => "eq_parametric",
                        "rill/graphic_eq" => "graphic_eq",
                        "rill/mono_to_stereo" => "mono_to_stereo",
                        "rill/moog_ladder" => "moog",
                        "rill/write_head" => "write_head",
                        "rill/read_head" => "read_head",
                        "rill/lofi_chip" => "ay38910",
                        _ => "",
                    };
                    if mapped.is_empty() {
                        None
                    } else {
                        registry.builtin_sig(mapped)
                    }
                })
                .ok_or_else(|| BuildError::UnknownNodeType(recipe.type_name.clone()))?;

            let arity = (sig.signal_ins(), sig.signal_outs);

            // Convert all recipe parameters to ParamDef.
            // Include non-f32 values (SignalSlab placeholders) so that
            // SetParameter can target them by name via param_maps.
            let param_defs: Vec<rill_lang::ir::ParamDef> = recipe
                .params
                .parameters
                .iter()
                .map(|(k, v)| {
                    let default = v.as_f32().unwrap_or(0.0) as f64;
                    rill_lang::ir::ParamDef {
                        name: k.clone(),
                        default,
                        min: f64::NEG_INFINITY,
                        max: f64::INFINITY,
                    }
                })
                .collect();

            // Name → recipe-index lookup for building param_bindings.
            let name_to_recipe_idx: HashMap<String, usize> = param_defs
                .iter()
                .enumerate()
                .map(|(i, pd)| (pd.name.clone(), i))
                .collect();

            let param_values: Vec<f64>;
            let param_bindings: Vec<(usize, usize)>;

            if sig.param_names.is_empty() {
                // Backward compat: no names → positional identity (HashMap order).
                param_values = recipe
                    .params
                    .parameters
                    .values()
                    .filter_map(|v| v.as_f32().map(|f| f as f64))
                    .collect();
                param_bindings = (0..param_defs.len()).map(|i| (i, i)).collect();
            } else {
                // Named params: match recipe param names to builtin arg positions.
                // param_values[i] = value for builtin arg i, in correct positional order.
                // param_bindings[(arg_pos, recipe_param_idx)] — used by push_builtin_params
                // to route SetParameter changes to the right builtin set_param(arg_pos, _) call.
                let num_args = sig.param_names.len();
                let mut values = vec![0.0; num_args];
                let mut bindings = Vec::with_capacity(num_args);
                for (arg_pos, builtin_name) in sig.param_names.iter().enumerate() {
                    if let Some(&recipe_idx) = name_to_recipe_idx.get(*builtin_name) {
                        values[arg_pos] = param_defs[recipe_idx].default;
                        bindings.push((arg_pos, recipe_idx));
                    }
                }
                param_values = values;
                param_bindings = bindings;
            }

            // Build BuiltinInstance: one builtin wrapping the recipe's type
            let builtin_name = sig.name.to_string();
            let builtin_instance = rill_lang::ir::BuiltinInstance {
                name: builtin_name,
                params: param_values,
                kind: sig.kind,
                signal_ins: arity.0,
                signal_outs: arity.1,
                param_bindings,
            };

            // Build instructions: one LoadInput (if the builtin has signal inputs)
            // followed by CallBlock. Use separate registers for input/output when
            // both exist — avoids register aliasing in exec_foreign_block where
            // taking the output register would clobber the input.
            let mut instrs = Vec::new();
            let mut output_reg = 0usize;
            let mut num_regs = 1usize;
            if arity.1 > 0 {
                if arity.0 > 0 {
                    instrs.push(rill_lang::ir::Instr::LoadInput { dst: 0, index: 0 });
                    num_regs = 2;
                    output_reg = 1;
                }
                #[cfg(feature = "debug")]
                {
                    // ProbePoint needs an extra register: output_reg + 1
                    num_regs += 1;
                }
                let srcs = if arity.0 > 0 { vec![0] } else { vec![] };
                instrs.push(rill_lang::ir::Instr::CallBlock {
                    dst: output_reg,
                    srcs,
                    instance: 0,
                });
                #[cfg(feature = "debug")]
                instrs.push(rill_lang::ir::Instr::ProbePoint {
                    id: idx as u32,
                    src: output_reg,
                    dst: output_reg.wrapping_add(1),
                });
            }

            let ir = rill_lang::ir::Ir {
                instrs,
                num_regs,
                output_reg,
                num_inputs: arity.0,
                num_outputs: arity.1,
                state: rill_lang::ir::StateLayout {
                    state_slots: 0,
                    delay_lens: vec![],
                    num_outputs: arity.1,
                },
                builtins: vec![builtin_instance],
                params: param_defs.clone(),
            };

            nodes.insert(
                name.clone(),
                GraphNode {
                    arity,
                    ir,
                    params: param_defs,
                    keep: false,
                    inline: false,
                    is_bridge: false,
                    feedback_read: vec![],
                    feedback_write: vec![],
                },
            );
        }

        // 3. Convert edges
        let mut edges = Vec::new();
        for (from_idx, from_port, to_idx, to_port) in &self.signal_edges {
            edges.push(GraphEdge {
                from_node: idx_to_name[from_idx].clone(),
                from_port: *from_port,
                to_node: idx_to_name[to_idx].clone(),
                to_port: *to_port,
                kind: EdgeKind::Signal,
            });
        }
        for (from_idx, from_port, to_idx, to_port) in &self.feedback_edges {
            edges.push(GraphEdge {
                from_node: idx_to_name[from_idx].clone(),
                from_port: *from_port,
                to_node: idx_to_name[to_idx].clone(),
                to_port: *to_port,
                kind: EdgeKind::Feedback,
            });
        }

        // 4. Compute topological order (Kahn's algorithm on signal edges only)
        let mut in_degree: HashMap<String, usize> = HashMap::new();
        for name in &node_list {
            in_degree.insert(name.clone(), 0);
        }
        for edge in &edges {
            if edge.kind == EdgeKind::Signal {
                *in_degree.get_mut(&edge.to_node).unwrap() += 1;
            }
        }

        let mut adj: HashMap<String, Vec<String>> = HashMap::new();
        for name in &node_list {
            adj.insert(name.clone(), vec![]);
        }
        for edge in &edges {
            if edge.kind == EdgeKind::Signal {
                adj.get_mut(&edge.from_node)
                    .unwrap()
                    .push(edge.to_node.clone());
            }
        }

        let mut queue: Vec<String> = in_degree
            .iter()
            .filter(|(_, &d)| d == 0)
            .map(|(n, _)| n.clone())
            .collect();
        let mut topo_order = Vec::new();

        while let Some(node) = queue.pop() {
            topo_order.push(node.clone());
            if let Some(neighbors) = adj.get(&node) {
                for neighbor in neighbors {
                    let deg = in_degree.get_mut(neighbor).unwrap();
                    *deg -= 1;
                    if *deg == 0 {
                        queue.push(neighbor.clone());
                    }
                }
            }
        }

        if topo_order.len() != node_list.len() {
            return Err(BuildError::CycleDetected);
        }

        // 5. Compute graph-level inputs/outputs from root/leaf nodes
        let inputs = topo_order
            .iter()
            .filter(|n| in_degree.get(*n) == Some(&0))
            .map(|n| nodes[n].arity.0)
            .sum();
        let outputs = topo_order
            .iter()
            .filter(|n| adj.get(*n).is_none_or(|v| v.is_empty()))
            .map(|n| nodes[n].arity.1)
            .sum();

        Ok(GraphIr {
            inputs,
            outputs,
            nodes,
            edges,
            topo_order,
        })
    }

    /// Convert the graph to an rill-lang AST `Program`.
    ///
    /// Each graph node becomes an [`Expr::Apply`](rill_lang::ast::Expr::Apply) with parameters ordered
    /// according to the builtin's `BuiltinSig::param_names`. Nodes are
    /// chained via [`BinOp::Seq`](rill_lang::ast::BinOp::Seq) according to their signal connections.
    ///
    /// Only simple chain topologies are supported (fan-out/fan-in will
    /// return [`BuildError::UnsupportedTopology`]).
    pub fn ast_from_def(
        &self,
        registry: &rill_lang::builtin::Registry<T>,
    ) -> Result<rill_lang::ast::Program, BuildError> {
        use rill_lang::ast::{BinOp, Def, Expr, Param, Program};
        use rill_lang::error::Span;

        let dummy = Span::new(0, 0);

        // Build id-to-index mapping
        let mut id_to_idx: HashMap<u32, usize> = HashMap::new();
        for (i, r) in self.recipes.iter().enumerate() {
            id_to_idx.insert(r.id, i);
        }

        // Resolve builtin names and parameter order for each recipe
        struct NodeMeta {
            builtin_name: String,
            param_values: Vec<f64>,
            param_names: Vec<String>,
        }

        let mut node_metas: Vec<NodeMeta> = Vec::with_capacity(self.recipes.len());

        for recipe in &self.recipes {
            let builtin_name = Self::resolve_builtin_name(&recipe.type_name, registry)
                .ok_or_else(|| BuildError::UnknownNodeType(recipe.type_name.clone()))?;

            let sig = registry.builtin_sig(&builtin_name).unwrap();

            // Build parameter values in builtin param_names order
            let param_names: Vec<String> = sig.param_names.iter().map(|n| n.to_string()).collect();
            let mut param_values = Vec::with_capacity(param_names.len());

            // Build a lookup from recipe param name to f64 value
            let recipe_defaults: HashMap<&str, f64> = recipe
                .params
                .parameters
                .iter()
                .filter_map(|(k, v)| v.as_f32().map(|f| (k.as_str(), f as f64)))
                .collect();

            for name in &param_names {
                let val = recipe_defaults.get(name.as_str()).copied().unwrap_or(0.0);
                param_values.push(val);
            }

            node_metas.push(NodeMeta {
                builtin_name,
                param_values,
                param_names,
            });
        }

        // Topological sort
        let mut in_degree: Vec<usize> = vec![0; self.recipes.len()];
        let mut adj: Vec<Vec<usize>> = vec![vec![]; self.recipes.len()];

        for (from_idx, _from_port, to_idx, _to_port) in &self.signal_edges {
            if *from_idx < self.recipes.len() && *to_idx < self.recipes.len() {
                adj[*from_idx].push(*to_idx);
                in_degree[*to_idx] += 1;
            }
        }

        let mut queue: Vec<usize> = (0..self.recipes.len())
            .filter(|i| in_degree[*i] == 0)
            .collect();
        let mut order: Vec<usize> = Vec::new();

        while let Some(u) = queue.pop() {
            order.push(u);
            for &v in &adj[u] {
                in_degree[v] -= 1;
                if in_degree[v] == 0 {
                    queue.push(v);
                }
            }
        }

        if order.len() != self.recipes.len() {
            return Err(BuildError::CycleDetected);
        }

        // Check for unsupported topologies
        for (i, targets) in adj.iter().enumerate() {
            if targets.len() > 1 {
                return Err(BuildError::UnsupportedTopology(format!(
                    "node {} fans out to {} destinations (split not yet supported)",
                    i,
                    targets.len()
                )));
            }
            let in_count = self
                .signal_edges
                .iter()
                .filter(|(_, _, to, _)| *to == i)
                .count();
            if in_count > 1 {
                return Err(BuildError::UnsupportedTopology(format!(
                    "node {} receives {} signal inputs (merge not yet supported)",
                    i, in_count
                )));
            }
        }

        // Build AST expressions for each node in topo order
        // Map recipe index → AST expression
        let mut node_exprs: Vec<Option<Expr>> = vec![None; self.recipes.len()];

        // Collect all parameter names for the main definition
        let mut all_param_names: Vec<String> = Vec::new();

        for &idx in &order {
            let meta = &node_metas[idx];

            // Find upstream signal connection
            let upstream_expr: Option<Expr> = self
                .signal_edges
                .iter()
                .find(|(_, _, to, _)| *to == idx)
                .and_then(|(from, _from_port, _to, _to_port)| node_exprs[*from].clone());

            // Build args: Float for static (first) params, Ref for dynamic (last) param.
            // Only expose the last (dynamic) param as a main definition parameter.
            //
            // Convention: the last param in builtin param_names is the SetParameter target.
            let mut args: Vec<Expr> = Vec::new();
            let n = meta.param_names.len();
            for (i, (&val, name)) in meta
                .param_values
                .iter()
                .zip(meta.param_names.iter())
                .enumerate()
            {
                if i < n - 1 {
                    // Static param: put Float constant, no main parameter
                    args.push(Expr::Float(val, dummy));
                } else {
                    // Dynamic param: use Ref + register on main definition
                    all_param_names.push(name.clone());
                    args.push(Expr::Ref(name.clone(), dummy));
                }
            }

            let apply = Expr::Apply {
                name: meta.builtin_name.clone(),
                args,
                span: dummy,
            };

            let expr = match upstream_expr {
                Some(up) => Expr::Bin {
                    op: BinOp::Seq,
                    lhs: Box::new(up),
                    rhs: Box::new(apply),
                    span: dummy,
                },
                None => apply,
            };

            node_exprs[idx] = Some(expr);
        }

        // Find the last node (sink/leaf) — the one with no downstream edges
        let leaf: usize = order
            .iter()
            .rfind(|&&i| adj[i].is_empty())
            .copied()
            .unwrap_or(0);

        let body = node_exprs[leaf].clone().unwrap_or(Expr::Wire(dummy));

        let params: Vec<Param> = all_param_names
            .into_iter()
            .map(|name| Param { name, span: dummy })
            .collect();

        Ok(Program {
            defs: vec![Def::Anchor {
                name: "main".to_string(),
                params,
                body,
                span: dummy,
                where_defs: vec![],
            }],
        })
    }

    /// Compile directly from the graph definition to a `CompiledGraphEngine`.
    ///
    /// Calls [`ast_from_def`](Self::ast_from_def) followed by rill-lang compilation.
    pub fn compile_def<const BUF: usize>(
        &self,
        registry: &rill_lang::builtin::Registry<T>,
        sample_rate: f32,
    ) -> Result<rill_lang::graph_engine::CompiledGraphEngine<T, BUF>, BuildError> {
        let program = self.ast_from_def(registry)?;
        rill_lang::compile_program::<T, BUF>(&program, registry, sample_rate)
            .map_err(|e| BuildError::CompilationFailed(format!("{e}")))
    }

    fn resolve_builtin_name(
        type_name: &str,
        registry: &rill_lang::builtin::Registry<T>,
    ) -> Option<String> {
        if registry.builtin_sig(type_name).is_some() {
            return Some(type_name.to_string());
        }
        if let Some(rest) = type_name.strip_prefix("rill/") {
            if registry.builtin_sig(rest).is_some() {
                return Some(rest.to_string());
            }
        }
        let mapped = match type_name {
            "rill/dry_wet_mix" => "dry_wet",
            "rill/parametric_eq" => "eq_parametric",
            "rill/graphic_eq" => "graphic_eq",
            "rill/mono_to_stereo" => "mono_to_stereo",
            "rill/moog_ladder" => "moog",
            "rill/write_head" => "write_head",
            "rill/read_head" => "read_head",
            "rill/lofi_chip" => "ay38910",
            _ => "",
        };
        if !mapped.is_empty() && registry.builtin_sig(mapped).is_some() {
            return Some(mapped.to_string());
        }
        None
    }
}