//! Generates Rust code under `crates/tangled-api/src/lexicon/` from the
//! vendored lexicon JSON files in `lexicons/`.
//!
//! Design rules:
//! - NSID `sh.tangled.repo.pull` maps to module `lexicon::sh::tangled::repo::pull`.
//! - `main` record defs become `struct Record` (+ `const NSID`).
//! - `main` object defs become a struct named after the last NSID segment.
//! - Non-main object defs become structs named after the def name.
//! - Queries/procedures emit `Params`/`Input`/`Output` types and a free
//!   function returning a `reqwest::RequestBuilder`, re-exported from the
//!   parent module.
//! - Anything the generator does not understand degrades to
//!   `serde_json::Value` with a warning; it never fails.

use std::collections::{BTreeMap, BTreeSet};
use std::fmt::Write as _;
use std::path::{Path, PathBuf};

use anyhow::{Context, Result};
use heck::{ToSnakeCase, ToUpperCamelCase};

use crate::lexicon::{BodyDef, LexDef, LexType, LexiconDoc, ObjectDef, ParamsDef, XrpcDef};

const RUST_KEYWORDS: &[&str] = &[
    "as", "break", "const", "continue", "crate", "dyn", "else", "enum", "extern", "false", "fn",
    "for", "if", "impl", "in", "let", "loop", "match", "mod", "move", "mut", "pub", "ref",
    "return", "static", "struct", "super", "trait", "true", "type", "unsafe", "use", "where",
    "while", "async", "await", "box", "try", "self",
];

/// Generated output: file paths relative to the `lexicon/` module root.
pub struct Generated {
    pub files: BTreeMap<PathBuf, String>,
    pub warnings: Vec<String>,
}

struct Doc {
    doc: LexiconDoc,
    /// Path of the source JSON relative to the lexicons dir.
    source: String,
    /// Snake-cased module segments derived from the NSID.
    mods: Vec<String>,
}

/// Maps `nsid` / `nsid#def` to a fully-qualified Rust type path.
struct TypeIndex {
    types: BTreeMap<String, String>,
    /// Token defs: refs to these become `String`.
    tokens: BTreeSet<String>,
}

impl TypeIndex {
    /// Resolves a lexicon ref (`#def`, `nsid`, `nsid#def`) to a Rust type.
    /// `current_id` is the NSID of the doc the ref appears in.
    fn resolve(&self, r: &str, current_id: &str) -> Option<String> {
        let key = if let Some(frag) = r.strip_prefix('#') {
            format!("{current_id}#{frag}")
        } else {
            r.to_string()
        };
        if self.tokens.contains(&key) {
            return Some("String".to_string());
        }
        self.types.get(&key).cloned()
    }
}

fn snake(seg: &str) -> String {
    let s = seg.to_snake_case();
    if RUST_KEYWORDS.contains(&s.as_str()) {
        format!("{s}_")
    } else {
        s
    }
}

/// Type names that would shadow prelude/std types used by the generated code.
const RESERVED_TYPE_NAMES: &[&str] = &["String", "Option", "Vec", "Box", "Default", "Clone", "Debug"];

fn pascal(name: &str) -> String {
    let s = name.to_upper_camel_case();
    if RESERVED_TYPE_NAMES.contains(&s.as_str()) {
        format!("{s}Def")
    } else {
        s
    }
}

fn field_ident(prop: &str) -> String {
    let s = prop.to_snake_case();
    if RUST_KEYWORDS.contains(&s.as_str()) {
        // `self` cannot be a raw identifier; everything else here can.
        if s == "self" {
            "self_".to_string()
        } else {
            format!("r#{s}")
        }
    } else {
        s
    }
}

fn doc_comment(out: &mut String, text: Option<&str>) {
    if let Some(text) = text {
        for line in text.lines() {
            let _ = writeln!(out, "/// {}", line.trim_end());
        }
    }
}

fn mods_of_nsid(id: &str) -> Vec<String> {
    id.split('.').map(snake).collect()
}

