darkly 0.5.0

//! Node-graph composable brush engine.

pub mod builtin_brushes;
pub mod bundle;
pub mod checkpoint_ring;
pub mod composite_pipeline;
pub mod curve_math;
pub mod eval;
pub mod gpu_context;
pub mod import;
pub mod interpolation;
pub mod library;
pub mod node;
pub mod nodes;
pub mod paint_info;
pub mod paint_target_ext;
pub mod pipeline;
pub mod portable;
pub mod preview_renderer;
pub mod save_points;
pub mod scratch;
pub mod spacing;
pub mod stabilizer;
pub mod stabilizers;
pub mod state;
pub mod stroke_buffer;
pub mod stroke_engine;
pub mod wgsl;
pub mod wire;

use std::collections::HashMap;
use std::sync::OnceLock;

use crate::nodegraph::NodeRegistration;
use wire::BrushWireType;

/// Reference dab dimension (width = height), in canvas pixels — the
/// UI/UX convention that defines what `brush size = 1.0` means in
/// canvas-pixel terms. Used by:
///
/// - `stamp`: the rendered dab's longer axis at `effective_size = 1.0`
///   is `DAB_REFERENCE_SIZE` canvas pixels (the slider's "100%" mark).
/// - `liquify`: `radius = size * DAB_REFERENCE_SIZE * 0.5`.
/// - `paint` / `watercolor`: same `effective_size *
///   DAB_REFERENCE_SIZE / 2` formula for `effective_radius`.
/// - `StrokeEngine::default_diameter() = DAB_REFERENCE_SIZE * 0.5` —
///   spacing fallback before any dab has reported its own size.
///
/// **Not a hard cap on dab footprint.** `effective_size` (= `size_input
/// × size`) can exceed 1.0 unbounded — both the slider and any wired
/// sensor (pen pressure, random, ...) can push past 100%. The system
/// must tolerate dab footprints arbitrarily larger than this value.
///
/// If you need a hard cap on something, introduce a separately-named
/// constant. Do not repurpose this one.
pub const DAB_REFERENCE_SIZE: u32 = 512;

/// Wrapper around `NodeRegistration<BrushWireType>` carrying the GPU
/// pipelines a brush node owns.  Used by `build.rs` auto-generation for
/// the nodes directory.
pub use node::BrushNodeRegistration;

/// HashMap-backed registry of all brush node types.
///
/// Populated from `nodes::registrations()` (auto-generated by `build.rs`).
/// Holds **static metadata** — port definitions, parameter schemas,
/// display names, GPU pipeline registrations, stroke-lifecycle hooks,
/// and an `evaluator: fn()` constructor per node type. The registry is
/// the single source of truth for "what nodes exist?": there is no
/// parallel hand-written evaluator map to keep in sync.
///
/// Stores a `BrushNodeRegistration` wrapper per type plus a parallel
/// `nodegraph_map` of just the bare `NodeRegistration<W>` for handing to
/// the nodegraph compiler, which is generic over `W` and doesn't know
/// about the brush-specific pipeline metadata. The registry itself is a
/// process-global `OnceLock` — see [`registry`].
pub struct BrushNodeRegistry {
    map: HashMap<String, BrushNodeRegistration>,
    nodegraph_map: HashMap<String, NodeRegistration<BrushWireType>>,
}

impl BrushNodeRegistry {
    /// Build the registry from the auto-discovered node modules.
    fn build() -> Self {
        let regs = nodes::registrations();
        let mut map = HashMap::with_capacity(regs.len());
        let mut nodegraph_map = HashMap::with_capacity(regs.len());
        for reg in regs {
            let type_id = reg.node.type_id.to_string();
            nodegraph_map.insert(type_id.clone(), reg.node.clone());
            map.insert(type_id, reg);
        }
        Self { map, nodegraph_map }
    }

    /// Look up a node registration by type_id.
    pub fn get(&self, type_id: &str) -> Option<&BrushNodeRegistration> {
        self.map.get(type_id)
    }

    /// All registered node types.
    pub fn types(&self) -> impl Iterator<Item = &BrushNodeRegistration> {
        self.map.values()
    }

