use mlf_lang::ast::*;
use mlf_lang::{ResolvedRef, Workspace};
use serde_json::{Map, Value, json};
use std::collections::HashMap;

/// A non-fatal advisory emitted during codegen. Mirrors the shape of
/// `ConversionWarning` on the JSON→MLF side so both directions use one
/// structured-warning model.
#[derive(Debug, Clone, PartialEq)]
pub struct CodegenWarning {
    /// Namespace of the lexicon that emitted the warning.
    pub namespace: String,
    /// Human-readable description of what was coerced and why.
    pub message: String,
}

/// Bundle of codegen output: the generated lexicon JSON plus any
/// advisory warnings produced while walking the AST.
#[derive(Debug, Clone)]
pub struct CodegenOutput {
    pub json: Value,
    pub warnings: Vec<CodegenWarning>,
}

// Re-export inventory for macros
#[doc(hidden)]
pub use inventory;

// Plugin system for code generators
pub mod plugin {
    use super::*;

    /// Context passed to code generators
    pub struct GeneratorContext<'a> {
        pub namespace: &'a str,
        pub lexicon: &'a Lexicon,
        pub workspace: &'a Workspace,
    }

    /// Trait for code generators
    pub trait CodeGenerator: Send + Sync {
        /// Unique identifier for this generator (e.g., "typescript", "python", "rust")
        fn name(&self) -> &'static str;

        /// Human-readable description
        fn description(&self) -> &'static str;

        /// File extension for generated files (e.g., ".ts", ".py", ".rs")
        fn file_extension(&self) -> &'static str;

        /// Generate code from a lexicon
        fn generate(&self, ctx: &GeneratorContext) -> Result<String, String>;
    }

    // Registry of code generators using inventory
    inventory::collect!(&'static dyn CodeGenerator);

    /// Get all registered generators as a Vec
    pub fn generators() -> Vec<&'static dyn CodeGenerator> {
        let mut result = Vec::new();
        for generator in inventory::iter::<&'static dyn CodeGenerator> {
            result.push(*generator);
        }
        result
    }

    /// Find a generator by name
    pub fn find_generator(name: &str) -> Option<&'static dyn CodeGenerator> {
        for generator in inventory::iter::<&'static dyn CodeGenerator> {
            if generator.name() == name {
                return Some(*generator);
            }
        }
        None
    }

    /// Macro to easily register a generator
    #[macro_export]
    macro_rules! register_generator {
        ($generator:expr) => {
            $crate::inventory::submit! {
                &$generator as &'static dyn $crate::plugin::CodeGenerator
            }
        };
    }
}

pub use plugin::{CodeGenerator, GeneratorContext};

fn has_main_annotation(annotations: &[Annotation]) -> bool {
    annotations.iter().any(|ann| ann.name.name == "main")
}

fn get_annotation_string_value(annotations: &[Annotation], name: &str) -> Option<String> {
    annotations
        .iter()
        .find(|ann| ann.name.name == name)
        .and_then(|ann| {
            // Get first positional argument if it exists
            ann.args.first().and_then(|arg| match arg {
                AnnotationArg::Positional(AnnotationValue::String(s)) => Some(s.clone()),
                _ => None,
            })
        })
}

fn get_encoding_annotation(annotations: &[Annotation], param_name: &str) -> Option<String> {
    annotations
        .iter()
        .find(|ann| ann.name.name == "encoding")
        .and_then(|ann| {
            // First check for named argument matching param_name
            for arg in &ann.args {
                if let AnnotationArg::Named { name, value } = arg {
                    if name.name == param_name {
                        if let AnnotationValue::String(s) = value {
                            return Some(s.clone());
                        }
                    }
                }
            }

            // Fall back to positional argument (applies to both input and output)
            for arg in &ann.args {
                if let AnnotationArg::Positional(AnnotationValue::String(s)) = arg {
                    return Some(s.clone());
                }
            }

            None
        })
}

