pr4xis 0.25.2

use std::collections::HashMap;
use std::fmt::Write;

use super::builder::{GenerateConfig, OntologyBuilder};

/// Generate Rust source code from ontology data.
///
/// Produces a module containing:
/// - Entity type (newtype over u32) implementing `pr4xis::category::Concept`
/// - Static `ENTITY_LABELS` + raw relation arrays (`RAW_TAXONOMY`,
///   `RAW_EQUIVALENCE`, `RAW_OPPOSITION`, `RAW_MEREOLOGY`, `RAW_CAUSATION`,
///   `RAW_REFERENCES`) as `&[(EntityRef<Marker>, EntityRef<Marker>)]`
/// - Word-lookup functions with cached adjacency maps
///
/// Per-def trait impls (`TaxonomyDef`, `EquivalenceDef`, `OppositionDef`,
/// `MereologyDef`, `CausalDef`) were deleted in #168; consumers read the
/// `RAW_*` arrays and feed them into kinded morphisms.
pub fn generate_rust(builder: &OntologyBuilder, config: &GenerateConfig) -> String {
    let mut out = String::new();

    // Build ID → index mapping for compact representation
    let id_map: HashMap<&str, u32> = builder
        .entities
        .iter()
        .enumerate()
        .map(|(i, e)| (e.id.as_str(), i as u32))
        .collect();

    write_header(&mut out, config, builder);
    write_entity_type(&mut out, config, builder, &id_map);
    write_entity_data(&mut out, config, builder, &id_map);
    write_word_index(&mut out, config, builder, &id_map);
    write_stats(&mut out, builder);
    write_codegen_data(&mut out, config, builder, &id_map);

    out
}

/// Leaf type name to use inside the generated arrays. For an external
/// marker like `"pr4xis_domains::cognitive::linguistics::english::English"`,
/// this is `"English"`. For a bare-ident marker (statute case, marker
/// defined inline) it's the path itself.
fn marker_leaf(config: &GenerateConfig) -> &str {
    let path = config.entity_marker_path.as_str();
    match path.rsplit_once("::") {
        Some((_, leaf)) => leaf,
        None => path,
    }
}

/// Whether to emit a `use <path>;` for the marker. Single-ident markers
/// resolve through the generated module's own scope (the statute case);
/// fully-qualified paths need an explicit import.
fn marker_needs_use(config: &GenerateConfig) -> bool {
    config.entity_marker_path.contains("::")
}

fn write_header(out: &mut String, config: &GenerateConfig, builder: &OntologyBuilder) {
    writeln!(out, "// Auto-generated by pr4xis::codegen").unwrap();
    writeln!(out, "// Module: {}", config.module_name).unwrap();
    writeln!(out, "// Entities: {}", builder.entities.len()).unwrap();
    writeln!(out, "// Relations: {}", builder.relation_count()).unwrap();
    writeln!(out, "// DO NOT EDIT — regenerate from source data").unwrap();
    writeln!(out).unwrap();
    writeln!(out, "use pr4xis::EntityRef;").unwrap();
    writeln!(out, "use pr4xis::category::{{Concept, FinitelyGenerated}};").unwrap();
    if marker_needs_use(config) {
        writeln!(out, "use {};", config.entity_marker_path).unwrap();
    }
    writeln!(out).unwrap();
}

fn write_entity_type(
    out: &mut String,
    config: &GenerateConfig,
    builder: &OntologyBuilder,
    _id_map: &HashMap<&str, u32>,
) {
    let ty = &config.entity_type_name;
    let count = builder.entities.len();

    writeln!(out, "#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]").unwrap();
    writeln!(out, "pub struct {ty}(pub u32);").unwrap();
    writeln!(out).unwrap();
    writeln!(out, "impl Concept for {ty} {{}}").unwrap();
    writeln!(out).unwrap();
    writeln!(out, "impl FinitelyGenerated for {ty} {{").unwrap();
    writeln!(out, "    fn variants() -> Vec<Self> {{").unwrap();
    writeln!(out, "        (0..{count}u32).map({ty}).collect()").unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();
}