    /// The bare `NodeRegistration<W>` map for the nodegraph compiler
    /// (which is generic over `W` and doesn't carry brush-specific
    /// pipeline metadata).
    pub fn as_map(&self) -> &HashMap<String, NodeRegistration<BrushWireType>> {
        &self.nodegraph_map
    }

    /// Build a fresh `type_id → Box<dyn BrushNodeEvaluator>` map by
    /// invoking each registration's evaluator constructor. Evaluators
    /// are trait objects (not `Clone`), so the registry can't memoize
    /// a single instance — the runner takes ownership of the returned
    /// map. Replaces the deleted `default_evaluators()` free function.
    pub fn evaluators(&self) -> HashMap<String, Box<dyn eval::BrushNodeEvaluator>> {
        let mut out = HashMap::with_capacity(self.map.len());
        for (type_id, reg) in &self.map {
            out.insert(type_id.clone(), (reg.evaluator)());
        }
        out
    }
}

/// Process-global brush-node registry. Built once on first access; every
/// subsequent caller borrows the same `&'static BrushNodeRegistry`. The
/// 20-plus former `BrushNodeRegistry::new()` sites all funnel through here
/// now, so node metadata is materialized once per process instead of per
/// graph compile / per-brush-builder.
pub fn registry() -> &'static BrushNodeRegistry {
    static REGISTRY: OnceLock<BrushNodeRegistry> = OnceLock::new();
    REGISTRY.get_or_init(BrushNodeRegistry::build)
}

/// Convenience over [`crate::nodegraph::Graph::find_terminal`] for
/// brush graphs: builds a fresh [`BrushNodeRegistry`] and delegates.
/// Use this from any graph-only call site (benches, tests, the WASM
/// bridge) where threading the registry through would be noise.
pub fn find_terminal(
    graph: &crate::nodegraph::Graph<BrushWireType>,
) -> Result<crate::nodegraph::NodeId, crate::nodegraph::FindTerminalError> {
    let registry = registry();
    graph.find_terminal(registry.as_map())
}

/// Build the default brush graph: pressure-sensitive disc through the
/// compiled `paint` terminal. Same shape as Round in
/// [`builtin_brushes`] — `pen → paint_color → shape (sine, amplitude 0)
/// → stamp → paint` — minus the per-brush configuration
/// closure (no exposed softness, no flow wire).
pub fn default_graph() -> crate::nodegraph::Graph<BrushWireType> {
    use crate::gpu::params::ParamValue;
    use crate::nodegraph::{Graph, PortRef};

    let registry = registry();
    let mut graph = Graph::new();

    let pen = graph.add_node(
        "pen_input",
        registry.get("pen_input").unwrap().ports.clone(),
        vec![],
    );
    let paint_color = graph.add_node(
        "paint_color",
        registry.get("paint_color").unwrap().ports.clone(),
        vec![],
    );
    let shape = graph.add_node(
        "shape",
        registry.get("shape").unwrap().ports.clone(),
        vec![ParamValue::Int(0)], // sine, amplitude 0 → plain disc
    );
    let stamp = graph.add_node(
        "stamp",
        registry.get("stamp").unwrap().ports.clone(),
        vec![],
    );
    let terminal = graph.add_node(
        "paint",
        registry.get("paint").unwrap().ports.clone(),
        vec![],
    );

    let wires = [
        (pen, "pressure", terminal, "flow"),
        (paint_color, "color", stamp, "color"),
        (shape, "mask", stamp, "tip"),
        (stamp, "dab", terminal, "rgba"),
        (pen, "position", terminal, "position"),
    ];
    for (fnode, fport, tnode, tport) in wires {
        graph
            .connect(
                PortRef {
                    node: fnode,
                    port: fport.into(),
                },
                PortRef {
                    node: tnode,
                    port: tport.into(),
                },
            )
            .unwrap();
    }

    graph
}