pub fn generate_lexicon(
    namespace: &str,
    lexicon: &Lexicon,
    workspace: &Workspace,
) -> CodegenOutput {
    let usage_counts = analyze_type_usage(lexicon);
    let eligibility = MainEligibility::for_lexicon(namespace, lexicon);

    let mut defs = Map::new();
    let mut self_description = String::new();
    let mut self_extensions: Vec<(String, Value)> = Vec::new();
    let mut warnings: Vec<CodegenWarning> = Vec::new();

    for item in &lexicon.items {
        match item {
            Item::Record(record) => {
                let mut value = generate_record_json(record, &usage_counts, workspace, namespace);
                apply_extension_annotations(
                    &mut value,
                    &record.annotations,
                    workspace,
                    namespace,
                    &mut warnings,
                );
                insert_def(
                    &mut defs,
                    &record.name.name,
                    eligibility.is_main(&record.name.name, &record.annotations),
                    value,
                );
            }
            Item::Query(query) => {
                let mut value = generate_query_json(query, &usage_counts, workspace, namespace);
                apply_extension_annotations(
                    &mut value,
                    &query.annotations,
                    workspace,
                    namespace,
                    &mut warnings,
                );
                insert_def(
                    &mut defs,
                    &query.name.name,
                    eligibility.is_main(&query.name.name, &query.annotations),
                    value,
                );
            }
            Item::Procedure(procedure) => {
                let mut value =
                    generate_procedure_json(procedure, &usage_counts, workspace, namespace);
                apply_extension_annotations(
                    &mut value,
                    &procedure.annotations,
                    workspace,
                    namespace,
                    &mut warnings,
                );
                insert_def(
                    &mut defs,
                    &procedure.name.name,
                    eligibility.is_main(&procedure.name.name, &procedure.annotations),
                    value,
                );
            }
            Item::Subscription(subscription) => {
                let mut value =
                    generate_subscription_json(subscription, &usage_counts, workspace, namespace);
                apply_extension_annotations(
                    &mut value,
                    &subscription.annotations,
                    workspace,
                    namespace,
                    &mut warnings,
                );
                insert_def(
                    &mut defs,
                    &subscription.name.name,
                    eligibility.is_main(&subscription.name.name, &subscription.annotations),
                    value,
                );
            }
            Item::DefType(def_type) => {
                let mut value =
                    generate_def_type_json(def_type, &usage_counts, workspace, namespace);
                apply_extension_annotations(
                    &mut value,
                    &def_type.annotations,
                    workspace,
                    namespace,
                    &mut warnings,
                );
                insert_def(
                    &mut defs,
                    &def_type.name.name,
                    eligibility.is_main(&def_type.name.name, &def_type.annotations),
                    value,
                );
            }
            Item::Token(token) => {
                let mut token_obj = Map::new();
                token_obj.insert("type".to_string(), json!("token"));
                insert_opt_str(&mut token_obj, "description", &extract_docs(&token.docs));
                let mut value = Value::Object(token_obj);
                apply_extension_annotations(
                    &mut value,
                    &token.annotations,
                    workspace,
                    namespace,
                    &mut warnings,
                );
                defs.insert(token.name.name.clone(), value);
            }
            Item::SelfItem(self_item) => {
                // The `self { }` item carries lexicon-level metadata. Its
                // docs become the top-level `description`; its extension
                // annotations become top-level JSON fields alongside
                // `lexicon`, `id`, `defs`.
                self_description = extract_docs(&self_item.docs);
                self_extensions = collect_extension_fields(
                    &self_item.annotations,
                    workspace,
                    namespace,
                    &mut warnings,
                );
            }
            // Inline types never appear in `defs` — they expand at their point
            // of use. Use statements are structural and not emitted.
            _ => {}
        }
    }

    let mut root = Map::new();
    root.insert("$type".to_string(), json!("com.atproto.lexicon.schema"));
    root.insert("lexicon".to_string(), json!(1));
    root.insert("id".to_string(), json!(namespace));
    insert_opt_str(&mut root, "description", &self_description);
    for (key, value) in self_extensions {
        root.insert(key, value);
    }
    root.insert("defs".to_string(), json!(defs));
    CodegenOutput {
        json: Value::Object(root),
        warnings,
    }
}

/// Extract `@const` / `@reference` annotations into (key, JSON value)
/// pairs. Used by the self-item handling to populate top-level lexicon
/// fields. Annotations of other names (including generator-scoped ones
/// like `@rust:deprecated`) are ignored.
fn collect_extension_fields(
    annotations: &[Annotation],
    workspace: &Workspace,
    current_namespace: &str,
    warnings: &mut Vec<CodegenWarning>,
) -> Vec<(String, Value)> {
    annotations
        .iter()
        .filter_map(|ann| extension_field(ann, workspace, current_namespace, warnings))
        .collect()
}

/// Mutate `value` (expected to be a JSON object) in place, adding any
/// extension fields derived from `@const` / `@reference` annotations.
/// Non-object values are left untouched.
fn apply_extension_annotations(
    value: &mut Value,
    annotations: &[Annotation],
    workspace: &Workspace,
    current_namespace: &str,
    warnings: &mut Vec<CodegenWarning>,
) {
    let Some(obj) = value.as_object_mut() else {
        return;
    };
    for (key, field_value) in
        collect_extension_fields(annotations, workspace, current_namespace, warnings)
    {
        obj.insert(key, field_value);
    }
}

/// If `annotation` is a recognised extension (`@const` or `@reference`),
/// return the `(key, JSON value)` it produces. Otherwise `None`.
///
/// `@const(key, value)` — value is a literal; rendered verbatim.
/// `@reference(key, path)` — path is resolved through the workspace and
/// emitted as an NSID string.
fn extension_field(
    annotation: &Annotation,
    workspace: &Workspace,
    current_namespace: &str,
    warnings: &mut Vec<CodegenWarning>,
) -> Option<(String, Value)> {
    if !annotation.selectors.is_empty() {
        // Generator-scoped annotation (e.g. `@rust:deprecated`); not an
        // extension.
        return None;
    }
    let name = annotation.name.name.as_str();
    let is_const = name == "const";
    let is_reference = name == "reference";
    if !is_const && !is_reference {
        return None;
    }
    // Both forms take exactly two positional args: (key: string, value).
    let mut args = annotation.args.iter();
    let key_arg = args.next()?;
    let value_arg = args.next()?;
    if args.next().is_some() {
        return None;
    }
    let AnnotationArg::Positional(AnnotationValue::String(key)) = key_arg else {
        return None;
    };
    let AnnotationArg::Positional(value) = value_arg else {
        return None;
    };
    let json_value = if is_const {
        annotation_value_to_json(value, key, current_namespace, warnings)
    } else {
        // @reference: expect a type path; emit the resolved NSID.
        let AnnotationValue::Reference(path) = value else {
            return None;
        };
        Value::String(resolve_ref_nsid(path, workspace, current_namespace))
    };
    Some((key.clone(), json_value))
}