fn write_entity_data(
    out: &mut String,
    config: &GenerateConfig,
    builder: &OntologyBuilder,
    _id_map: &HashMap<&str, u32>,
) {
    let ty = &config.entity_type_name;

    // Entity labels (for display/debug)
    writeln!(out, "static ENTITY_LABELS: &[&str] = &[").unwrap();
    for entity in &builder.entities {
        let label = entity.label.replace('"', "\\\"");
        writeln!(out, "    \"{label}\",").unwrap();
    }
    writeln!(out, "];").unwrap();
    writeln!(out).unwrap();

    // Entity IDs (original string IDs for lookup)
    writeln!(out, "static ENTITY_IDS: &[&str] = &[").unwrap();
    for entity in &builder.entities {
        let id = entity.id.replace('"', "\\\"");
        writeln!(out, "    \"{id}\",").unwrap();
    }
    writeln!(out, "];").unwrap();
    writeln!(out).unwrap();

    // Entity-kind tags. POS for lexical entities (`"n"`, `"v"`, …),
    // synthetic kinds for statute terms (`"statute_term"`), USLM
    // element names for the forthcoming UsCode codegen.
    writeln!(out, "static ENTITY_KIND: &[&str] = &[").unwrap();
    for entity in &builder.entities {
        let kind = entity.pos.as_deref().unwrap_or("");
        writeln!(out, "    \"{kind}\",").unwrap();
    }
    writeln!(out, "];").unwrap();
    writeln!(out).unwrap();

    // Definitions (first definition per entity, empty string if none)
    writeln!(out, "static ENTITY_DEFS: &[&str] = &[").unwrap();
    for entity in &builder.entities {
        let def = entity
            .definitions
            .first()
            .map(|d| d.replace('\\', "\\\\").replace('"', "\\\""))
            .unwrap_or_default();
        writeln!(out, "    \"{def}\",").unwrap();
    }
    writeln!(out, "];").unwrap();
    writeln!(out).unwrap();

    // Lookup functions
    writeln!(out, "impl {ty} {{").unwrap();
    writeln!(out, "    pub fn label(&self) -> &'static str {{").unwrap();
    writeln!(
        out,
        "        ENTITY_LABELS.get(self.0 as usize).copied().unwrap_or(\"unknown\")"
    )
    .unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out).unwrap();
    writeln!(out, "    pub fn original_id(&self) -> &'static str {{").unwrap();
    writeln!(
        out,
        "        ENTITY_IDS.get(self.0 as usize).copied().unwrap_or(\"unknown\")"
    )
    .unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out).unwrap();
    writeln!(out, "    pub fn kind(&self) -> &'static str {{").unwrap();
    writeln!(
        out,
        "        ENTITY_KIND.get(self.0 as usize).copied().unwrap_or(\"\")"
    )
    .unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out).unwrap();
    writeln!(out, "    pub fn definition(&self) -> &'static str {{").unwrap();
    writeln!(
        out,
        "        ENTITY_DEFS.get(self.0 as usize).copied().unwrap_or(\"\")"
    )
    .unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();
}

fn write_word_index(
    out: &mut String,
    config: &GenerateConfig,
    builder: &OntologyBuilder,
    id_map: &HashMap<&str, u32>,
) {
    // Always emit WORD_INDEX (possibly empty) — `CODEGEN_DATA` references
    // it unconditionally. A source whose terms carry no lemmas (e.g.,
    // statute structural extractions before adjunction-to-English
    // codegen runs) still needs the symbol to resolve.

    let marker = marker_leaf(config);

    // Group words by text for multi-sense lookup.
    let mut by_word: HashMap<&str, Vec<u32>> = HashMap::new();
    for (word, entity_id) in &builder.word_index {
        if let Some(&idx) = id_map.get(entity_id.as_str()) {
            by_word.entry(word.as_str()).or_default().push(idx);
        }
    }

    // Sort for binary search.
    let mut sorted_words: Vec<(&str, &Vec<u32>)> = by_word.iter().map(|(&k, v)| (k, v)).collect();
    sorted_words.sort_by_key(|(w, _)| *w);

    // Per-row id slice statics so each WORD_INDEX entry references a
    // `&'static [EntityRef<Marker>]` (the same type the runtime
    // `CodegenData<P>` expects). Inline array literals can't be coerced
    // to a slice inside a `&'static [...]` tuple element.
    for (i, (_, ids)) in sorted_words.iter().enumerate() {
        write!(
            out,
            "static WORD_INDEX_IDS_{i}: &[EntityRef<{marker}>] = &["
        )
        .unwrap();
        for (j, id) in ids.iter().enumerate() {
            if j > 0 {
                write!(out, ", ").unwrap();
            }
            write!(out, "EntityRef::<{marker}>::new({id}u64)").unwrap();
        }
        writeln!(out, "];").unwrap();
    }
    if !sorted_words.is_empty() {
        writeln!(out).unwrap();
    }

    writeln!(
        out,
        "static WORD_INDEX: &[(&str, &[EntityRef<{marker}>])] = &["
    )
    .unwrap();
    for (i, (word, _)) in sorted_words.iter().enumerate() {
        let word_escaped = word.replace('\\', "\\\\").replace('"', "\\\"");
        writeln!(out, "    (\"{word_escaped}\", WORD_INDEX_IDS_{i}),").unwrap();
    }
    writeln!(out, "];").unwrap();
    writeln!(out).unwrap();

    writeln!(
        out,
        "/// Look up entity handles by word text (binary search)."
    )
    .unwrap();
    writeln!(
        out,
        "pub fn lookup(word: &str) -> &'static [EntityRef<{marker}>] {{"
    )
    .unwrap();
    writeln!(
        out,
        "    match WORD_INDEX.binary_search_by_key(&word, |(w, _)| w) {{"
    )
    .unwrap();
    writeln!(out, "        Ok(idx) => WORD_INDEX[idx].1,").unwrap();
    writeln!(out, "        Err(_) => &[],").unwrap();
    writeln!(out, "    }}").unwrap();
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();
}

