pr4xis 0.6.0

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::Entity
/// - Static arrays for all entities and relations
/// - TaxonomyDef, EquivalenceDef, OppositionDef, MereologyDef, CausalDef implementations
/// - Word lookup functions with cached adjacency maps
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, builder, &id_map);

    if let Some(name) = &config.taxonomy_name {
        write_relation_impl(
            &mut out,
            name,
            "TaxonomyDef",
            &config.entity_type_name,
            &builder.taxonomy,
            &id_map,
        );
    }
    if let Some(name) = &config.equivalence_name {
        write_relation_impl(
            &mut out,
            name,
            "EquivalenceDef",
            &config.entity_type_name,
            &builder.equivalence,
            &id_map,
        );
    }
    if let Some(name) = &config.opposition_name {
        write_relation_impl(
            &mut out,
            name,
            "OppositionDef",
            &config.entity_type_name,
            &builder.opposition,
            &id_map,
        );
    }
    if let Some(name) = &config.mereology_name {
        write_relation_impl(
            &mut out,
            name,
            "MereologyDef",
            &config.entity_type_name,
            &builder.mereology,
            &id_map,
        );
    }
    if let Some(name) = &config.causation_name {
        write_relation_impl(
            &mut out,
            name,
            "CausalDef",
            &config.entity_type_name,
            &builder.causation,
            &id_map,
        );
    }

    write_stats(&mut out, builder);
    write_codegen_data(&mut out, config, builder, &id_map);

    out
}

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::category::Entity;").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 Entity 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();

    // POS tags (if available)
    writeln!(out, "static ENTITY_POS: &[&str] = &[").unwrap();
    for entity in &builder.entities {
        let pos = entity.pos.as_deref().unwrap_or("");
        writeln!(out, "    \"{pos}\",").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 pos(&self) -> &'static str {{").unwrap();
    writeln!(
        out,
        "        ENTITY_POS.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, builder: &OntologyBuilder, id_map: &HashMap<&str, u32>) {
    if builder.word_index.is_empty() {
        return;
    }

    // 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);

    writeln!(out, "static WORD_INDEX: &[(&str, &[u32])] = &[").unwrap();
    for (word, ids) in &sorted_words {
        let word_escaped = word.replace('\\', "\\\\").replace('"', "\\\"");
        let id_list: Vec<String> = ids.iter().map(|id| id.to_string()).collect();
        writeln!(out, "    (\"{word_escaped}\", &[{}]),", id_list.join(", ")).unwrap();
    }
    writeln!(out, "];").unwrap();
    writeln!(out).unwrap();

    writeln!(out, "/// Look up synsets by word text (binary search).").unwrap();
    writeln!(out, "pub fn lookup(word: &str) -> &'static [u32] {{").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_relation_impl(
    out: &mut String,
    struct_name: &str,
    trait_name: &str,
    entity_type: &str,
    relations: &[(String, String)],
    id_map: &HashMap<&str, u32>,
) {
    // Static array of relation pairs
    let array_name = format!("{}_RELATIONS", struct_name.to_uppercase());
    writeln!(
        out,
        "static {array_name}: &[({entity_type}, {entity_type})] = &["
    )
    .unwrap();

    let mut count = 0;
    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, "    ({entity_type}({a_idx}), {entity_type}({b_idx})),").unwrap();
            count += 1;
        }
    }
    writeln!(out, "];").unwrap();
    writeln!(out).unwrap();

    // Struct + trait impl
    writeln!(out, "pub struct {struct_name};").unwrap();
    writeln!(out).unwrap();

    // Map trait name to the correct use path and method
    let (use_path, method_sig) = match trait_name {
        "TaxonomyDef" => (
            "pr4xis::ontology::reasoning::taxonomy::TaxonomyDef",
            format!(
                "    type Entity = {entity_type};\n    fn relations() -> Vec<({entity_type}, {entity_type})> {{ {array_name}.to_vec() }}"
            ),
        ),
        "EquivalenceDef" => (
            "pr4xis::ontology::reasoning::equivalence::EquivalenceDef",
            format!(
                "    type Entity = {entity_type};\n    fn pairs() -> Vec<({entity_type}, {entity_type})> {{ {array_name}.to_vec() }}"
            ),
        ),
        "OppositionDef" => (
            "pr4xis::ontology::reasoning::opposition::OppositionDef",
            format!(
                "    type Entity = {entity_type};\n    fn pairs() -> Vec<({entity_type}, {entity_type})> {{ {array_name}.to_vec() }}"
            ),
        ),
        "MereologyDef" => (
            "pr4xis::ontology::reasoning::mereology::MereologyDef",
            format!(
                "    type Entity = {entity_type};\n    fn relations() -> Vec<({entity_type}, {entity_type})> {{ {array_name}.to_vec() }}"
            ),
        ),
        "CausalDef" => (
            "pr4xis::ontology::reasoning::causation::CausalDef",
            format!(
                "    type Entity = {entity_type};\n    fn relations() -> Vec<({entity_type}, {entity_type})> {{ {array_name}.to_vec() }}"
            ),
        ),
        _ => return,
    };

    writeln!(out, "impl {use_path} for {struct_name} {{").unwrap();
    writeln!(out, "{method_sig}").unwrap();
    writeln!(out, "}}").unwrap();
    writeln!(out).unwrap();

    eprintln!("  codegen: {struct_name} ({trait_name}) — {count} relations");
}

fn write_codegen_data(
    out: &mut String,
    config: &GenerateConfig,
    builder: &OntologyBuilder,
    id_map: &HashMap<&str, u32>,
) {
    // Write raw (u32, u32) relation arrays for CodegenData
    let write_raw_relations = |out: &mut String,
                               name: &str,
                               relations: &[(String, String)],
                               id_map: &HashMap<&str, u32>| {
        writeln!(out, "static {name}: &[(u32, u32)] = &[").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, "    ({a_idx}, {b_idx}),").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);

    writeln!(out).unwrap();
    writeln!(
        out,
        "/// Language-agnostic codegen data — consumed by Language functor."
    )
    .unwrap();
    writeln!(
        out,
        "pub static CODEGEN_DATA: pr4xis::codegen_data::CodegenData = pr4xis::codegen_data::CodegenData {{"
    )
    .unwrap();
    writeln!(out, "    entity_count: {},", builder.entities.len()).unwrap();
    writeln!(out, "    entity_ids: ENTITY_IDS,").unwrap();
    writeln!(out, "    entity_pos: ENTITY_POS,").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, "}};").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 WORD_INDEX_COUNT: usize = {};",
        builder.word_index.len()
    )
    .unwrap();
    writeln!(out, "}}").unwrap();
}