fn main_type_name(id: &str, def: &LexDef) -> Option<String> {
    match def {
        LexDef::Record(_) => Some("Record".to_string()),
        LexDef::Object(_) => Some(pascal(id.rsplit('.').next().unwrap())),
        _ => None,
    }
}

pub fn generate(lexicons_dir: &Path) -> Result<Generated> {
    let mut docs = Vec::new();
    let mut warnings = Vec::new();
    let mut json_files = Vec::new();
    collect_json_files(lexicons_dir, &mut json_files)?;
    json_files.sort();

    for path in &json_files {
        let text = std::fs::read_to_string(path)
            .with_context(|| format!("reading {}", path.display()))?;
        let doc: LexiconDoc = serde_json::from_str(&text)
            .with_context(|| format!("parsing {}", path.display()))?;
        let source = path
            .strip_prefix(lexicons_dir)
            .unwrap()
            .to_string_lossy()
            .replace('\\', "/");
        if matches!(doc.defs.get("main"), Some(LexDef::Subscription(_))) {
            warnings.push(format!(
                "skipping {} ({}): subscriptions are not supported",
                source, doc.id
            ));
            continue;
        }
        let mods = mods_of_nsid(&doc.id);
        docs.push(Doc { doc, source, mods });
    }

    // Build the type index across all docs.
    let mut index = TypeIndex {
        types: BTreeMap::new(),
        tokens: BTreeSet::new(),
    };
    for d in &docs {
        let rust_mod = format!("crate::lexicon::{}", d.mods.join("::"));
        for (name, def) in &d.doc.defs {
            let key = if name == "main" {
                d.doc.id.clone()
            } else {
                format!("{}#{name}", d.doc.id)
            };
            match def {
                LexDef::Token(_) => {
                    index.tokens.insert(key);
                }
                LexDef::Record(_) | LexDef::Object(_) => {
                    let type_name = if name == "main" {
                        main_type_name(&d.doc.id, def).unwrap()
                    } else {
                        pascal(name)
                    };
                    index.types.insert(key, format!("{rust_mod}::{type_name}"));
                }
                LexDef::Type(_) => {
                    let type_name = pascal(name);
                    index.types.insert(key, format!("{rust_mod}::{type_name}"));
                }
                _ => {}
            }
        }
    }

    // Build the module tree.
    #[derive(Default)]
    struct Node {
        children: BTreeMap<String, Node>,
        doc: Option<usize>,
    }
    let mut root = Node::default();
    for (i, d) in docs.iter().enumerate() {
        let mut node = &mut root;
        for seg in &d.mods {
            node = node.children.entry(seg.clone()).or_default();
        }
        node.doc = Some(i);
    }

    // Emit files.
    let mut files = BTreeMap::new();
    fn emit_node(
        node: &Node,
        path: &[String],
        docs: &[Doc],
        index: &TypeIndex,
        files: &mut BTreeMap<PathBuf, String>,
        warnings: &mut Vec<String>,
    ) {
        let file_path = if node.children.is_empty() {
            // Leaf: <parent...>/<seg>.rs
            let mut p = PathBuf::new();
            for seg in &path[..path.len() - 1] {
                p.push(seg);
            }
            p.push(format!("{}.rs", path.last().unwrap()));
            p
        } else {
            let mut p = PathBuf::new();
            for seg in path {
                p.push(seg);
            }
            p.push("mod.rs");
            p
        };

        let mut out = String::new();
        out.push_str("// @generated by `cargo xtask codegen` \u{2014} do not edit.\n");
        if let Some(i) = node.doc {
            let d = &docs[i];
            let _ = writeln!(out, "// source: lexicons/{} ({})", d.source, d.doc.id);
        }
        out.push('\n');
        out.push_str("#![allow(clippy::all)]\n\n");

        for (seg, child) in &node.children {
            let _ = writeln!(out, "pub mod {seg};");
            // Re-export builder functions for xrpc children so call sites read
            // `repo::create(...)` instead of `repo::create::create(...)`.
            if let Some(ci) = child.doc {
                if matches!(
                    docs[ci].doc.defs.get("main"),
                    Some(LexDef::Query(_) | LexDef::Procedure(_))
                ) {
                    let _ = writeln!(out, "pub use {seg}::{seg};");
                }
            }
        }
        if !node.children.is_empty() {
            out.push('\n');
        }

        if let Some(i) = node.doc {
            emit_doc(&docs[i], index, &mut out, warnings);
        }

        files.insert(file_path, out);

        for (seg, child) in &node.children {
            let mut child_path = path.to_vec();
            child_path.push(seg.clone());
            emit_node(child, &child_path, docs, index, files, warnings);
        }
    }

    // Root mod.rs.
    {
        let mut out = String::new();
        out.push_str("// @generated by `cargo xtask codegen` \u{2014} do not edit.\n\n");
        for seg in root.children.keys() {
            let _ = writeln!(out, "pub mod {seg};");
        }
        files.insert(PathBuf::from("mod.rs"), out);
    }
    for (seg, child) in &root.children {
        emit_node(
            child,
            std::slice::from_ref(seg),
            &docs,
            &index,
            &mut files,
            &mut warnings,
        );
    }

    Ok(Generated { files, warnings })
}

