//! Phase 3c — language-level dispatch for `grad_graph_*` builtins.
//!
//! Routed from both `cjc-eval` and `cjc-mir-exec` *after* the shared
//! `cjc_runtime::dispatch_builtin` and after `cjc_quantum::dispatch_quantum`,
//! so user `.cjcl` code can construct GradGraph nodes, run forward/backward,
//! and read tensors — all from the language.
//!
//! # Handle representation
//!
//! Phase 3c picks **Option B** from the brief: a single ambient `GradGraph`
//! per execution thread, addressed by `usize` node index. The graph lives in
//! a `thread_local!` so both executors share the same instance within a
//! single test thread, which is exactly what AST↔MIR parity testing needs.
//!
//! Node indices cross the language boundary as `Value::Int(i64)`. This avoids
//! introducing a new `Value` variant, preserving `Value` enum layout and the
//! MIR register layout (HARD RULE #1 from the Phase 3c brief).
//!
//! # Why a satellite dispatch and not `cjc-runtime/builtins.rs`?
//!
//! `cjc-ad` already depends on `cjc-runtime`. Adding the reverse dependency
//! would create a cycle. The `cjc-quantum::dispatch_quantum` precedent shows
//! satellite dispatch is the canonical pattern for this exact case.

use std::cell::RefCell;
use std::rc::Rc;

use cjc_runtime::value::Value;
use cjc_runtime::tensor::Tensor;

use crate::idx::NodeIdx;
use crate::GradGraph;
use crate::pinn::Activation;

thread_local! {
    /// Ambient graph for the current thread. Reset by `grad_graph_new()`.
    static AMBIENT: RefCell<GradGraph> = RefCell::new(GradGraph::new());
}

/// Run a closure with mutable access to the ambient graph.
///
/// Borrow is held only for the duration of the closure, which never re-enters
/// `dispatch_grad_graph` — so nested-borrow panics are structurally impossible.
pub fn with_ambient<R>(f: impl FnOnce(&mut GradGraph) -> R) -> R {
    AMBIENT.with(|cell| f(&mut *cell.borrow_mut()))
}

/// Reset the ambient graph to a fresh empty state.
///
/// Called by `grad_graph_new()`. Tests should call this at the top of any
/// `.cjcl` source that constructs a graph, since thread-locals can outlive
/// individual test functions when cargo runs tests on a thread pool.
pub fn reset_ambient() {
    AMBIENT.with(|cell| {
        *cell.borrow_mut() = GradGraph::new();
    });
}

// ─── Argument helpers ──────────────────────────────────────────────────────

fn arg_count(name: &str, args: &[Value], expected: usize) -> Result<(), String> {
    if args.len() != expected {
        Err(format!(
            "{}: expected {} arguments, got {}",
            name,
            expected,
            args.len()
        ))
    } else {
        Ok(())
    }
}

/// Decode a `Value::Int(i64)` into a typed [`NodeIdx`].
///
/// **Phase 2a typed boundary:** the dispatch layer speaks `NodeIdx` end
/// to end; the underlying `GradGraph` still takes `usize`, so each arm
/// converts via `idx.index()` at the call site. A future Phase 2b PR
/// will migrate `GradGraph` itself, at which point this function's
/// return type is the API the rest of the dispatch already uses.
fn arg_idx(name: &str, val: &Value) -> Result<NodeIdx, String> {
    match val {
        Value::Int(i) if *i >= 0 => Ok(NodeIdx::from_usize(*i as usize)),
        Value::Int(i) => Err(format!("{}: node index must be non-negative, got {}", name, i)),
        other => Err(format!(
            "{}: expected Int node index, got {}",
            name,
            other.type_name()
        )),
    }
}

/// Like `arg_idx`, but additionally asserts that the index is in range for
/// the current ambient graph. Returning `Err` here is what makes the
/// dispatch layer fuzz-safe: out-of-range indices surface as language-level
/// errors rather than panicking deep inside `GradGraph::*` index ops.
fn arg_idx_checked(name: &str, val: &Value) -> Result<NodeIdx, String> {
    let idx = arg_idx(name, val)?;
    let len = with_ambient(|g| g.len());
    if idx.index() >= len {
        return Err(format!(
            "{}: node index {} out of range (graph has {} nodes)",
            name, idx, len
        ));
    }
    Ok(idx)
}

