//! KDL emitter — byte-equivalent to `cores/mizan-python/src/mizan_core/ir.py`.
//!
//! The Python emitter is the spec; this is the second implementation under
//! the same contract. Any divergence is a bug here, not a contract change.

use crate::ir::{DefaultValue, NamedType, Primitive, StructField, TypeShape};
use crate::registry::{CONTEXTS, FUNCTIONS, TYPES};
use crate::traits::FunctionSpec;
use std::collections::BTreeMap;

const INDENT: &str = "    ";

/// Escape a string for KDL — same escape set as the Python emitter.
fn kdl_string(s: &str) -> String {
    let mut out = String::with_capacity(s.len() + 2);
    out.push('"');
    for c in s.chars() {
        match c {
            '\\' => out.push_str("\\\\"),
            '"' => out.push_str("\\\""),
            '\n' => out.push_str("\\n"),
            '\r' => out.push_str("\\r"),
            '\t' => out.push_str("\\t"),
            other => out.push(other),
        }
    }
    out.push('"');
    out
}

fn kdl_bool(b: bool) -> &'static str {
    if b {
        "#true"
    } else {
        "#false"
    }
}

fn kdl_default(v: &DefaultValue) -> String {
    match v {
        DefaultValue::Null => "#null".into(),
        DefaultValue::Boolean(b) => kdl_bool(*b).into(),
        DefaultValue::Integer(i) => i.to_string(),
        DefaultValue::Number(f) => {
            // Match Python's `repr(float)` for whole-number-equal-but-float
            // values: e.g. 1.0 → "1.0", not "1".
            if f.fract() == 0.0 && f.is_finite() {
                format!("{f:.1}")
            } else {
                f.to_string()
            }
        }
        DefaultValue::String(s) => kdl_string(s),
    }
}

/// Convert snake_case to camelCase. Matches Python's `_snake_to_camel`.
pub fn snake_to_camel(name: &str) -> String {
    let normalized = name.replace('.', "_").replace('-', "_");
    let mut parts = normalized.split('_');
    let mut out = String::new();
    if let Some(first) = parts.next() {
        out.push_str(first);
    }
    for part in parts {
        if part.is_empty() {
            continue;
        }
        let mut chars = part.chars();
        if let Some(c) = chars.next() {
            out.extend(c.to_uppercase());
            out.push_str(chars.as_str());
        }
    }
    out
}

struct Emitter<'a> {
    lines: Vec<String>,
    /// Types whose references should be substituted with their inline
    /// shape at the use site (and which don't emit as their own
    /// `type "X" { ... }` entries). Populated from `IrSnapshot::inlines`.
    inlines: &'a BTreeMap<&'static str, TypeShape>,
}

impl<'a> Emitter<'a> {
    fn new(inlines: &'a BTreeMap<&'static str, TypeShape>) -> Self {
        Self {
            lines: Vec::new(),
            inlines,
        }
    }

    fn prefix(&self, indent: usize) -> String {
        INDENT.repeat(indent)
    }

    fn leaf(&mut self, indent: usize, parts: &[&str]) {
        let mut line = self.prefix(indent);
        line.push_str(&parts.join(" "));
        self.lines.push(line);
    }

    fn open(&mut self, indent: usize, parts: &[&str]) {
        let mut line = self.prefix(indent);
        line.push_str(&parts.join(" "));
        line.push_str(" {");
        self.lines.push(line);
    }

    fn close(&mut self, indent: usize) {
        let mut line = self.prefix(indent);
        line.push('}');
        self.lines.push(line);
    }

    fn blank(&mut self) {
        self.lines.push(String::new());
    }