fn collect_json_files(dir: &Path, out: &mut Vec<PathBuf>) -> Result<()> {
    for entry in std::fs::read_dir(dir).with_context(|| format!("reading {}", dir.display()))? {
        let entry = entry?;
        let path = entry.path();
        if path.is_dir() {
            collect_json_files(&path, out)?;
        } else if path.extension().is_some_and(|e| e == "json") {
            out.push(path);
        }
    }
    Ok(())
}

fn emit_doc(d: &Doc, index: &TypeIndex, out: &mut String, warnings: &mut Vec<String>) {
    let mut used_names: BTreeSet<String> = BTreeSet::new();
    // Reserve names for all defs up front so nested-struct synthesis can't
    // collide with them.
    for (name, def) in &d.doc.defs {
        let n = if name == "main" {
            main_type_name(&d.doc.id, def)
        } else {
            Some(pascal(name))
        };
        if let Some(n) = n {
            if !used_names.insert(n.clone()) {
                warnings.push(format!("{}: duplicate type name {n}", d.doc.id));
            }
        }
    }

    for (name, def) in &d.doc.defs {
        match def {
            LexDef::Record(r) => {
                let _ = writeln!(out, "pub const NSID: &str = \"{}\";\n", d.doc.id);
                doc_comment(out, r.description.as_deref());
                emit_struct(
                    d,
                    index,
                    out,
                    warnings,
                    &mut used_names,
                    "Record",
                    &r.record,
                    StructKind::Record,
                );
            }
            LexDef::Object(o) => {
                let type_name = if name == "main" {
                    main_type_name(&d.doc.id, def).unwrap()
                } else {
                    pascal(name)
                };
                emit_struct(
                    d,
                    index,
                    out,
                    warnings,
                    &mut used_names,
                    &type_name,
                    o,
                    StructKind::Embedded,
                );
            }
            LexDef::Token(t) => {
                doc_comment(out, t.description.as_deref());
                if name == "main" {
                    let _ = writeln!(out, "pub const NSID: &str = \"{}\";\n", d.doc.id);
                } else {
                    let const_name = name.to_snake_case().to_uppercase();
                    let _ = writeln!(
                        out,
                        "pub const {const_name}: &str = \"{}#{name}\";\n",
                        d.doc.id
                    );
                }
            }
            LexDef::Query(x) | LexDef::Procedure(x) => {
                if name != "main" {
                    warnings.push(format!(
                        "{}: non-main query/procedure def {name} skipped",
                        d.doc.id
                    ));
                    continue;
                }
                let is_query = matches!(def, LexDef::Query(_));
                emit_xrpc(d, index, out, warnings, &mut used_names, x, is_query);
            }
            LexDef::Subscription(_) => {
                warnings.push(format!("{}: subscription def {name} skipped", d.doc.id));
            }
            LexDef::Type(t) => {
                let type_name = pascal(name);
                doc_comment(out, t.description());
                let mut nested = Vec::new();
                let rust = rust_type(d, index, t, &type_name, warnings, &mut nested);
                let _ = writeln!(out, "pub type {type_name} = {rust};\n");
                emit_nested(
                    d,
                    index,
                    out,
                    warnings,
                    &mut used_names,
                    nested,
                    StructKind::Embedded,
                );
            }
        }
    }
}