fn arg_tensor<'a>(name: &str, val: &'a Value) -> Result<&'a Tensor, String> {
    match val {
        Value::Tensor(t) => Ok(t),
        other => Err(format!(
            "{}: expected Tensor, got {}",
            name,
            other.type_name()
        )),
    }
}

fn arg_f64(name: &str, val: &Value) -> Result<f64, String> {
    match val {
        Value::Float(f) => Ok(*f),
        Value::Int(i) => Ok(*i as f64),
        other => Err(format!(
            "{}: expected number, got {}",
            name,
            other.type_name()
        )),
    }
}

/// Decode `Value::Array` of `Value::Int` (each non-negative) into a
/// `Vec<usize>`, **and bound-check each entry against the ambient graph
/// length**. This is the right helper for arrays of *node indices*
/// (Phase 3b's `cat`, `backward_collect`); shape and gather-index arrays
/// use the lighter `arg_usize_array` since they're not graph-relative.
fn arg_idx_array(name: &str, val: &Value) -> Result<Vec<usize>, String> {
    let raw = arg_usize_array(name, val)?;
    let len = with_ambient(|g| g.len());
    for (i, &idx) in raw.iter().enumerate() {
        if idx >= len {
            return Err(format!(
                "{name}: node index [{i}]={idx} out of range (graph has {len} nodes)"
            ));
        }
    }
    Ok(raw)
}

/// Decode `Value::Array` of `Value::Int` (each ≥ 0) into a `Vec<usize>`.
/// Used by shape and gather-index arguments where bounds are op-specific
/// rather than graph-relative.
fn arg_usize_array(name: &str, val: &Value) -> Result<Vec<usize>, String> {
    match val {
        Value::Array(elems) => {
            let mut out = Vec::with_capacity(elems.len());
            for (i, e) in elems.iter().enumerate() {
                match e {
                    Value::Int(n) if *n >= 0 => out.push(*n as usize),
                    Value::Int(n) => {
                        return Err(format!(
                            "{name}: shape element [{i}] must be non-negative, got {n}"
                        ))
                    }
                    other => {
                        return Err(format!(
                            "{name}: shape element [{i}] must be Int, got {}",
                            other.type_name()
                        ))
                    }
                }
            }
            Ok(out)
        }
        other => Err(format!(
            "{name}: expected Array of Int (shape), got {}",
            other.type_name()
        )),
    }
}

fn arg_str<'a>(name: &str, val: &'a Value) -> Result<&'a str, String> {
    match val {
        Value::String(s) => Ok(s.as_str()),
        other => Err(format!(
            "{}: expected String, got {}",
            name,
            other.type_name()
        )),
    }
}

fn parse_activation(name: &str, s: &str) -> Result<Activation, String> {
    match s {
        "tanh" => Ok(Activation::Tanh),
        "sigmoid" => Ok(Activation::Sigmoid),
        "relu" => Ok(Activation::Relu),
        "none" => Ok(Activation::None),
        "gelu" => Ok(Activation::Gelu),
        "silu" => Ok(Activation::Silu),
        "elu" => Ok(Activation::Elu),
        "selu" => Ok(Activation::Selu),
        "sin" => Ok(Activation::SinAct),
        other => Err(format!(
            "{}: unknown activation {:?} (expected one of: tanh sigmoid relu none gelu silu elu selu sin)",
            name, other
        )),
    }
}

/// Convert a typed [`NodeIdx`] back to a `Value::Int(i64)` for return
/// across the language boundary. Construction methods on `GradGraph`
/// currently return `usize`; arms therefore wrap with
/// `idx_value(NodeIdx::from_usize(g.input(t)))` until the bulk
/// `GradGraph` migration in Phase 2b.
fn idx_value(idx: NodeIdx) -> Value {
    Value::Int(idx.0 as i64)
}

fn tensor_value(t: Tensor) -> Value {
    Value::Tensor(t)
}