/// Compile any brush graph into a ready-to-run runner.
///
/// Generates the per-brush WGSL fragment shader and attaches it to
/// the runner when the graph has a terminal. Errors from the
/// topological compile or the WGSL compile (e.g. an upstream node
/// with no `compile_wgsl` implementation) are returned verbatim as
/// human-readable strings.
pub fn compile_graph(
    graph: &crate::nodegraph::Graph<BrushWireType>,
) -> Result<eval::BrushGraphRunner, String> {
    let registry = registry();
    // Rewrite Switch nodes on a throwaway clone before any downstream
    // consumer reads the graph. Every consumer below (runner, topo
    // compile, WGSL compile) must see the same rewritten shape so a
    // Switch never leaks past this point. The persisted graph is
    // untouched — save/load, undo, and the brush-builder UI keep seeing
    // the original placement.
    let rewritten = {
        let mut g = graph.clone();
        nodes::switch::apply_to(&mut g);
        g
    };
    let evaluators = registry.evaluators();
    let mut runner = eval::BrushGraphRunner::new(&rewritten, registry.as_map(), evaluators)
        .map_err(|e| format!("{e}"))?;
    if runner.has_terminal() {
        let plan =
            crate::nodegraph::compile(&rewritten, registry.as_map()).map_err(|e| format!("{e}"))?;
        // Build a fresh evaluators map for the compiler — the runner
        // owns the live one. Cheap (just trait-object constructors).
        let compile_evals = registry.evaluators();
        let compiled = wgsl::compile_brush_to_wgsl(&rewritten, &plan, &compile_evals)
            .map_err(|e| format!("paint WGSL compilation failed: {e}"))?;
        runner.set_compiled_brush(std::sync::Arc::new(compiled));
    }
    Ok(runner)
}

/// Compile the default brush graph into a ready-to-run runner.
pub fn default_runner() -> Result<eval::BrushGraphRunner, String> {
    compile_graph(&default_graph())
}

/// Deserialize a brush graph from JSON and compile it.
///
/// Returns the compiled runner or a human-readable error string.
pub fn compile_from_json(json: &str) -> Result<eval::BrushGraphRunner, String> {
    let graph: crate::nodegraph::Graph<BrushWireType> =
        serde_json::from_str(json).map_err(|e| format!("invalid graph JSON: {e}"))?;
    compile_graph(&graph)
}

/// Reset every exposed input port on every node back to its registration
/// default. This produces a "canonical" view of the graph that's
/// independent of any user-facing scrubs (size, opacity, hardness, etc.).
///
/// Used by the active-dab thumbnail render so the brush icon represents
/// the brush's identity (shape, texture, dynamics) rather than the user's
/// momentary parameter adjustments. Generalizes per-port pinning so adding
/// a new exposed scrub (opacity hotkey, hardness hotkey, …) requires no
/// changes here.
///
/// Ports flagged with [`crate::nodegraph::PortDef::persist_in_thumbnail`]
/// are exempt — orientation knobs (rotation, phase) visibly change the
/// dab itself, so the icon should reflect them.
pub fn reset_exposed_scrubs(graph: &mut crate::nodegraph::Graph<BrushWireType>) {
    let registry = registry();
    let mut resets: Vec<(crate::nodegraph::NodeId, String, f32)> = Vec::new();
    // Walk every node and consult the *registration*'s `.exposed()`
    // flag — only ports the node type declares as user-tunable by
    // default get neutralized. Author-added entries in
    // `graph.exposed_ports` (size, opacity, etc. that the author
    // surfaced themselves) are part of the brush's identity and stay
    // at the author's configured values for the thumbnail.
    for (id, node) in graph.nodes() {
        let Some(reg) = registry.get(&node.type_id) else {
            continue;
        };
        for port in &reg.ports {
            if port.exposed
                && port.dir == crate::nodegraph::PortDir::Input
                && !port.persist_in_thumbnail
            {
                resets.push((*id, port.name.clone(), port.default));
            }
        }
    }
    for (id, name, default) in resets {
        let _ = graph.set_port_default(id, &name, default);
    }
}

