//! Mizan IR — the canonical KDL document every backend adapter emits and
//! every codegen target consumes. See `docs/AFI_ARCHITECTURE.md` and
//! `cores/mizan-python/src/mizan_core/ir.py` for the locked grammar.

use anyhow::{anyhow, bail, Context, Result};
use indexmap::IndexMap;
use kdl::{KdlDocument, KdlNode, KdlValue};


#[derive(Debug, Default)]
pub struct MizanIR {
    pub types: IndexMap<String, NamedType>,
    pub functions: Vec<MizanFunction>,
    pub contexts: IndexMap<String, MizanContext>,
    pub channels: Vec<MizanChannel>,
}


// ─── Type system ────────────────────────────────────────────────────────────


#[derive(Debug, Clone)]
pub enum NamedType {
    Struct(Vec<StructField>),
    List(TypeShape),
    Enum(Vec<String>),
    Alias(TypeShape),
}


#[derive(Debug, Clone)]
pub struct StructField {
    pub name: String,
    pub required: bool,
    pub default: Option<DefaultValue>,
    pub shape: TypeShape,
}


#[derive(Debug, Clone)]
pub enum TypeShape {
    Primitive(Primitive),
    Ref(String),
    List(Box<TypeShape>),
    Optional(Box<TypeShape>),
    Enum(Vec<String>),
    /// Multi-arm union with two or more non-null branches.
    Union(Vec<TypeShape>),
}


#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Primitive { Integer, Number, Boolean, String }


impl Primitive {
    fn parse(s: &str) -> Result<Self> {
        match s {
            "integer" => Ok(Primitive::Integer),
            "number" => Ok(Primitive::Number),
            "boolean" => Ok(Primitive::Boolean),
            "string" => Ok(Primitive::String),
            other => bail!("unknown primitive {other:?}"),
        }
    }
}


#[derive(Debug, Clone)]
pub enum DefaultValue {
    Integer(i64),
    Number(f64),
    Boolean(bool),
    String(String),
    Null,
}


// ─── Functions ──────────────────────────────────────────────────────────────


#[derive(Debug, Clone)]
pub struct MizanFunction {
    pub name: String,
    pub camel_name: String,
    pub has_input: bool,
    pub input_type: Option<String>,
    pub output_type: String,
    pub output_nullable: bool,
    pub transport: Transport,
    pub is_context: IsContext,
    pub is_form: bool,
    pub form_name: Option<String>,
    pub form_role: Option<String>,
    pub affects: Vec<AffectTarget>,
    pub merge: Vec<String>,
}


#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Transport { Http, Websocket, Both }


impl Transport {
    fn parse(s: &str) -> Result<Self> {
        match s {
            "http" => Ok(Transport::Http),
            "websocket" => Ok(Transport::Websocket),
            "both" => Ok(Transport::Both),
            other => bail!("unknown transport {other:?}"),
        }
    }
}


#[derive(Debug, Clone, PartialEq, Eq)]
pub enum IsContext { No, Yes(String) }


impl IsContext {
    pub fn as_str(&self) -> Option<&str> {
        match self {
            IsContext::No => None,
            IsContext::Yes(s) => Some(s.as_str()),
        }
    }
}


#[derive(Debug, Clone)]
pub struct AffectTarget {
    pub kind: AffectKind,
    pub name: String,
    pub context: Option<String>,
}


#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AffectKind { Context, Function }


// ─── Contexts ───────────────────────────────────────────────────────────────


#[derive(Debug, Clone, Default)]
pub struct MizanContext {
    pub functions: Vec<String>,
    pub params: IndexMap<String, ContextParam>,
}


#[derive(Debug, Clone)]
pub struct ContextParam {
    pub ty: Primitive,
    pub required: bool,
    pub shared_by: Vec<String>,
}


// ─── Channels (Django-only) ─────────────────────────────────────────────────


#[derive(Debug, Clone)]
pub struct MizanChannel {
    pub name: String,
    pub pascal_name: String,
    pub params_type: Option<String>,
    pub react_message_type: Option<String>,
    pub django_message_type: Option<String>,
}


impl MizanChannel {
    pub fn has_params(&self) -> bool { self.params_type.is_some() }
    pub fn has_react_message(&self) -> bool { self.react_message_type.is_some() }
    pub fn has_django_message(&self) -> bool { self.django_message_type.is_some() }
}


// ─── KDL parsing ────────────────────────────────────────────────────────────


pub fn parse_ir(source: &str) -> Result<MizanIR> {
    let doc: KdlDocument = source.parse()
        .map_err(|e| anyhow!("KDL parse error: {e}"))?;

    let mut ir = MizanIR::default();
    for node in doc.nodes() {
        match node.name().value() {
            "type" => {
                let (name, ty) = parse_named_type(node)?;
                ir.types.insert(name, ty);
            }
            "function" => ir.functions.push(parse_function(node)?),
            "context" => {
                let (name, ctx) = parse_context(node)?;
                ir.contexts.insert(name, ctx);
            }
            "channel" => ir.channels.push(parse_channel(node)?),
            other => bail!("unknown top-level KDL node {other:?}"),
        }
    }
    Ok(ir)
}