// ─── Dispatch ─────────────────────────────────────────────────────────────

/// Dispatch table for language-level `grad_graph_*` builtins.
///
/// Returns `Ok(None)` if `name` is not one of ours, so callers can fall
/// through to other dispatch tables.
pub fn dispatch_grad_graph(name: &str, args: &[Value]) -> Result<Option<Value>, String> {
    let result: Value = match name {
        // ── Construction ────────────────────────────────────────────
        "grad_graph_new" => {
            arg_count(name, args, 0)?;
            reset_ambient();
            Value::Void
        }
        "grad_graph_param" => {
            arg_count(name, args, 1)?;
            let t = arg_tensor(name, &args[0])?.clone();
            let idx = NodeIdx::from_usize(with_ambient(|g| g.parameter(t)));
            idx_value(idx)
        }
        "grad_graph_input" => {
            arg_count(name, args, 1)?;
            let t = arg_tensor(name, &args[0])?.clone();
            let idx = NodeIdx::from_usize(with_ambient(|g| g.input(t)));
            idx_value(idx)
        }
        // const = non-trainable input; same node behavior, kept as a separate
        // name so user code documents intent.
        "grad_graph_const" => {
            arg_count(name, args, 1)?;
            let t = arg_tensor(name, &args[0])?.clone();
            let idx = NodeIdx::from_usize(with_ambient(|g| g.input(t)));
            idx_value(idx)
        }

        // ── Pointwise ops ───────────────────────────────────────────
        // Phase 2a typed boundary: arg_idx_checked returns NodeIdx; the
        // underlying GradGraph still takes usize, hence `.index()` at
        // each call site. Phase 2b will migrate GradGraph itself.
        "grad_graph_add" => {
            arg_count(name, args, 2)?;
            let a = arg_idx_checked(name, &args[0])?;
            let b = arg_idx_checked(name, &args[1])?;
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.add(a.index(), b.index()))))
        }
        "grad_graph_sub" => {
            arg_count(name, args, 2)?;
            let a = arg_idx_checked(name, &args[0])?;
            let b = arg_idx_checked(name, &args[1])?;
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.sub(a.index(), b.index()))))
        }
        "grad_graph_mul" => {
            arg_count(name, args, 2)?;
            let a = arg_idx_checked(name, &args[0])?;
            let b = arg_idx_checked(name, &args[1])?;
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.mul(a.index(), b.index()))))
        }
        "grad_graph_div" => {
            arg_count(name, args, 2)?;
            let a = arg_idx_checked(name, &args[0])?;
            let b = arg_idx_checked(name, &args[1])?;
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.div(a.index(), b.index()))))
        }
        "grad_graph_neg" => {
            arg_count(name, args, 1)?;
            let a = arg_idx_checked(name, &args[0])?;
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.neg(a.index()))))
        }
        "grad_graph_scalar_mul" => {
            arg_count(name, args, 2)?;
            let a = arg_idx_checked(name, &args[0])?;
            let s = arg_f64(name, &args[1])?;
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.scalar_mul(a.index(), s))))
        }
        "grad_graph_pow" => {
            arg_count(name, args, 2)?;
            let a = arg_idx_checked(name, &args[0])?;
            // Brief specifies integer power, but cjc-ad's `pow` takes f64.
            // Accept either i64 or f64; both narrow to the same call.
            let n = arg_f64(name, &args[1])?;
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.pow(a.index(), n))))
        }
        "grad_graph_exp" => {
            arg_count(name, args, 1)?;
            let a = arg_idx_checked(name, &args[0])?;
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.exp(a.index()))))
        }
        "grad_graph_ln" => {
            arg_count(name, args, 1)?;
            let a = arg_idx_checked(name, &args[0])?;
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.ln(a.index()))))
        }
        "grad_graph_sqrt" => {
            arg_count(name, args, 1)?;
            let a = arg_idx_checked(name, &args[0])?;
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.sqrt(a.index()))))
        }
        "grad_graph_sin" => {
            arg_count(name, args, 1)?;
            let a = arg_idx_checked(name, &args[0])?;
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.sin(a.index()))))
        }
        "grad_graph_cos" => {
            arg_count(name, args, 1)?;
            let a = arg_idx_checked(name, &args[0])?;
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.cos(a.index()))))
        }
        "grad_graph_tanh" => {
            arg_count(name, args, 1)?;
            let a = arg_idx_checked(name, &args[0])?;
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.tanh_act(a.index()))))
        }

        // ── Reductions / matmul ─────────────────────────────────────
        "grad_graph_sum" => {
            arg_count(name, args, 1)?;
            let a = arg_idx_checked(name, &args[0])?;
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.sum(a.index()))))
        }
        "grad_graph_mean" => {
            arg_count(name, args, 1)?;
            let a = arg_idx_checked(name, &args[0])?;
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.mean(a.index()))))
        }
        "grad_graph_matmul" => {
            arg_count(name, args, 2)?;
            let a = arg_idx_checked(name, &args[0])?;
            let b = arg_idx_checked(name, &args[1])?;
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.matmul(a.index(), b.index()))))
        }

        // ── Fused MLP layer ────────────────────────────────────────
        // Signature: grad_graph_mlp_layer(input, weight, bias, activation_str) -> NodeIdx
        "grad_graph_mlp_layer" => {
            arg_count(name, args, 4)?;
            let inp = arg_idx_checked(name, &args[0])?;
            let w = arg_idx_checked(name, &args[1])?;
            let b = arg_idx_checked(name, &args[2])?;
            let act_str = arg_str(name, &args[3])?;
            let act = parse_activation(name, act_str)?;
            idx_value(NodeIdx::from_usize(with_ambient(|g| {
                g.mlp_layer(inp.index(), w.index(), b.index(), act)
            })))
        }

        // ── Forward / state read ────────────────────────────────────
        "grad_graph_forward" => {
            // The graph is forward-evaluated eagerly on each op call, so
            // forward() at a node is just a tensor read. Kept under this name
            // for spec parity with the brief.
            arg_count(name, args, 1)?;
            let idx = arg_idx_checked(name, &args[0])?;
            tensor_value(with_ambient(|g| g.tensor(idx.index())))
        }
        "grad_graph_set_tensor" => {
            arg_count(name, args, 2)?;
            let idx = arg_idx_checked(name, &args[0])?;
            let t = arg_tensor(name, &args[1])?.clone();
            with_ambient(|g| g.set_tensor(idx.index(), t));
            Value::Void
        }
        "grad_graph_param_grad" => {
            arg_count(name, args, 1)?;
            let idx = arg_idx_checked(name, &args[0])?;
            let g_opt = with_ambient(|g| g.grad(idx.index()));
            match g_opt {
                Some(t) => tensor_value(t),
                None => {
                    return Err(format!(
                        "grad_graph_param_grad: node {} has no gradient (not a parameter, or backward not run)",
                        idx
                    ))
                }
            }
        }

        // ── Backward / optimizer support ────────────────────────────
        "grad_graph_zero_grad" => {
            arg_count(name, args, 0)?;
            with_ambient(|g| g.zero_grad());
            Value::Void
        }
        "grad_graph_backward" => {
            arg_count(name, args, 1)?;
            let loss_idx = arg_idx_checked(name, &args[0])?;
            // Backward requires a scalar leaf. If the loss node is non-scalar
            // (e.g., an unreduced add/mul), the inner GradGraph would assert.
            // Surface that as a clean Err instead of unwinding.
            let shape = with_ambient(|g| g.tensor(loss_idx.index()).shape().to_vec());
            let total: usize = shape.iter().product();
            if total != 1 {
                return Err(format!(
                    "grad_graph_backward: loss node {} has shape {:?} (numel={}), but backward requires a scalar (numel=1). Apply grad_graph_sum or grad_graph_mean first.",
                    loss_idx, shape, total
                ));
            }
            with_ambient(|g| g.backward(loss_idx.index()));
            Value::Void
        }
        "grad_graph_clip_grad_norm" => {
            arg_count(name, args, 1)?;
            let max_norm = arg_f64(name, &args[0])?;
            let actual = with_ambient(|g| g.clip_grad_norm(max_norm));
            Value::Float(actual)
        }

        // ── Phase 3a: transformer-backbone activations & ops ────────
        // All map directly to existing `cjc_ad::GradGraph` Rust methods.
        // Determinism contract is unchanged: `softmax`, `cross_entropy`,
        // and `layer_norm` use Kahan accumulators internally; `gelu`,
        // `silu`, `reshape` are pointwise / shape-only and inherit the
        // tensor kernel's bit-exact behavior.

        "grad_graph_softmax" => {
            arg_count(name, args, 1)?;
            let a = arg_idx_checked(name, &args[0])?;
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.softmax(a.index()))))
        }
        "grad_graph_cross_entropy" => {
            arg_count(name, args, 2)?;
            let logits = arg_idx_checked(name, &args[0])?;
            let targets = arg_idx_checked(name, &args[1])?;
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.cross_entropy(logits.index(), targets.index()))))
        }
        "grad_graph_layer_norm" => {
            arg_count(name, args, 1)?;
            let a = arg_idx_checked(name, &args[0])?;
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.layer_norm(a.index()))))
        }
        "grad_graph_gelu" => {
            arg_count(name, args, 1)?;
            let a = arg_idx_checked(name, &args[0])?;
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.gelu(a.index()))))
        }
        "grad_graph_silu" => {
            arg_count(name, args, 1)?;
            let a = arg_idx_checked(name, &args[0])?;
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.silu(a.index()))))
        }
        // grad_graph_reshape(node, shape_array) -> NodeIdx
        // shape_array is a Value::Array of Value::Int (positive).
        "grad_graph_reshape" => {
            arg_count(name, args, 2)?;
            let a = arg_idx_checked(name, &args[0])?;
            let shape = arg_usize_array(name, &args[1])?;
            // Materialize new tensor; shape mismatch surfaces as a clean Err
            // rather than the inner `expect()` panic.
            let cur_numel: usize = with_ambient(|g| g.tensor(a.index()).shape().iter().product());
            let new_numel: usize = shape.iter().product();
            if new_numel != cur_numel {
                return Err(format!(
                    "grad_graph_reshape: cannot reshape tensor with {cur_numel} elements into shape {shape:?} ({new_numel} elements)"
                ));
            }
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.reshape(a.index(), &shape))))
        }

        // ── Phase 3b: array-arg & state-recovery ops ────────────────
        // - batch_norm: trivial 1-arg activation, completes the
        //   normalization pair with layer_norm
        // - gather/cat: take Array<Int>; reuse `arg_usize_array` for the
        //   non-graph-relative shape/index array, with bounds-checking
        //   added at this boundary so OOB doesn't panic in the kernel
        // - reforward: takes (start, end) range; recomputes forward-pass
        //   tensors so users can `set_tensor()` a parameter and re-run
        //   the graph in place (the v2.5 PINN-graph-reuse path)
        // - backward_collect: zero_grad + backward + gather grads in one
        //   call; returns Value::Array of Tensors. None entries (i.e. a
        //   passed index that isn't a parameter or has no gradient yet)
        //   are surfaced as Err — the user explicitly listed which nodes
        //   they expect gradients for, so a None means a contract bug

        "grad_graph_batch_norm" => {
            arg_count(name, args, 1)?;
            let a = arg_idx_checked(name, &args[0])?;
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.batch_norm(a.index()))))
        }

        // grad_graph_gather(node, indices_array, axis) -> NodeIdx
        "grad_graph_gather" => {
            arg_count(name, args, 3)?;
            let a = arg_idx_checked(name, &args[0])?;
            let indices = arg_usize_array(name, &args[1])?;
            // axis is a tensor-axis number, not a graph-relative node id;
            // strip the typed wrapper as soon as it crosses the boundary.
            let axis = arg_idx(name, &args[2])?.index();
            // Bounds-check axis and indices against the input tensor shape
            // so OOB doesn't panic deep in `Tensor` indexing.
            let (ndim, axis_dim) = with_ambient(|g| {
                let t = g.tensor(a.index());
                let s = t.shape();
                (s.len(), if axis < s.len() { Some(s[axis]) } else { None })
            });
            if axis >= ndim {
                return Err(format!(
                    "grad_graph_gather: axis {axis} out of range (tensor has {ndim} dimensions)"
                ));
            }
            let dim_size = axis_dim.unwrap();
            for (i, &ix) in indices.iter().enumerate() {
                if ix >= dim_size {
                    return Err(format!(
                        "grad_graph_gather: index [{i}]={ix} out of range for axis {axis} (size {dim_size})"
                    ));
                }
            }
            idx_value(NodeIdx::from_usize(
                with_ambient(|g| g.gather(a.index(), &indices, axis)),
            ))
        }

        // grad_graph_cat(node_array, axis) -> NodeIdx
        "grad_graph_cat" => {
            arg_count(name, args, 2)?;
            let inputs = arg_idx_array(name, &args[0])?;
            // axis is a tensor-axis number, not a graph-relative node id.
            let axis = arg_idx(name, &args[1])?.index();
            if inputs.is_empty() {
                return Err(
                    "grad_graph_cat: input list must contain at least one tensor".to_string(),
                );
            }
            // Validate that axis is in range for the first input, and
            // that all inputs share the same ndim. Per-axis size mismatch
            // is handled by the kernel.
            let (ndim, _shape0) = with_ambient(|g| {
                let s = g.tensor(inputs[0]).shape().to_vec();
                (s.len(), s)
            });
            if axis >= ndim {
                return Err(format!(
                    "grad_graph_cat: axis {axis} out of range (tensors have {ndim} dimensions)"
                ));
            }
            for (i, &idx) in inputs.iter().enumerate().skip(1) {
                let other_ndim = with_ambient(|g| g.tensor(idx).shape().len());
                if other_ndim != ndim {
                    return Err(format!(
                        "grad_graph_cat: input [{i}] has {other_ndim} dimensions, expected {ndim}"
                    ));
                }
            }
            idx_value(NodeIdx::from_usize(with_ambient(|g| g.cat(&inputs, axis))))
        }

        // grad_graph_reforward(start, end) -> Void
        // Recomputes forward-pass tensors for ops in the inclusive range
        // [start, end]. Used by PINN-style graph reuse: set_tensor() a
        // parameter, then reforward to update intermediate values before
        // calling backward again.
        "grad_graph_reforward" => {
            arg_count(name, args, 2)?;
            let start = arg_idx_checked(name, &args[0])?;
            let end = arg_idx_checked(name, &args[1])?;
            if start > end {
                return Err(format!(
                    "grad_graph_reforward: start ({start}) must be ≤ end ({end})"
                ));
            }
            with_ambient(|g| g.reforward(start.index(), end.index()));
            Value::Void
        }

        // grad_graph_backward_collect(loss, param_indices_array) -> Array<Tensor>
        "grad_graph_backward_collect" => {
            arg_count(name, args, 2)?;
            let loss_idx = arg_idx_checked(name, &args[0])?;
            let param_indices = arg_idx_array(name, &args[1])?;
            // Same scalar-shape guard as grad_graph_backward — backward
            // requires a numel=1 leaf.
            let shape = with_ambient(|g| g.tensor(loss_idx.index()).shape().to_vec());
            let total: usize = shape.iter().product();
            if total != 1 {
                return Err(format!(
                    "grad_graph_backward_collect: loss node {loss_idx} has shape {shape:?} (numel={total}), but backward requires a scalar (numel=1). Apply grad_graph_sum or grad_graph_mean first."
                ));
            }
            let grads = with_ambient(|g| g.backward_collect(loss_idx.index(), &param_indices));
            // Surface a None entry as Err — the caller explicitly named
            // which nodes they expect gradients for, so a None means
            // either a non-parameter was listed or backward was not run
            // yet (the latter shouldn't happen because backward_collect
            // calls backward internally).
            let mut out: Vec<Value> = Vec::with_capacity(grads.len());
            for (i, g_opt) in grads.into_iter().enumerate() {
                match g_opt {
                    Some(t) => out.push(Value::Tensor(t)),
                    None => {
                        return Err(format!(
                            "grad_graph_backward_collect: param_indices[{i}]={} has no gradient (not a parameter, or unreachable from loss)",
                            param_indices[i]
                        ));
                    }
                }
            }
            Value::Array(std::rc::Rc::new(out))
        }

        // ── Introspection ───────────────────────────────────────────
        "grad_graph_len" => {
            arg_count(name, args, 0)?;
            Value::Int(with_ambient(|g| g.len()) as i64)
        }

        // Not one of ours — let the caller fall through.
        _ => return Ok(None),
    };

    Ok(Some(result))
}