#[derive(Clone, Copy, PartialEq)]
enum StructKind {
    /// A record or an object embedded in records: preserves unknown fields
    /// via a flattened `extra` map so read-modify-write round-trips are safe.
    Record,
    Embedded,
    /// Request/response bodies and query params: no `extra` map.
    Io,
}

struct Nested<'a> {
    name: String,
    def: &'a ObjectDef,
}

#[allow(clippy::too_many_arguments)]
fn emit_struct(
    d: &Doc,
    index: &TypeIndex,
    out: &mut String,
    warnings: &mut Vec<String>,
    used_names: &mut BTreeSet<String>,
    name: &str,
    def: &ObjectDef,
    kind: StructKind,
) {
    doc_comment(out, def.description.as_deref());
    out.push_str("#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]\n");
    let _ = writeln!(out, "pub struct {name} {{");

    let mut nested = Vec::new();
    let has_type_prop = def.properties.contains_key("type");

    if kind == StructKind::Record && !has_type_prop {
        out.push_str("    /// Record type discriminator (`$type`); set to [`NSID`] when writing.\n");
        out.push_str(
            "    #[serde(rename = \"$type\", skip_serializing_if = \"Option::is_none\", default)]\n",
        );
        out.push_str("    pub r#type: Option<String>,\n");
    } else if kind == StructKind::Record {
        warnings.push(format!(
            "{}: record has a `type` property; no $type field emitted",
            d.doc.id
        ));
    }

    for (prop, lex) in &def.properties {
        let ident = field_ident(prop);
        let optional = !def.required.contains(prop) || def.nullable.contains(prop);
        let parent_hint = format!("{name}{}", pascal(prop));
        let base_ty = rust_type(d, index, lex, &parent_hint, warnings, &mut nested);
        let ty = if optional {
            format!("Option<{base_ty}>")
        } else {
            base_ty
        };

        if let Some(desc) = lex.description() {
            for line in desc.lines() {
                let _ = writeln!(out, "    /// {}", line.trim_end());
            }
        }
        let mut serde_parts = Vec::new();
        if ident.trim_start_matches("r#") != prop {
            serde_parts.push(format!("rename = \"{prop}\""));
        }
        if optional {
            serde_parts.push("skip_serializing_if = \"Option::is_none\"".to_string());
            serde_parts.push("default".to_string());
        }
        if !serde_parts.is_empty() {
            let _ = writeln!(out, "    #[serde({})]", serde_parts.join(", "));
        }
        let _ = writeln!(out, "    pub {ident}: {ty},");
    }

    if kind != StructKind::Io {
        out.push_str("    /// Unknown fields, preserved across read-modify-write round-trips.\n");
        out.push_str("    #[serde(flatten)]\n");
        out.push_str("    pub extra: serde_json::Map<String, serde_json::Value>,\n");
    }
    out.push_str("}\n\n");

    let nested_kind = if kind == StructKind::Io {
        StructKind::Io
    } else {
        StructKind::Embedded
    };
    emit_nested(d, index, out, warnings, used_names, nested, nested_kind);
}

fn emit_nested<'a>(
    d: &Doc,
    index: &TypeIndex,
    out: &mut String,
    warnings: &mut Vec<String>,
    used_names: &mut BTreeSet<String>,
    nested: Vec<Nested<'a>>,
    kind: StructKind,
) {
    for n in nested {
        if !used_names.insert(n.name.clone()) {
            warnings.push(format!(
                "{}: nested struct name collision: {}",
                d.doc.id, n.name
            ));
            continue;
        }
        emit_struct(d, index, out, warnings, used_names, &n.name, n.def, kind);
    }
}