fn parse_named_type(node: &KdlNode) -> Result<(String, NamedType)> {
    let name = first_string_arg(node)
        .context("`type` requires a name as its first argument")?;
    let children = node.children()
        .ok_or_else(|| anyhow!("type {name:?}: missing children block"))?;
    let kind_node = single_child(children, &format!("type {name:?}"))?;
    let kind = match kind_node.name().value() {
        "struct" => NamedType::Struct(parse_struct_fields(kind_node)?),
        "list" => NamedType::List(type_child_of(kind_node, &format!("type {name:?} list"))?),
        "enum" => NamedType::Enum(parse_string_args(kind_node)),
        "alias" => NamedType::Alias(type_child_of(kind_node, &format!("type {name:?} alias"))?),
        other => bail!("type {name:?}: unknown shape node {other:?}"),
    };
    Ok((name, kind))
}


fn parse_struct_fields(struct_node: &KdlNode) -> Result<Vec<StructField>> {
    let mut fields = Vec::new();
    let Some(children) = struct_node.children() else { return Ok(fields); };
    for child in children.nodes() {
        if child.name().value() != "field" {
            bail!("struct: unexpected node {:?}", child.name().value());
        }
        fields.push(parse_struct_field(child)?);
    }
    Ok(fields)
}


fn parse_struct_field(field_node: &KdlNode) -> Result<StructField> {
    let name = first_string_arg(field_node).context("`field` requires a name")?;
    let required = bool_prop(field_node, "required").unwrap_or(true);
    let default = field_node.entry("default")
        .map(|e| parse_default_value(e.value()))
        .transpose()?;
    let shape = type_child_of(field_node, &format!("field {name:?}"))?;
    Ok(StructField { name, required, default, shape })
}


fn parse_default_value(v: &KdlValue) -> Result<DefaultValue> {
    if v.is_null() { return Ok(DefaultValue::Null); }
    if let Some(b) = v.as_bool() { return Ok(DefaultValue::Boolean(b)); }
    if let Some(i) = v.as_integer() { return Ok(DefaultValue::Integer(i as i64)); }
    if let Some(f) = v.as_float() { return Ok(DefaultValue::Number(f)); }
    if let Some(s) = v.as_string() { return Ok(DefaultValue::String(s.to_string())); }
    bail!("unsupported default literal: {v:?}")
}


fn type_child_of(parent: &KdlNode, label: &str) -> Result<TypeShape> {
    let children = parent.children()
        .ok_or_else(|| anyhow!("{label}: missing children for type-shape"))?;
    let nodes = children.nodes();
    if nodes.len() != 1 {
        bail!("{label}: expected exactly one type-shape child, got {}", nodes.len());
    }
    parse_type_shape(&nodes[0])
}


fn parse_type_shape(node: &KdlNode) -> Result<TypeShape> {
    match node.name().value() {
        "primitive" => Ok(TypeShape::Primitive(Primitive::parse(&first_string_arg(node)?)?)),
        "ref" => Ok(TypeShape::Ref(first_string_arg(node)?)),
        "list" => Ok(TypeShape::List(Box::new(type_child_of(node, "list")?))),
        "optional" => Ok(TypeShape::Optional(Box::new(type_child_of(node, "optional")?))),
        "enum" => Ok(TypeShape::Enum(parse_string_args(node))),
        "union" => {
            let children = node.children()
                .ok_or_else(|| anyhow!("union: missing children"))?;
            let branches: Result<Vec<TypeShape>> = children.nodes().iter()
                .map(parse_type_shape).collect();
            Ok(TypeShape::Union(branches?))
        }
        other => bail!("unknown type-shape node {other:?}"),
    }
}


fn parse_function(node: &KdlNode) -> Result<MizanFunction> {
    let name = first_string_arg(node)
        .context("`function` requires a name as its first argument")?;
    let children = node.children()
        .ok_or_else(|| anyhow!("function {name:?}: missing children"))?;

    let mut camel = None;
    let mut has_input = false;
    let mut input_type = None;
    let mut output_type = None;
    let mut output_nullable = false;
    let mut transport = Transport::Http;
    let mut is_context = IsContext::No;
    let mut is_form = false;
    let mut form_name = None;
    let mut form_role = None;
    let mut affects: Vec<AffectTarget> = Vec::new();
    let mut merge: Vec<String> = Vec::new();

    for child in children.nodes() {
        match child.name().value() {
            "camel" => camel = Some(string_arg(child, "camel")?),
            "has-input" => has_input = bool_arg(child, "has-input")?,
            "input" => input_type = Some(string_arg(child, "input")?),
            "output" => output_type = Some(string_arg(child, "output")?),
            "output-nullable" => output_nullable = bool_arg(child, "output-nullable")?,
            "transport" => transport = Transport::parse(&string_arg(child, "transport")?)?,
            "context" => is_context = IsContext::Yes(string_arg(child, "context")?),
            "is-form" => is_form = bool_arg(child, "is-form")?,
            "form-name" => form_name = Some(string_arg(child, "form-name")?),
            "form-role" => form_role = Some(string_arg(child, "form-role")?),
            "affects" => affects.push(AffectTarget {
                kind: AffectKind::Context,
                name: string_arg(child, "affects")?,
                context: None,
            }),
            "merge" => merge.push(string_arg(child, "merge")?),
            other => bail!("function {name:?}: unknown child {other:?}"),
        }
    }

    Ok(MizanFunction {
        name: name.clone(),
        camel_name: camel.ok_or_else(|| anyhow!("function {name:?}: missing `camel`"))?,
        has_input,
        input_type,
        output_type: output_type.ok_or_else(|| anyhow!("function {name:?}: missing `output`"))?,
        output_nullable,
        transport,
        is_context,
        is_form,
        form_name,
        form_role,
        affects,
        merge,
    })
}