/// Literal-to-JSON conversion for `@const` values. Coerces whole-number
/// f64 to i64 (so `@const("revision", 3)` emits `3`, not `3.0`) and
/// transparently stringifies genuinely fractional numbers — ATProto's
/// data model has no floats, so the spec-compliant representation is a
/// string. `key` and `namespace` are used to tag any emitted warnings.
fn annotation_value_to_json(
    value: &AnnotationValue,
    key: &str,
    namespace: &str,
    warnings: &mut Vec<CodegenWarning>,
) -> Value {
    match value {
        AnnotationValue::String(s) => Value::String(s.clone()),
        AnnotationValue::Number(n) => {
            if n.is_finite() && n.fract() == 0.0 && *n >= i64::MIN as f64 && *n <= i64::MAX as f64 {
                Value::Number(serde_json::Number::from(*n as i64))
            } else if n.is_finite() {
                warnings.push(CodegenWarning {
                    namespace: namespace.to_string(),
                    message: format!(
                        "@const({:?}, {}): ATProto's data model has no floats; \
                         emitting as string {:?} to stay spec-compliant",
                        key,
                        n,
                        n.to_string()
                    ),
                });
                Value::String(n.to_string())
            } else {
                warnings.push(CodegenWarning {
                    namespace: namespace.to_string(),
                    message: format!(
                        "@const({:?}, {}): non-finite number is not representable in JSON; \
                         emitting `null`",
                        key, n
                    ),
                });
                Value::Null
            }
        }
        AnnotationValue::Boolean(b) => Value::Bool(*b),
        AnnotationValue::Null => Value::Null,
        AnnotationValue::Array(items) => Value::Array(
            items
                .iter()
                .map(|item| annotation_value_to_json(item, key, namespace, warnings))
                .collect(),
        ),
        AnnotationValue::Object(entries) => {
            let mut obj = Map::new();
            for (entry_key, entry_value) in entries {
                obj.insert(
                    entry_key.clone(),
                    annotation_value_to_json(entry_value, key, namespace, warnings),
                );
            }
            Value::Object(obj)
        }
        // A `Reference` shouldn't appear in a `@const` value position —
        // it would mean the author wrote a bare identifier where they
        // meant a literal. Fall back to emitting the path as-is; the
        // parser could also reject this earlier, but defensively handle
        // it here.
        AnnotationValue::Reference(path) => Value::String(path.to_string()),
    }
}

/// Decides whether a given def should be emitted under the key `main` or
/// under its own name. The rules are lexicon-level — they depend on
/// whether the namespace is a `.defs` container, whether any item carries
/// an explicit `@main`, and how many main-eligible items exist — so we
/// compute the context once and consult it per item.
struct MainEligibility {
    /// True when at least one item has `@main`; in that case only
    /// annotated items are main, and all heuristics are skipped.
    has_explicit_main: bool,
    /// Total count of items eligible to be main.
    eligible_count: usize,
    /// The last segment of the namespace (e.g. `profile` for
    /// `app.bsky.actor.profile`). Used for the "name matches namespace"
    /// heuristic.
    expected_main_name: String,
    /// True when the namespace ends in `.defs`. By convention these
    /// lexicons are containers for named defs and never have a main
    /// entry; heuristic promotion is suppressed here to preserve that
    /// convention on roundtrip.
    is_defs_namespace: bool,
}

impl MainEligibility {
    fn for_lexicon(namespace: &str, lexicon: &Lexicon) -> Self {
        let expected_main_name = namespace.rsplit('.').next().unwrap_or("").to_string();
        let is_defs_namespace = expected_main_name == "defs";

        let headers: Vec<_> = lexicon.items.iter().filter_map(item_header).collect();
        let eligible_count = headers.len();
        let has_explicit_main = headers.iter().any(|(_, ann)| has_main_annotation(ann));

        Self {
            has_explicit_main,
            eligible_count,
            expected_main_name,
            is_defs_namespace,
        }
    }

    fn is_main(&self, name: &str, annotations: &[Annotation]) -> bool {
        if self.has_explicit_main {
            return has_main_annotation(annotations);
        }
        // `.defs` lexicons are pure containers — never promote a def to
        // `main`, not even when it's the only one in the file.
        if self.is_defs_namespace {
            return false;
        }
        self.eligible_count == 1 || name == self.expected_main_name
    }
}