/// Validate a brush graph from JSON without compiling.
///
/// Returns Ok(()) or a human-readable error string.
pub fn validate_graph_json(json: &str) -> Result<(), String> {
    let mut graph: crate::nodegraph::Graph<BrushWireType> =
        serde_json::from_str(json).map_err(|e| format!("invalid graph JSON: {e}"))?;
    // Mirror compile_graph: Switch rewriting must happen before the topo
    // pass for validation to match what compilation will actually see.
    nodes::switch::apply_to(&mut graph);
    let registry = registry();
    crate::nodegraph::compile(&graph, registry.as_map())
        .map_err(|e| format!("graph validation failed: {e}"))?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::eval::BrushGraphRunner;
    use super::paint_info::{PaintInformation, StrokeRecord};
    use super::wire::{BrushWireType, ScalarValue};
    use super::*;
    use crate::gpu::params::ParamValue;
    use crate::nodegraph::{Graph, NodeId, PortRef};

    /// Helper: build a graph with pen_input.pressure → multiply.a,
    /// and multiply.b pinned to 0.5 via a port default (no extra node).
    fn pressure_times_half_graph() -> (Graph<BrushWireType>, NodeId, NodeId) {
        let registry = registry();
        let mut graph = Graph::new();

        let pen_reg = registry.get("pen_input").unwrap();
        let pen = graph.add_node("pen_input", pen_reg.ports.clone(), vec![]);

        let mul_reg = registry.get("multiply").unwrap();
        let multiply = graph.add_node("multiply", mul_reg.ports.clone(), vec![]);

        // pen.pressure → multiply.a
        graph
            .connect(
                PortRef {
                    node: pen,
                    port: "pressure".into(),
                },
                PortRef {
                    node: multiply,
                    port: "a".into(),
                },
            )
            .unwrap();

        // multiply.b defaults to 0.5 — no wire, just a port default.
        graph.set_port_default(multiply, "b", 0.5).unwrap();

        (graph, pen, multiply)
    }

    #[test]
    fn pressure_times_port_default() {
        let (graph, _pen, multiply) = pressure_times_half_graph();
        let registry = registry();
        let evaluators = registry.evaluators();

        let mut runner = BrushGraphRunner::new(&graph, registry.as_map(), evaluators).unwrap();

        // Evaluate with pressure = 0.8.
        let info = PaintInformation {
            pressure: 0.8,
            ..Default::default()
        };
        runner.seed_sensors(&info, [0.0, 0.0, 0.0, 1.0], 42, 0);
        runner.execute_cpu();

        // multiply.result should be 0.8 * 0.5 = 0.4.
        let slot = runner.find_node_output_slot(multiply, "result").unwrap();
        let result = runner.read_slot(slot).unwrap();
        match result {
            ScalarValue::Scalar(v) => assert!((v - 0.4).abs() < 1e-6, "expected 0.4, got {v}"),
            other => panic!("expected Scalar, got {:?}", other),
        }
    }

    #[test]
    fn curve_spline_identity() {
        let registry = registry();
        let mut graph = Graph::new();

        let pen_reg = registry.get("pen_input").unwrap();
        let pen = graph.add_node("pen_input", pen_reg.ports.clone(), vec![]);

        // Default identity curve: [(0,0), (1,1)] → output ≈ input
        let curve_reg = registry.get("curve").unwrap();
        let curve = graph.add_node(
            "curve",
            curve_reg.ports.clone(),
            vec![ParamValue::Curve(vec![[0.0, 0.0], [1.0, 1.0]])],
        );

        graph
            .connect(
                PortRef {
                    node: pen,
                    port: "pressure".into(),
                },
                PortRef {
                    node: curve,
                    port: "input".into(),
                },
            )
            .unwrap();

        let evaluators = registry.evaluators();
        let mut runner = BrushGraphRunner::new(&graph, registry.as_map(), evaluators).unwrap();

        let info = PaintInformation {
            pressure: 0.5,
            ..Default::default()
        };
        runner.seed_sensors(&info, [0.0, 0.0, 0.0, 1.0], 42, 0);
        runner.execute_cpu();

        let slot = runner.find_node_output_slot(curve, "output").unwrap();
        let result = runner.read_slot(slot).unwrap();
        match result {
            ScalarValue::Scalar(v) => {
                assert!((v - 0.5).abs() < 0.02, "identity curve at 0.5: got {v}");
            }
            other => panic!("expected Scalar, got {:?}", other),
        }
    }

    #[test]
    fn curve_spline_s_curve() {
        let registry = registry();
        let mut graph = Graph::new();

        let pen_reg = registry.get("pen_input").unwrap();
        let pen = graph.add_node("pen_input", pen_reg.ports.clone(), vec![]);

        // S-curve: compress low values
        let curve_reg = registry.get("curve").unwrap();
        let curve = graph.add_node(
            "curve",
            curve_reg.ports.clone(),
            vec![ParamValue::Curve(vec![[0.0, 0.0], [0.5, 0.2], [1.0, 1.0]])],
        );

        graph
            .connect(
                PortRef {
                    node: pen,
                    port: "pressure".into(),
                },
                PortRef {
                    node: curve,
                    port: "input".into(),
                },
            )
            .unwrap();

        let evaluators = registry.evaluators();
        let mut runner = BrushGraphRunner::new(&graph, registry.as_map(), evaluators).unwrap();

        let info = PaintInformation {
            pressure: 0.5,
            ..Default::default()
        };
        runner.seed_sensors(&info, [0.0, 0.0, 0.0, 1.0], 42, 0);
        runner.execute_cpu();

        let slot = runner.find_node_output_slot(curve, "output").unwrap();
        let result = runner.read_slot(slot).unwrap();
        match result {
            ScalarValue::Scalar(v) => {
                assert!(
                    (v - 0.2).abs() < 0.05,
                    "s-curve at 0.5: expected ~0.2, got {v}"
                );
            }
            other => panic!("expected Scalar, got {:?}", other),
        }
    }

    #[test]
    fn stroke_record_accumulates() {
        let mut record = StrokeRecord::new([1.0, 0.0, 0.0, 1.0], "test_brush".into());
        assert!(record.is_empty());

        for i in 0..10 {
            record.push(PaintInformation {
                pos: [i as f32 * 10.0, 50.0],
                pressure: 0.5 + (i as f32) * 0.05,
                index: i,
                ..Default::default()
            });
        }

        assert_eq!(record.len(), 10);
        assert_eq!(record.events[0].index, 0);
        assert_eq!(record.events[9].index, 9);
        assert!((record.events[0].pressure - 0.5).abs() < 1e-6);
    }

    #[test]
    fn disconnected_input_uses_default() {
        // multiply with no wires — both inputs should use their defaults (1.0).
        let registry = registry();
        let mut graph = Graph::new();

        let mul_reg = registry.get("multiply").unwrap();
        let multiply = graph.add_node("multiply", mul_reg.ports.clone(), vec![]);

        let evaluators = registry.evaluators();
        let mut runner = BrushGraphRunner::new(&graph, registry.as_map(), evaluators).unwrap();

        runner.execute_cpu();

        let slot = runner.find_node_output_slot(multiply, "result").unwrap();
        let result = runner.read_slot(slot).unwrap();
        match result {
            ScalarValue::Scalar(v) => {
                assert!(
                    (v - 1.0).abs() < 1e-6,
                    "expected 1.0 (default * default), got {v}"
                );
            }
            other => panic!("expected Scalar, got {:?}", other),
        }
    }

    #[test]
    fn paint_color_output() {
        let registry = registry();
        let mut graph = Graph::new();

        let color_reg = registry.get("paint_color").unwrap();
        let color_node = graph.add_node("paint_color", color_reg.ports.clone(), vec![]);

        let evaluators = registry.evaluators();
        let mut runner = BrushGraphRunner::new(&graph, registry.as_map(), evaluators).unwrap();

        let fg_color = [1.0, 0.2, 0.3, 1.0];
        let info = PaintInformation::default();
        runner.seed_sensors(&info, fg_color, 42, 0);
        runner.execute_cpu();

        let slot = runner.find_node_output_slot(color_node, "color").unwrap();
        let result = runner.read_slot(slot).unwrap();
        match result {
            ScalarValue::Vec4(c) => assert_eq!(c, fg_color),
            other => panic!("expected Vec4, got {:?}", other),
        }
    }

    #[test]
    fn default_runner_compiles() {
        let runner = super::default_runner();
        assert!(
            runner.is_ok(),
            "default_runner() failed: {:?}",
            runner.err()
        );
    }
}