    fn emit_type_child(&mut self, indent: usize, shape: &TypeShape) {
        match shape {
            TypeShape::Primitive(p) => {
                let name = kdl_string(p.name());
                self.leaf(indent, &["primitive", &name]);
            }
            TypeShape::Ref(name) => {
                // Inline-substitute when the referenced type is a
                // primitive-alias or string-enum. Matches Python's
                // Pydantic Literal/alias inlining.
                if let Some(inline_shape) = self.inlines.get(name).cloned() {
                    self.emit_type_child(indent, &inline_shape);
                    return;
                }
                let n = kdl_string(name);
                self.leaf(indent, &["ref", &n]);
            }
            TypeShape::List(inner) => {
                self.open(indent, &["list"]);
                self.emit_type_child(indent + 1, inner);
                self.close(indent);
            }
            TypeShape::Optional(inner) => {
                self.open(indent, &["optional"]);
                self.emit_type_child(indent + 1, inner);
                self.close(indent);
            }
            TypeShape::Enum(variants) => {
                let mut parts: Vec<String> = vec!["enum".into()];
                for v in variants {
                    parts.push(kdl_string(v));
                }
                let line: Vec<&str> = parts.iter().map(String::as_str).collect();
                self.leaf(indent, &line);
            }
            TypeShape::Union(branches) => {
                self.open(indent, &["union"]);
                for b in branches {
                    self.emit_type_child(indent + 1, b);
                }
                self.close(indent);
            }
        }
    }

    fn emit_named_type(&mut self, indent: usize, name: &str, body: &NamedType) {
        let name_lit = kdl_string(name);
        self.open(indent, &["type", &name_lit]);
        match body {
            NamedType::Struct(fields) => {
                self.open(indent + 1, &["struct"]);
                for field in fields {
                    self.emit_struct_field(indent + 2, field);
                }
                self.close(indent + 1);
            }
            NamedType::Alias(inner) => {
                self.open(indent + 1, &["alias"]);
                self.emit_type_child(indent + 2, inner);
                self.close(indent + 1);
            }
            NamedType::Enum(variants) => {
                let mut parts: Vec<String> = vec!["enum".into()];
                for v in variants {
                    parts.push(kdl_string(v));
                }
                let line: Vec<&str> = parts.iter().map(String::as_str).collect();
                self.leaf(indent + 1, &line);
            }
        }
        self.close(indent);
    }

    fn emit_struct_field(&mut self, indent: usize, field: &StructField) {
        let name = kdl_string(field.name);
        let mut header: Vec<String> = vec!["field".into(), name];
        if !field.required {
            header.push(format!("required={}", kdl_bool(false)));
            if let Some(default) = &field.default {
                header.push(format!("default={}", kdl_default(default)));
            }
        }
        let line_parts: Vec<&str> = header.iter().map(String::as_str).collect();
        self.open(indent, &line_parts);
        self.emit_type_child(indent + 1, &field.shape);
        self.close(indent);
    }

    fn emit_function(&mut self, indent: usize, fn_spec: &dyn FunctionSpec) {
        let name = kdl_string(fn_spec.name());
        self.open(indent, &["function", &name]);

        let camel = kdl_string(fn_spec.camel_name());
        self.leaf(indent + 1, &["camel", &camel]);

        self.leaf(indent + 1, &["has-input", kdl_bool(fn_spec.has_input())]);

        if let Some(input_type) = fn_spec.input_type() {
            let lit = kdl_string(input_type);
            self.leaf(indent + 1, &["input", &lit]);
        }

        let output_lit = kdl_string(fn_spec.output_type());
        self.leaf(indent + 1, &["output", &output_lit]);

        if fn_spec.output_nullable() {
            self.leaf(indent + 1, &["output-nullable", kdl_bool(true)]);
        }

        let transport_lit = kdl_string(fn_spec.transport().name());
        self.leaf(indent + 1, &["transport", &transport_lit]);

        if let Some(ctx) = fn_spec.context() {
            let lit = kdl_string(ctx);
            self.leaf(indent + 1, &["context", &lit]);
        }

        for affect in fn_spec.affects() {
            // Mirror Python's behavior: only context-typed affects make it
            // into the KDL `affects` leaf. Function-typed affects are
            // reserved for a future IR extension.
            if let crate::ir::AffectTarget::Context(name) = affect {
                let lit = kdl_string(name);
                self.leaf(indent + 1, &["affects", &lit]);
            }
        }

        for merge in fn_spec.merge() {
            let lit = kdl_string(merge);
            self.leaf(indent + 1, &["merge", &lit]);
        }

        if fn_spec.is_form() {
            self.leaf(indent + 1, &["is-form", kdl_bool(true)]);
            if let Some(form_name) = fn_spec.form_name() {
                let lit = kdl_string(form_name);
                self.leaf(indent + 1, &["form-name", &lit]);
            }
            if let Some(form_role) = fn_spec.form_role() {
                let lit = kdl_string(form_role);
                self.leaf(indent + 1, &["form-role", &lit]);
            }
        }

        self.close(indent);
    }