/// Unified accessor for the main-eligible item kinds. Returns the def's
/// declared name and annotations — the only facts the main-promotion
/// decision needs from any given item. Returns `None` for item kinds
/// that never participate in `main` eligibility (e.g. inline types,
/// tokens, use statements).
fn item_header(item: &Item) -> Option<(&str, &[Annotation])> {
    match item {
        Item::Record(r) => Some((&r.name.name, &r.annotations)),
        Item::Query(q) => Some((&q.name.name, &q.annotations)),
        Item::Procedure(p) => Some((&p.name.name, &p.annotations)),
        Item::Subscription(s) => Some((&s.name.name, &s.annotations)),
        Item::DefType(d) => Some((&d.name.name, &d.annotations)),
        _ => None,
    }
}

/// Insert a def into the lexicon's `defs` map under the canonical key —
/// `"main"` for the main def, otherwise the def's own name.
fn insert_def(defs: &mut Map<String, Value>, name: &str, is_main: bool, value: Value) {
    let key = if is_main {
        "main".to_string()
    } else {
        name.to_string()
    };
    defs.insert(key, value);
}

fn analyze_type_usage(lexicon: &Lexicon) -> HashMap<String, usize> {
    let mut usage_counts = HashMap::new();

    for item in &lexicon.items {
        match item {
            Item::Record(record) => {
                for field in &record.fields {
                    count_type_references(&field.ty, &mut usage_counts);
                }
            }
            Item::Query(query) => {
                for param in &query.params {
                    count_type_references(&param.ty, &mut usage_counts);
                }
                match &query.returns {
                    ReturnType::None { .. } => {}
                    ReturnType::Type(ty) => count_type_references(ty, &mut usage_counts),
                    ReturnType::TypeWithErrors { success, .. } => {
                        count_type_references(success, &mut usage_counts)
                    }
                }
            }
            Item::Procedure(procedure) => {
                for param in &procedure.params {
                    count_type_references(&param.ty, &mut usage_counts);
                }
                match &procedure.returns {
                    ReturnType::None { .. } => {}
                    ReturnType::Type(ty) => count_type_references(ty, &mut usage_counts),
                    ReturnType::TypeWithErrors { success, .. } => {
                        count_type_references(success, &mut usage_counts)
                    }
                }
            }
            Item::Subscription(subscription) => {
                for param in &subscription.params {
                    count_type_references(&param.ty, &mut usage_counts);
                }
                if let Some(messages) = &subscription.messages {
                    count_type_references(messages, &mut usage_counts);
                }
            }
            Item::InlineType(inline_type) => {
                count_type_references(&inline_type.ty, &mut usage_counts);
            }
            Item::DefType(def_type) => {
                count_type_references(&def_type.ty, &mut usage_counts);
            }
            _ => {}
        }
    }

    usage_counts
}

fn count_type_references(ty: &Type, counts: &mut HashMap<String, usize>) {
    match ty {
        Type::Reference { path, .. } => {
            if path.segments.len() == 1 {
                let name = &path.segments[0].name;
                *counts.entry(name.clone()).or_insert(0) += 1;
            }
        }
        Type::Array { inner, .. } => count_type_references(inner, counts),
        Type::Union { types, .. } => {
            for t in types {
                count_type_references(t, counts);
            }
        }
        Type::Object { fields, .. } => {
            for field in fields {
                count_type_references(&field.ty, counts);
            }
        }
        Type::Parenthesized { inner, .. } => count_type_references(inner, counts),
        Type::Constrained { base, .. } => count_type_references(base, counts),
        _ => {}
    }
}

fn extract_docs(docs: &[DocComment]) -> String {
    docs.iter()
        .map(|d| d.text.trim())
        .collect::<Vec<_>>()
        .join("\n")
}

///// Insert `key: value` into `map` only when `value` is non-empty. Several
/// ATProto fields (`description`, per-type `format`, etc.) are optional and
/// should be omitted when empty — emitting `""` causes spurious roundtrip
/// diffs against authoritative lexicons.
fn insert_opt_str(map: &mut Map<String, Value>, key: &str, value: &str) {
    if !value.is_empty() {
        map.insert(key.to_string(), Value::String(value.to_string()));
    }
}

/// Insert `key: values` only when the slice is non-empty. Mirrors the
/// ATProto convention for `required`, `enum`, and similar list fields.
fn insert_opt_list(map: &mut Map<String, Value>, key: &str, values: &[String]) {
    if !values.is_empty() {
        map.insert(key.to_string(), json!(values));
    }
}

/// If `docs` render to a non-empty string, insert it as the `description`
/// field of `obj` (which is expected to be a JSON object). No-op for
/// non-object values or empty doc sets.
fn add_description_from_docs(obj: &mut Value, docs: &[DocComment]) {
    let text = extract_docs(docs);
    if text.is_empty() {
        return;
    }
    if let Some(map) = obj.as_object_mut() {
        map.insert("description".to_string(), Value::String(text));
    }
}

/// Materialise a query/procedure/subscription parameter list into
/// `(properties, required)`. Each parameter is emitted through
/// [`generate_type_json`], annotated with its doc comment, and collected
/// into `properties`; the names of non-optional parameters are collected
/// into `required` in source order.
fn build_param_properties(
    params: &[Field],
    usage_counts: &HashMap<String, usize>,
    workspace: &Workspace,
    current_namespace: &str,
) -> (Map<String, Value>, Vec<String>) {
    let mut properties = Map::new();
    let mut required = Vec::new();
    for param in params {
        if !param.optional {
            required.push(param.name.name.clone());
        }
        let mut param_json =
            generate_type_json(&param.ty, usage_counts, workspace, current_namespace);
        add_description_from_docs(&mut param_json, &param.docs);
        properties.insert(param.name.name.clone(), param_json);
    }
    (properties, required)
}