fn write_codegen_data(
    out: &mut String,
    config: &GenerateConfig,
    builder: &OntologyBuilder,
    id_map: &HashMap<&str, u32>,
) {
    let marker = marker_leaf(config);

    // Write typed `(EntityRef<Marker>, EntityRef<Marker>)` relation
    // arrays for CodegenData<Marker>.
    let write_raw_relations = |out: &mut String,
                               name: &str,
                               relations: &[(String, String)],
                               id_map: &HashMap<&str, u32>| {
        writeln!(
            out,
            "static {name}: &[(EntityRef<{marker}>, EntityRef<{marker}>)] = &["
        )
        .unwrap();
        for (a, b) in relations {
            if let (Some(&a_idx), Some(&b_idx)) = (id_map.get(a.as_str()), id_map.get(b.as_str())) {
                writeln!(
                    out,
                    "    (EntityRef::<{marker}>::new({a_idx}u64), EntityRef::<{marker}>::new({b_idx}u64)),"
                )
                .unwrap();
            }
        }
        writeln!(out, "];").unwrap();
        writeln!(out).unwrap();
    };

    write_raw_relations(out, "RAW_TAXONOMY", &builder.taxonomy, id_map);
    write_raw_relations(out, "RAW_MEREOLOGY", &builder.mereology, id_map);
    write_raw_relations(out, "RAW_OPPOSITION", &builder.opposition, id_map);
    write_raw_relations(out, "RAW_EQUIVALENCE", &builder.equivalence, id_map);
    write_raw_relations(out, "RAW_CAUSATION", &builder.causation, id_map);
    write_raw_relations(out, "RAW_REFERENCES", &builder.references, id_map);

    writeln!(out).unwrap();
    writeln!(
        out,
        "/// Ontology-agnostic codegen data — consumed by the matching `from_codegen` functor."
    )
    .unwrap();
    writeln!(
        out,
        "pub static CODEGEN_DATA: pr4xis::codegen_data::CodegenData<{marker}> = pr4xis::codegen_data::CodegenData {{"
    )
    .unwrap();
    writeln!(out, "    entity_count: {},", builder.entities.len()).unwrap();
    writeln!(out, "    entity_ids: ENTITY_IDS,").unwrap();
    writeln!(out, "    entity_kind: ENTITY_KIND,").unwrap();
    writeln!(out, "    entity_labels: ENTITY_LABELS,").unwrap();
    writeln!(out, "    entity_defs: ENTITY_DEFS,").unwrap();
    writeln!(out, "    word_index: WORD_INDEX,").unwrap();
    writeln!(out, "    taxonomy: RAW_TAXONOMY,").unwrap();
    writeln!(out, "    mereology: RAW_MEREOLOGY,").unwrap();
    writeln!(out, "    opposition: RAW_OPPOSITION,").unwrap();
    writeln!(out, "    equivalence: RAW_EQUIVALENCE,").unwrap();
    writeln!(out, "    causation: RAW_CAUSATION,").unwrap();
    writeln!(out, "    references: RAW_REFERENCES,").unwrap();
    writeln!(out, "}};").unwrap();
    writeln!(out).unwrap();

    // Suppress unused warnings for the config entity type
    let _ = config;
}

fn write_stats(out: &mut String, builder: &OntologyBuilder) {
    writeln!(out, "/// Generated ontology statistics.").unwrap();
    writeln!(out, "pub mod stats {{").unwrap();
    writeln!(
        out,
        "    pub const ENTITY_COUNT: usize = {};",
        builder.entities.len()
    )
    .unwrap();
    writeln!(
        out,
        "    pub const TAXONOMY_COUNT: usize = {};",
        builder.taxonomy.len()
    )
    .unwrap();
    writeln!(
        out,
        "    pub const EQUIVALENCE_COUNT: usize = {};",
        builder.equivalence.len()
    )
    .unwrap();
    writeln!(
        out,
        "    pub const OPPOSITION_COUNT: usize = {};",
        builder.opposition.len()
    )
    .unwrap();
    writeln!(
        out,
        "    pub const MEREOLOGY_COUNT: usize = {};",
        builder.mereology.len()
    )
    .unwrap();
    writeln!(
        out,
        "    pub const CAUSATION_COUNT: usize = {};",
        builder.causation.len()
    )
    .unwrap();
    writeln!(
        out,
        "    pub const REFERENCES_COUNT: usize = {};",
        builder.references.len()
    )
    .unwrap();
    writeln!(
        out,
        "    pub const WORD_INDEX_COUNT: usize = {};",
        builder.word_index.len()
    )
    .unwrap();
    writeln!(out, "}}").unwrap();
}