/// Maps a lexicon type to a Rust type, collecting inline objects that need
/// their own struct definitions.
fn rust_type<'a>(
    d: &Doc,
    index: &TypeIndex,
    lex: &'a LexType,
    name_hint: &str,
    warnings: &mut Vec<String>,
    nested: &mut Vec<Nested<'a>>,
) -> String {
    match lex {
        LexType::String { .. } => "String".to_string(),
        LexType::Integer { .. } => "i64".to_string(),
        LexType::Boolean { .. } => "bool".to_string(),
        LexType::Array { items, .. } => {
            let inner = rust_type(d, index, items, name_hint, warnings, nested);
            format!("Vec<{inner}>")
        }
        LexType::Object(o) => {
            nested.push(Nested {
                name: name_hint.to_string(),
                def: o,
            });
            name_hint.to_string()
        }
        LexType::Ref { target } => match index.resolve(target, &d.doc.id) {
            Some(ty) => ty,
            None => {
                warnings.push(format!(
                    "{}: unresolved ref {target}; using serde_json::Value",
                    d.doc.id
                ));
                "serde_json::Value".to_string()
            }
        },
        // TODO: unions could become $type-tagged enums.
        LexType::Union { .. } => "serde_json::Value".to_string(),
        LexType::Blob { .. } => "crate::types::BlobRef".to_string(),
        LexType::Unknown { .. } | LexType::Bytes { .. } | LexType::CidLink { .. } => {
            "serde_json::Value".to_string()
        }
        LexType::Other(v) => {
            warnings.push(format!(
                "{}: unsupported lexicon type {}; using serde_json::Value",
                d.doc.id,
                v.get("type").and_then(|t| t.as_str()).unwrap_or("?")
            ));
            "serde_json::Value".to_string()
        }
    }
}

fn emit_params(
    d: &Doc,
    index: &TypeIndex,
    out: &mut String,
    warnings: &mut Vec<String>,
    used_names: &mut BTreeSet<String>,
    params: &ParamsDef,
) {
    doc_comment(out, params.description.as_deref());
    out.push_str("#[derive(Debug, Clone, Default, serde::Serialize)]\n");
    out.push_str("pub struct Params {\n");
    let mut nested = Vec::new();
    for (prop, lex) in &params.properties {
        let ident = field_ident(prop);
        let optional = !params.required.contains(prop);
        let base_ty = rust_type(
            d,
            index,
            lex,
            &format!("Params{}", pascal(prop)),
            warnings,
            &mut nested,
        );
        let ty = if optional {
            format!("Option<{base_ty}>")
        } else {
            base_ty
        };
        if let Some(desc) = lex.description() {
            for line in desc.lines() {
                let _ = writeln!(out, "    /// {}", line.trim_end());
            }
        }
        let mut serde_parts = Vec::new();
        if ident.trim_start_matches("r#") != prop {
            serde_parts.push(format!("rename = \"{prop}\""));
        }
        if optional {
            serde_parts.push("skip_serializing_if = \"Option::is_none\"".to_string());
        }
        if !serde_parts.is_empty() {
            let _ = writeln!(out, "    #[serde({})]", serde_parts.join(", "));
        }
        let _ = writeln!(out, "    pub {ident}: {ty},");
    }
    out.push_str("}\n\n");
    emit_nested(d, index, out, warnings, used_names, nested, StructKind::Io);
}

enum BodyKind {
    None,
    Json(String), // Rust type name for the body
    Binary,
}