/// Shared iteration for records and inline object types. Returns
/// `(properties, required, nullable)`:
///
/// * `properties` — field-name → emitted JSON type, in declaration order.
/// * `required` — non-optional field names.
/// * `nullable` — fields whose MLF type is `T | null` (ATProto spec encodes
///   nullability as a sibling `nullable` array on the enclosing object,
///   not as a union member). The null is stripped from the emitted type.
fn collect_object_fields(
    fields: &[Field],
    usage_counts: &HashMap<String, usize>,
    workspace: &Workspace,
    current_namespace: &str,
) -> (Map<String, Value>, Vec<String>, Vec<String>) {
    let mut properties = Map::new();
    let mut required = Vec::new();
    let mut nullable = Vec::new();
    for field in fields {
        if !field.optional {
            required.push(field.name.name.clone());
        }
        let (effective_ty, is_nullable) = match strip_nullable(&field.ty) {
            Some(stripped) => (stripped, true),
            None => (field.ty.clone(), false),
        };
        if is_nullable {
            nullable.push(field.name.name.clone());
        }
        let mut field_json =
            generate_type_json(&effective_ty, usage_counts, workspace, current_namespace);
        add_description_from_docs(&mut field_json, &field.docs);
        properties.insert(field.name.name.clone(), field_json);
    }
    (properties, required, nullable)
}

/// If `ty` is a union that includes a `null` primitive, return the type
/// with the null stripped — the ATProto-side lowering of MLF's `T | null`
/// pattern, which is how we express field nullability in the surface
/// language. Returns `None` when `ty` isn't a nullable union.
///
/// Unwraps parenthesized types transparently so `(T | null)` behaves the
/// same as `T | null`. Collapses a single-remaining-member union to just
/// that member.
fn strip_nullable(ty: &Type) -> Option<Type> {
    if let Type::Parenthesized { inner, .. } = ty {
        return strip_nullable(inner);
    }
    let Type::Union {
        types,
        closed,
        span,
    } = ty
    else {
        return None;
    };
    let has_null = types.iter().any(is_null_primitive);
    if !has_null {
        return None;
    }
    let remaining: Vec<Type> = types
        .iter()
        .filter(|t| !is_null_primitive(t))
        .cloned()
        .collect();
    // A single-member union degenerates to the member itself — the same
    // invariant our Lexicon→MLF pass maintains, kept symmetric here.
    if remaining.len() == 1 {
        return Some(remaining.into_iter().next().unwrap());
    }
    Some(Type::Union {
        types: remaining,
        closed: *closed,
        span: *span,
    })
}

fn is_null_primitive(ty: &Type) -> bool {
    match ty {
        Type::Primitive {
            kind: PrimitiveType::Null,
            ..
        } => true,
        Type::Parenthesized { inner, .. } => is_null_primitive(inner),
        _ => false,
    }
}

/// Build a `type: "params"` JSON object from the output of
/// [`build_param_properties`]. `required` is omitted when empty, and an
/// empty `properties` object is still emitted (queries always carry an
/// empty `parameters.properties` per ATProto convention).
fn build_params_object(properties: Map<String, Value>, required: &[String]) -> Value {
    let mut obj = Map::new();
    obj.insert("type".to_string(), json!("params"));
    insert_opt_list(&mut obj, "required", required);
    obj.insert("properties".to_string(), Value::Object(properties));
    Value::Object(obj)
}

/// Like [`build_params_object`] but returns `None` when there are no
/// parameters — used by subscriptions, where the entire `parameters`
/// field is elided rather than emitted as an empty object.
fn build_params_object_opt(properties: Map<String, Value>, required: &[String]) -> Option<Value> {
    if properties.is_empty() {
        None
    } else {
        Some(build_params_object(properties, required))
    }
}

/// Resolve a typed reference path to the ATProto NSID string it names.
///
/// The three valid output forms — local `#foo`, external `ns.path#foo`,
/// and bare `ns.path` (implicit-main) — map one-to-one onto the
/// [`ResolvedRef`] variants, so there's no guessing about the def name
/// or whether to append a fragment. Unresolvable paths (typically a
/// typo caught elsewhere by validation) fall through to a best-effort
/// rendering so the JSON stays self-describing for tooling.
fn resolve_ref_nsid(path: &Path, workspace: &Workspace, current_namespace: &str) -> String {
    match workspace.resolve_ref(path, current_namespace) {
        Some(ResolvedRef::Local { def_name }) => format!("#{}", def_name),
        Some(ResolvedRef::ImplicitMain { namespace, .. }) => namespace,
        Some(ResolvedRef::External {
            namespace,
            def_name,
        }) => format!("{}#{}", namespace, def_name),
        None => unresolved_ref_fallback(path),
    }
}