    fn emit_context(&mut self, indent: usize, ctx_name: &str, members: &[&'static dyn FunctionSpec]) {
        let name_lit = kdl_string(ctx_name);
        self.open(indent, &["context", &name_lit]);

        // Function membership in registration order.
        for fn_spec in members {
            let lit = kdl_string(fn_spec.name());
            self.leaf(indent + 1, &["function", &lit]);
        }

        // Param info — collect across every member, then emit alphabetized
        // by param name to match Python.
        struct ParamSlot {
            primitive: Primitive,
            shared_by: Vec<&'static str>,
        }
        let mut params: BTreeMap<&'static str, ParamSlot> = BTreeMap::new();
        for fn_spec in members {
            for p in fn_spec.input_params() {
                let slot = params.entry(p.name).or_insert(ParamSlot {
                    primitive: p.primitive,
                    shared_by: Vec::new(),
                });
                slot.primitive = p.primitive;
                slot.shared_by.push(fn_spec.name());
            }
        }

        let member_count = members.len();
        for (param_name, slot) in params.iter() {
            let name_lit = kdl_string(param_name);
            self.open(indent + 1, &["param", &name_lit]);
            let type_lit = kdl_string(slot.primitive.name());
            self.leaf(indent + 2, &["type", &type_lit]);
            let required = slot.shared_by.len() == member_count;
            self.leaf(indent + 2, &["required", kdl_bool(required)]);
            for sharer in &slot.shared_by {
                let lit = kdl_string(sharer);
                self.leaf(indent + 2, &["shared-by", &lit]);
            }
            self.close(indent + 1);
        }

        self.close(indent);
    }

    fn into_string(mut self) -> String {
        // Trim trailing blanks, then add a single terminating newline.
        while matches!(self.lines.last(), Some(s) if s.is_empty()) {
            self.lines.pop();
        }
        let mut out = self.lines.join("\n");
        out.push('\n');
        out
    }
}

/// Collected typed registries view used by `build_ir`.
pub(crate) struct IrSnapshot {
    pub types: BTreeMap<&'static str, NamedType>,
    pub functions: Vec<&'static dyn FunctionSpec>,
    pub contexts: Vec<(&'static str, Vec<&'static dyn FunctionSpec>)>,
    /// Types that inline to a `TypeShape` at every reference site rather
    /// than emitting as their own `type "X" { ... }` entry. Populated from
    /// `Alias(Primitive(_))` and `Enum` named types — both are
    /// information-zero indirections that the codegen consumer doesn't
    /// gain anything from naming. Matches the Python emitter's behavior
    /// (Pydantic `FigureId = str` and `Literal["..."]` inline; they don't
    /// materialize as named types).
    pub inlines: BTreeMap<&'static str, TypeShape>,
}

impl IrSnapshot {
    pub(crate) fn collect() -> Self {
        // Types: alphabetized for byte-equivalence with Python's `sorted(named_types)`.
        let mut all_types: BTreeMap<&'static str, NamedType> = BTreeMap::new();
        for entry in TYPES {
            all_types.insert(entry.name, (entry.shape_fn)());
        }

        // Partition into emit-candidate types vs inlines. An inline is a
        // named type whose shape collapses to a single `TypeShape` at the
        // field site — primitive aliases and string enums.
        let mut candidates: BTreeMap<&'static str, NamedType> = BTreeMap::new();
        let mut inlines: BTreeMap<&'static str, TypeShape> = BTreeMap::new();
        for (name, body) in all_types {
            match &body {
                NamedType::Alias(TypeShape::Primitive(p)) => {
                    inlines.insert(name, TypeShape::Primitive(*p));
                }
                NamedType::Enum(variants) => {
                    inlines.insert(name, TypeShape::Enum(variants.clone()));
                }
                _ => {
                    candidates.insert(name, body);
                }
            }
        }

        // Tree-shake: keep only types reachable from a registered function's
        // input/output. The function macro registers canonical-named
        // entries (e.g. `userPrefsOutput`); derive registers original-named
        // entries (`UserPrefs`, `BrushSettings`, …). Only those reached
        // via Ref-walk from a function's input/output names belong in the
        // emitted IR. Mirrors Python's `_collect_named_types`.
        let mut reachable: std::collections::HashSet<&'static str> =
            std::collections::HashSet::new();
        let mut frontier: Vec<&'static str> = Vec::new();
        for fn_spec in FUNCTIONS {
            if fn_spec.private() {
                continue;
            }
            if let Some(input_name) = fn_spec.input_type() {
                if reachable.insert(input_name) {
                    frontier.push(input_name);
                }
            }
            let output_name = fn_spec.output_type();
            if reachable.insert(output_name) {
                frontier.push(output_name);
            }
        }
        while let Some(name) = frontier.pop() {
            // Inlines don't carry refs we care about (Primitive/Enum); skip.
            if inlines.contains_key(name) {
                continue;
            }
            let body = match candidates.get(name) {
                Some(b) => b.clone(),
                None => continue,
            };
            collect_refs(&body, &mut |r| {
                if reachable.insert(r) {
                    frontier.push(r);
                }
            });
        }
        let types: BTreeMap<&'static str, NamedType> = candidates
            .into_iter()
            .filter(|(name, _)| reachable.contains(name))
            .collect();

        // Functions: alphabetical by wire name (canonical IR ordering,
        // matches the Python emitter's `sorted(functions)`). Skip `private`.
        let mut functions: Vec<&'static dyn FunctionSpec> = FUNCTIONS
            .iter()
            .copied()
            .filter(|f| !f.private())
            .collect();
        functions.sort_by_key(|f| f.name());

        // Contexts: alphabetical by name (canonical IR ordering), each with
        // its members sorted alphabetically too.
        let mut context_names: Vec<&'static str> = CONTEXTS.iter().map(|c| c.name).collect();
        context_names.sort();
        let mut contexts: Vec<(&'static str, Vec<&'static dyn FunctionSpec>)> = Vec::new();
        for name in context_names {
            let mut members: Vec<&'static dyn FunctionSpec> = functions
                .iter()
                .copied()
                .filter(|f| f.context() == Some(name))
                .collect();
            members.sort_by_key(|f| f.name());
            if !members.is_empty() {
                contexts.push((name, members));
            }
        }

        Self {
            types,
            functions,
            contexts,
            inlines,
        }
    }
}

/// Walk every Ref in a NamedType's shape and call `visit` for each name.
fn collect_refs<F: FnMut(&'static str)>(body: &NamedType, visit: &mut F) {
    match body {
        NamedType::Struct(fields) => {
            for field in fields {
                walk_shape_refs(&field.shape, visit);
            }
        }
        NamedType::Alias(inner) => walk_shape_refs(inner, visit),
        NamedType::Enum(_) => {}
    }
}

fn walk_shape_refs<F: FnMut(&'static str)>(shape: &TypeShape, visit: &mut F) {
    match shape {
        TypeShape::Ref(name) => visit(name),
        TypeShape::List(inner) | TypeShape::Optional(inner) => walk_shape_refs(inner, visit),
        TypeShape::Union(branches) => {
            for b in branches {
                walk_shape_refs(b, visit);
            }
        }
        TypeShape::Primitive(_) | TypeShape::Enum(_) => {}
    }
}

/// Build the Mizan IR for every registered type/function/context. Returns KDL.
pub fn build_ir() -> String {
    crate::graph_check::verify_invariants();
    let snap = IrSnapshot::collect();
    let mut em = Emitter::new(&snap.inlines);

    // Type definitions
    let types_emitted = !snap.types.is_empty();
    for (name, body) in &snap.types {
        em.emit_named_type(0, name, body);
    }
    if types_emitted {
        em.blank();
    }

    // Functions
    let fns_emitted = !snap.functions.is_empty();
    for fn_spec in &snap.functions {
        em.emit_function(0, *fn_spec);
    }
    if fns_emitted {
        em.blank();
    }

    // Contexts
    let ctxs_emitted = !snap.contexts.is_empty();
    for (ctx_name, members) in &snap.contexts {
        em.emit_context(0, ctx_name, members);
    }
    if ctxs_emitted {
        em.blank();
    }

    // Future: channels — once channel registry lands on the Rust side.

    em.into_string()
}