// Make `Rc` reachable to suppress unused-import warning when no Rc paths
// are exercised in this module today; it stays available for future arms
// that need to wrap collections.
#[allow(dead_code)]
fn _rc_marker() -> Rc<()> {
    Rc::new(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use cjc_runtime::tensor::Tensor;

    fn t(data: &[f64], shape: &[usize]) -> Value {
        Value::Tensor(Tensor::from_vec(data.to_vec(), shape).unwrap())
    }

    #[test]
    fn new_resets_ambient() {
        // Build something, reset, confirm len drops to 0.
        let _ = dispatch_grad_graph("grad_graph_new", &[]).unwrap();
        let _ = dispatch_grad_graph(
            "grad_graph_param",
            &[t(&[1.0, 2.0], &[2])],
        )
        .unwrap();
        let len_before = match dispatch_grad_graph("grad_graph_len", &[])
            .unwrap()
            .unwrap()
        {
            Value::Int(n) => n,
            _ => panic!("expected Int"),
        };
        assert!(len_before >= 1);
        let _ = dispatch_grad_graph("grad_graph_new", &[]).unwrap();
        let len_after = match dispatch_grad_graph("grad_graph_len", &[])
            .unwrap()
            .unwrap()
        {
            Value::Int(n) => n,
            _ => panic!("expected Int"),
        };
        assert_eq!(len_after, 0);
    }

    #[test]
    fn add_two_params_then_backward_returns_unit_grads() {
        let _ = dispatch_grad_graph("grad_graph_new", &[]).unwrap();
        let a = match dispatch_grad_graph("grad_graph_param", &[t(&[3.0], &[1])])
            .unwrap()
            .unwrap()
        {
            Value::Int(i) => i,
            _ => panic!(),
        };
        let b = match dispatch_grad_graph("grad_graph_param", &[t(&[4.0], &[1])])
            .unwrap()
            .unwrap()
        {
            Value::Int(i) => i,
            _ => panic!(),
        };
        let s = match dispatch_grad_graph(
            "grad_graph_add",
            &[Value::Int(a), Value::Int(b)],
        )
        .unwrap()
        .unwrap()
        {
            Value::Int(i) => i,
            _ => panic!(),
        };
        // Forward value: 3 + 4 = 7
        let s_t = match dispatch_grad_graph("grad_graph_forward", &[Value::Int(s)])
            .unwrap()
            .unwrap()
        {
            Value::Tensor(t) => t,
            _ => panic!(),
        };
        assert_eq!(s_t.to_vec(), vec![7.0]);
        // Backward: dS/da = dS/db = 1
        let _ = dispatch_grad_graph("grad_graph_zero_grad", &[]).unwrap();
        let _ = dispatch_grad_graph("grad_graph_backward", &[Value::Int(s)]).unwrap();
        let ga = match dispatch_grad_graph("grad_graph_param_grad", &[Value::Int(a)])
            .unwrap()
            .unwrap()
        {
            Value::Tensor(t) => t,
            _ => panic!(),
        };
        let gb = match dispatch_grad_graph("grad_graph_param_grad", &[Value::Int(b)])
            .unwrap()
            .unwrap()
        {
            Value::Tensor(t) => t,
            _ => panic!(),
        };
        assert_eq!(ga.to_vec(), vec![1.0]);
        assert_eq!(gb.to_vec(), vec![1.0]);
    }

    #[test]
    fn unknown_name_returns_none() {
        let r = dispatch_grad_graph("not_a_grad_graph_builtin", &[]).unwrap();
        assert!(r.is_none());
    }
}