/// Render a reference path the workspace couldn't resolve. Multi-
/// segment paths are emitted as `ns.path#def_name` on the assumption
/// that the author wrote a full NSID; single-segment paths render as a
/// local `#name` fragment.
fn unresolved_ref_fallback(path: &Path) -> String {
    if path.segments.len() == 1 {
        return format!("#{}", path.segments[0].name);
    }
    let namespace = path.segments[..path.segments.len() - 1]
        .iter()
        .map(|s| s.name.as_str())
        .collect::<Vec<_>>()
        .join(".");
    let def_name = &path.segments.last().unwrap().name;
    format!("{}#{}", namespace, def_name)
}

fn generate_record_json(
    record: &Record,
    usage_counts: &HashMap<String, usize>,
    workspace: &Workspace,
    current_namespace: &str,
) -> Value {
    let (properties, required, nullable) =
        collect_object_fields(&record.fields, usage_counts, workspace, current_namespace);

    let mut record_obj = Map::new();
    record_obj.insert("type".to_string(), json!("object"));
    insert_opt_list(&mut record_obj, "required", &required);
    insert_opt_list(&mut record_obj, "nullable", &nullable);
    record_obj.insert("properties".to_string(), Value::Object(properties));

    // Check for @key annotation, default to "tid"
    let key = get_annotation_string_value(&record.annotations, "key")
        .unwrap_or_else(|| "tid".to_string());

    let mut record_top = Map::new();
    record_top.insert("type".to_string(), json!("record"));
    insert_opt_str(&mut record_top, "description", &extract_docs(&record.docs));
    record_top.insert("key".to_string(), Value::String(key));
    record_top.insert("record".to_string(), Value::Object(record_obj));
    Value::Object(record_top)
}

fn generate_query_json(
    query: &Query,
    usage_counts: &HashMap<String, usize>,
    workspace: &Workspace,
    current_namespace: &str,
) -> Value {
    let (params_properties, params_required) =
        build_param_properties(&query.params, usage_counts, workspace, current_namespace);
    let params = build_params_object(params_properties, &params_required);

    // Check for @encoding annotation (output only for queries), default to "application/json"
    let output_encoding = get_encoding_annotation(&query.annotations, "output")
        .unwrap_or_else(|| "application/json".to_string());

    let (output, errors) = match &query.returns {
        ReturnType::None { .. } => (None, None),
        ReturnType::Type(ty) => {
            let mut output_obj = Map::new();
            output_obj.insert("encoding".to_string(), json!(output_encoding));
            output_obj.insert(
                "schema".to_string(),
                generate_type_json(ty, usage_counts, workspace, current_namespace),
            );
            (Some(Value::Object(output_obj)), None)
        }
        ReturnType::TypeWithErrors {
            success, errors, ..
        } => {
            let mut error_array = Vec::new();
            for error in errors {
                let error_docs = extract_docs(&error.docs);
                let error_obj = if error_docs.is_empty() {
                    json!({ "name": error.name.name.clone() })
                } else {
                    json!({
                        "name": error.name.name.clone(),
                        "description": error_docs
                    })
                };
                error_array.push(error_obj);
            }

            let mut output_obj = Map::new();
            output_obj.insert("encoding".to_string(), json!(output_encoding));
            output_obj.insert(
                "schema".to_string(),
                generate_type_json(success, usage_counts, workspace, current_namespace),
            );
            (
                Some(Value::Object(output_obj)),
                Some(Value::Array(error_array)),
            )
        }
    };

    let mut query_obj = Map::new();
    query_obj.insert("type".to_string(), json!("query"));
    insert_opt_str(&mut query_obj, "description", &extract_docs(&query.docs));
    query_obj.insert("parameters".to_string(), params);
    if let Some(output_val) = output {
        query_obj.insert("output".to_string(), output_val);
    }
    if let Some(errors_val) = errors {
        query_obj.insert("errors".to_string(), errors_val);
    }
    Value::Object(query_obj)
}