fn body_kind(
    d: &Doc,
    index: &TypeIndex,
    out: &mut String,
    warnings: &mut Vec<String>,
    used_names: &mut BTreeSet<String>,
    body: &Option<BodyDef>,
    type_name: &str,
) -> BodyKind {
    let Some(body) = body else {
        return BodyKind::None;
    };
    if body.encoding != "application/json" {
        return BodyKind::Binary;
    }
    match &body.schema {
        Some(LexType::Object(o)) => {
            emit_struct(
                d,
                index,
                out,
                warnings,
                used_names,
                type_name,
                o,
                StructKind::Io,
            );
            BodyKind::Json(type_name.to_string())
        }
        Some(other) => {
            let mut nested = Vec::new();
            let rust = rust_type(d, index, other, type_name, warnings, &mut nested);
            if rust != type_name {
                let _ = writeln!(out, "pub type {type_name} = {rust};\n");
            }
            emit_nested(d, index, out, warnings, used_names, nested, StructKind::Io);
            BodyKind::Json(type_name.to_string())
        }
        None => {
            warnings.push(format!(
                "{}: application/json body without schema; treating as untyped",
                d.doc.id
            ));
            let _ = writeln!(out, "pub type {type_name} = serde_json::Value;\n");
            BodyKind::Json(type_name.to_string())
        }
    }
}

fn emit_xrpc(
    d: &Doc,
    index: &TypeIndex,
    out: &mut String,
    warnings: &mut Vec<String>,
    used_names: &mut BTreeSet<String>,
    x: &XrpcDef,
    is_query: bool,
) {
    let _ = writeln!(out, "pub const NSID: &str = \"{}\";\n", d.doc.id);

    if let Some(params) = &x.parameters {
        used_names.insert("Params".to_string());
        emit_params(d, index, out, warnings, used_names, params);
    }
    used_names.insert("Input".to_string());
    used_names.insert("Output".to_string());
    let input = body_kind(d, index, out, warnings, used_names, &x.input, "Input");
    let output = body_kind(d, index, out, warnings, used_names, &x.output, "Output");

    let fn_name = snake(d.doc.id.rsplit('.').next().unwrap());

    doc_comment(out, x.description.as_deref());
    if !x.errors.is_empty() {
        out.push_str("///\n/// Errors:\n");
        for e in &x.errors {
            match &e.description {
                Some(desc) => {
                    let _ = writeln!(out, "/// - `{}`: {}", e.name, desc);
                }
                None => {
                    let _ = writeln!(out, "/// - `{}`", e.name);
                }
            }
        }
    }
    match &output {
        BodyKind::Binary => {
            out.push_str("///\n/// The response body is binary; use [`crate::xrpc::send_bytes`].\n")
        }
        BodyKind::None => {
            out.push_str("///\n/// The response has no body; use [`crate::xrpc::send_unit`].\n")
        }
        BodyKind::Json(_) => {}
    }

    let mut args = vec![
        "http: &reqwest::Client".to_string(),
        "base: &str".to_string(),
    ];
    if x.parameters.is_some() {
        args.push("params: &Params".to_string());
    }
    match &input {
        BodyKind::Json(ty) => args.push(format!("input: &{ty}")),
        BodyKind::Binary => {
            args.push("body: reqwest::Body".to_string());
            args.push("content_type: &str".to_string());
        }
        BodyKind::None => {}
    }

    let _ = writeln!(
        out,
        "pub fn {fn_name}({}) -> reqwest::RequestBuilder {{",
        args.join(", ")
    );
    let verb = if is_query { "get" } else { "post" };
    let _ = writeln!(
        out,
        "    let rb = http.{verb}(crate::xrpc::xrpc_url(base, NSID));"
    );
    if matches!(output, BodyKind::Json(_)) {
        out.push_str("    let rb = rb.header(reqwest::header::ACCEPT, \"application/json\");\n");
    }
    if x.parameters.is_some() {
        out.push_str("    let rb = crate::xrpc::push_query(rb, params);\n");
    }
    match &input {
        BodyKind::Json(_) => out.push_str("    let rb = rb.json(input);\n"),
        BodyKind::Binary => out.push_str(
            "    let rb = rb.header(reqwest::header::CONTENT_TYPE, content_type).body(body);\n",
        ),
        BodyKind::None => {}
    }
    out.push_str("    rb\n}\n\n");
}