fn parse_context(node: &KdlNode) -> Result<(String, MizanContext)> {
    let name = first_string_arg(node).context("`context` requires a name")?;
    let mut ctx = MizanContext::default();
    let Some(children) = node.children() else { return Ok((name, ctx)); };
    for child in children.nodes() {
        match child.name().value() {
            "function" => ctx.functions.push(string_arg(child, "function")?),
            "param" => {
                let (pname, param) = parse_context_param(child)?;
                ctx.params.insert(pname, param);
            }
            other => bail!("context {name:?}: unknown child {other:?}"),
        }
    }
    Ok((name, ctx))
}


fn parse_context_param(node: &KdlNode) -> Result<(String, ContextParam)> {
    let pname = first_string_arg(node).context("`param` requires a name")?;
    let children = node.children()
        .ok_or_else(|| anyhow!("param {pname:?}: missing children"))?;
    let mut ty = None;
    let mut required = false;
    let mut shared_by = Vec::new();
    for child in children.nodes() {
        match child.name().value() {
            "type" => ty = Some(Primitive::parse(&string_arg(child, "type")?)?),
            "required" => required = bool_arg(child, "required")?,
            "shared-by" => shared_by.push(string_arg(child, "shared-by")?),
            other => bail!("param {pname:?}: unknown child {other:?}"),
        }
    }
    Ok((pname.clone(), ContextParam {
        ty: ty.ok_or_else(|| anyhow!("param {pname:?}: missing `type`"))?,
        required,
        shared_by,
    }))
}


fn parse_channel(node: &KdlNode) -> Result<MizanChannel> {
    let name = first_string_arg(node).context("`channel` requires a name")?;
    let children = node.children()
        .ok_or_else(|| anyhow!("channel {name:?}: missing children"))?;
    let mut pascal_name = None;
    let mut params_type = None;
    let mut react_message_type = None;
    let mut django_message_type = None;
    for child in children.nodes() {
        match child.name().value() {
            "pascal-name" => pascal_name = Some(string_arg(child, "pascal-name")?),
            "params" => params_type = Some(string_arg(child, "params")?),
            "react-message" => react_message_type = Some(string_arg(child, "react-message")?),
            "django-message" => django_message_type = Some(string_arg(child, "django-message")?),
            other => bail!("channel {name:?}: unknown child {other:?}"),
        }
    }
    Ok(MizanChannel {
        name: name.clone(),
        pascal_name: pascal_name.ok_or_else(|| anyhow!("channel {name:?}: missing `pascal-name`"))?,
        params_type,
        react_message_type,
        django_message_type,
    })
}


// ─── KDL accessor helpers ───────────────────────────────────────────────────


fn first_string_arg(node: &KdlNode) -> Result<String> {
    let entry = node.entries().iter()
        .find(|e| e.name().is_none())
        .ok_or_else(|| anyhow!("node {:?}: missing positional argument", node.name().value()))?;
    entry.value().as_string()
        .map(str::to_string)
        .ok_or_else(|| anyhow!("node {:?}: positional argument is not a string", node.name().value()))
}


fn string_arg(node: &KdlNode, label: &str) -> Result<String> {
    first_string_arg(node).context(format!("{label}: requires a string argument"))
}


fn bool_arg(node: &KdlNode, label: &str) -> Result<bool> {
    node.entries().iter()
        .find(|e| e.name().is_none())
        .and_then(|e| e.value().as_bool())
        .ok_or_else(|| anyhow!("{label}: missing positional bool argument"))
}


fn bool_prop(node: &KdlNode, key: &str) -> Option<bool> {
    node.entry(key).and_then(|e| e.value().as_bool())
}


fn parse_string_args(node: &KdlNode) -> Vec<String> {
    node.entries().iter()
        .filter(|e| e.name().is_none())
        .filter_map(|e| e.value().as_string().map(str::to_string))
        .collect()
}


fn single_child<'a>(children: &'a KdlDocument, label: &str) -> Result<&'a KdlNode> {
    let nodes = children.nodes();
    if nodes.len() != 1 {
        bail!("{label}: expected exactly one child node, got {}", nodes.len());
    }
    Ok(&nodes[0])
}


// ─── Library entry point ────────────────────────────────────────────────────


pub fn parse_ir_from_str(source: &str) -> Result<MizanIR> {
    parse_ir(source)
}