fn generate_procedure_json(
    procedure: &Procedure,
    usage_counts: &HashMap<String, usize>,
    workspace: &Workspace,
    current_namespace: &str,
) -> Value {
    let (params_properties, params_required) = build_param_properties(
        &procedure.params,
        usage_counts,
        workspace,
        current_namespace,
    );

    // Check for @encoding annotation with "input" parameter, default to "application/json"
    let input_encoding = get_encoding_annotation(&procedure.annotations, "input")
        .unwrap_or_else(|| "application/json".to_string());

    let input = if !params_properties.is_empty() {
        let mut schema_obj = Map::new();
        schema_obj.insert("type".to_string(), json!("object"));
        insert_opt_list(&mut schema_obj, "required", &params_required);
        schema_obj.insert("properties".to_string(), Value::Object(params_properties));

        let mut input_obj = Map::new();
        input_obj.insert("encoding".to_string(), json!(input_encoding));
        input_obj.insert("schema".to_string(), Value::Object(schema_obj));
        Some(Value::Object(input_obj))
    } else {
        None
    };

    // Check for @encoding annotation with "output" parameter, default to "application/json"
    let output_encoding = get_encoding_annotation(&procedure.annotations, "output")
        .unwrap_or_else(|| "application/json".to_string());

    let (output, errors) = match &procedure.returns {
        ReturnType::None { .. } => (None, None),
        ReturnType::Type(ty) => {
            let mut output_obj = Map::new();
            output_obj.insert("encoding".to_string(), json!(output_encoding));
            output_obj.insert(
                "schema".to_string(),
                generate_type_json(ty, usage_counts, workspace, current_namespace),
            );
            (Some(Value::Object(output_obj)), None)
        }
        ReturnType::TypeWithErrors {
            success, errors, ..
        } => {
            let mut error_array = Vec::new();
            for error in errors {
                let error_docs = extract_docs(&error.docs);
                let error_obj = if error_docs.is_empty() {
                    json!({ "name": error.name.name.clone() })
                } else {
                    json!({
                        "name": error.name.name.clone(),
                        "description": error_docs
                    })
                };
                error_array.push(error_obj);
            }

            let mut output_obj = Map::new();
            output_obj.insert("encoding".to_string(), json!(output_encoding));
            output_obj.insert(
                "schema".to_string(),
                generate_type_json(success, usage_counts, workspace, current_namespace),
            );
            (
                Some(Value::Object(output_obj)),
                Some(Value::Array(error_array)),
            )
        }
    };

    let mut result = Map::new();
    result.insert("type".to_string(), json!("procedure"));
    insert_opt_str(&mut result, "description", &extract_docs(&procedure.docs));
    if let Some(input_val) = input {
        result.insert("input".to_string(), input_val);
    }
    if let Some(output_val) = output {
        result.insert("output".to_string(), output_val);
    }
    if let Some(errors_val) = errors {
        result.insert("errors".to_string(), errors_val);
    }
    Value::Object(result)
}

fn generate_subscription_json(
    subscription: &Subscription,
    usage_counts: &HashMap<String, usize>,
    workspace: &Workspace,
    current_namespace: &str,
) -> Value {
    let (params_properties, params_required) = build_param_properties(
        &subscription.params,
        usage_counts,
        workspace,
        current_namespace,
    );

    let mut result = Map::new();
    result.insert("type".to_string(), json!("subscription"));
    insert_opt_str(
        &mut result,
        "description",
        &extract_docs(&subscription.docs),
    );
    let mut result = Value::Object(result);

    if let Some(messages) = &subscription.messages {
        let message = json!({
            "schema": generate_type_json(messages, usage_counts, workspace, current_namespace)
        });
        result["message"] = message;
    }

    if let Some(parameters) = build_params_object_opt(params_properties, &params_required) {
        result["parameters"] = parameters;
    }

    result
}

fn generate_def_type_json(
    def_type: &DefType,
    usage_counts: &HashMap<String, usize>,
    workspace: &Workspace,
    current_namespace: &str,
) -> Value {
    let mut field_json =
        generate_type_json(&def_type.ty, usage_counts, workspace, current_namespace);
    // def types insert `description` at position 1 (right after `type`) to
    // match the canonical field order in published lexicons.
    let text = extract_docs(&def_type.docs);
    if !text.is_empty() {
        if let Some(obj) = field_json.as_object_mut() {
            obj.shift_insert(1, "description".to_string(), Value::String(text));
        }
    }
    field_json
}

fn generate_type_json(
    ty: &Type,
    usage_counts: &HashMap<String, usize>,
    workspace: &Workspace,
    current_namespace: &str,
) -> Value {
    match ty {
        Type::Primitive { kind, .. } => generate_primitive_json(*kind),
        Type::Reference { path, .. } => {
            // Inline types expand in place; other references resolve to an NSID.
            if workspace.is_inline_type(path) {
                if let Some(resolved_ty) = workspace.resolve_type_reference(path) {
                    return generate_type_json(
                        &resolved_ty,
                        usage_counts,
                        workspace,
                        current_namespace,
                    );
                }
            }
            let nsid = resolve_ref_nsid(path, workspace, current_namespace);
            json!({ "type": "ref", "ref": nsid })
        }
        Type::Array { inner, .. } => {
            json!({
                "type": "array",
                "items": generate_type_json(inner, usage_counts, workspace, current_namespace)
            })
        }
        Type::Union { types, closed, .. } => {
            // ATProto unions carry an array of NSID *strings*. For ref
            // members we resolve the NSID directly; for non-ref members
            // (e.g. `null` produced by MLF's nullable-via-union pattern)
            // we fall through to the typed form. That fallback is
            // non-spec — ATProto unions strictly contain refs — but
            // preserves the member's semantics until a proper nullable
            // lowering is added upstream.
            let refs: Vec<Value> = types
                .iter()
                .map(|t| match t {
                    Type::Reference { path, .. } => {
                        Value::String(resolve_ref_nsid(path, workspace, current_namespace))
                    }
                    _ => generate_type_json(t, usage_counts, workspace, current_namespace),
                })
                .collect();
            let mut union_obj = Map::new();
            union_obj.insert("type".to_string(), json!("union"));
            union_obj.insert("refs".to_string(), Value::Array(refs));
            // Only emit "closed" when true; open unions omit it per spec.
            if *closed {
                union_obj.insert("closed".to_string(), json!(true));
            }
            Value::Object(union_obj)
        }
        Type::Object { fields, .. } => {
            let (properties, required, nullable) =
                collect_object_fields(fields, usage_counts, workspace, current_namespace);

            let mut obj = Map::new();
            obj.insert("type".to_string(), json!("object"));
            insert_opt_list(&mut obj, "required", &required);
            insert_opt_list(&mut obj, "nullable", &nullable);
            obj.insert("properties".to_string(), Value::Object(properties));
            Value::Object(obj)
        }
        Type::Parenthesized { inner, .. } => {
            // Parentheses are just for grouping - unwrap and process inner type
            generate_type_json(inner, usage_counts, workspace, current_namespace)
        }
        Type::Constrained {
            base, constraints, ..
        } => {
            let mut base_json =
                generate_type_json(base, usage_counts, workspace, current_namespace);

            if let Some(obj) = base_json.as_object_mut() {
                for constraint in constraints {
                    apply_constraint_to_json(obj, constraint);
                }
            }

            base_json
        }
        Type::Unknown { .. } => {
            json!({ "type": "unknown" })
        }
    }
}

fn generate_primitive_json(kind: PrimitiveType) -> Value {
    match kind {
        PrimitiveType::Null => json!({ "type": "null" }),
        PrimitiveType::Boolean => json!({ "type": "boolean" }),
        PrimitiveType::Integer => json!({ "type": "integer" }),
        PrimitiveType::String => json!({ "type": "string" }),
        PrimitiveType::Bytes => json!({ "type": "bytes" }),
        PrimitiveType::Blob => json!({ "type": "blob" }),
    }
}

fn apply_constraint_to_json(obj: &mut Map<String, Value>, constraint: &Constraint) {
    match constraint {
        Constraint::MinLength { value, .. } => {
            obj.insert("minLength".to_string(), json!(value));
        }
        Constraint::MaxLength { value, .. } => {
            obj.insert("maxLength".to_string(), json!(value));
        }
        Constraint::MinGraphemes { value, .. } => {
            obj.insert("minGraphemes".to_string(), json!(value));
        }
        Constraint::MaxGraphemes { value, .. } => {
            obj.insert("maxGraphemes".to_string(), json!(value));
        }
        Constraint::Minimum { value, .. } => {
            obj.insert("minimum".to_string(), json!(value));
        }
        Constraint::Maximum { value, .. } => {
            obj.insert("maximum".to_string(), json!(value));
        }
        Constraint::Format { value, .. } => {
            obj.insert("format".to_string(), json!(value));
        }
        Constraint::Enum { values, .. } => {
            let enum_vals: Vec<String> = values
                .iter()
                .map(|v| match v {
                    mlf_lang::ast::ValueRef::Literal(s) => s.clone(),
                    mlf_lang::ast::ValueRef::Reference(path) => path.to_string(),
                })
                .collect();
            obj.insert("enum".to_string(), json!(enum_vals));
        }
        Constraint::KnownValues { values, .. } => {
            let known_vals: Vec<String> = values
                .iter()
                .map(|v| match v {
                    mlf_lang::ast::ValueRef::Literal(s) => s.clone(),
                    mlf_lang::ast::ValueRef::Reference(path) => path.to_string(),
                })
                .collect();
            obj.insert("knownValues".to_string(), json!(known_vals));
        }
        Constraint::Accept { mimes, .. } => {
            obj.insert("accept".to_string(), json!(mimes));
        }
        Constraint::MaxSize { value, .. } => {
            obj.insert("maxSize".to_string(), json!(value));
        }
        Constraint::Default { value, .. } => {
            let default_val = match value {
                ConstraintValue::String(s) => json!(s),
                ConstraintValue::Integer(i) => json!(i),
                ConstraintValue::Boolean(b) => json!(b),
                ConstraintValue::Reference(path) => json!(path.to_string()),
            };
            obj.insert("default".to_string(), default_val);
        }
        Constraint::Const { value, .. } => {
            let const_val = match value {
                ConstraintValue::String(s) => json!(s),
                ConstraintValue::Integer(i) => json!(i),
                ConstraintValue::Boolean(b) => json!(b),
                ConstraintValue::Reference(path) => json!(path.to_string()),
            };
            obj.insert("const".to_string(), const_val);
        }
    }
}

// Example built-in generator: JSON Lexicon
pub struct JsonLexiconGenerator;

impl CodeGenerator for JsonLexiconGenerator {
    fn name(&self) -> &'static str {
        "json"
    }

    fn description(&self) -> &'static str {
        "Generate AT Protocol JSON lexicon format"
    }

    fn file_extension(&self) -> &'static str {
        ".json"
    }

    fn generate(&self, ctx: &GeneratorContext) -> Result<String, String> {
        // The plugin `CodeGenerator` trait doesn't carry a warning
        // channel today; the callers that care about warnings call
        // `generate_lexicon` directly instead of going through this shim.
        let output = generate_lexicon(ctx.namespace, ctx.lexicon, ctx.workspace);
        serde_json::to_string_pretty(&output.json)
            .map_err(|e| format!("Failed to serialize JSON: {}", e))
    }
}

// Register the JSON generator as a static instance
static JSON_GENERATOR: JsonLexiconGenerator = JsonLexiconGenerator;
register_generator!(JSON_GENERATOR);