cimoxide-gen 0.3.0

use std::collections::{HashMap, HashSet};
use std::fmt::Write as FmtWrite;
use std::fs;
use std::path::Path;

use crate::schema::model::*;

pub fn generate_python_stubs(
    spec: &CimSpecification,
    output_dir: &Path,
) -> Result<(), Box<dyn std::error::Error>> {
    fs::create_dir_all(output_dir)?;

    let stubs = render_stubs(spec);
    fs::write(output_dir.join("types.pyi"), stubs)?;
    Ok(())
}

fn render_stubs(spec: &CimSpecification) -> String {
    let mut s = String::new();
    writeln!(s, "# Generated by cimgen — do not edit by hand.").unwrap();
    writeln!(s, "from __future__ import annotations").unwrap();
    writeln!(s, "from typing import TypedDict").unwrap();
    writeln!(s).unwrap();

    // Enums first as Literal aliases
    let mut enum_ids: Vec<&String> = spec.enums.keys().collect();
    enum_ids.sort();
    for id in &enum_ids {
        let e = &spec.enums[*id];
        if e.values.is_empty() {
            continue;
        }
        render_enum_stub(&mut s, e);
    }

    // Structs in topological order (parent before child)
    let sorted = topological_sort(&spec.types);
    for id in &sorted {
        if let Some(t) = spec.types.get(id) {
            render_type_stub(&mut s, t, spec);
        }
    }

    s
}

fn render_enum_stub(s: &mut String, e: &CimEnum) {
    let values: Vec<String> = e
        .values
        .iter()
        .map(|v| format!("\"{}\"", v.label.replace('"', "\\\"")))
        .collect();
    writeln!(s, "{} = {}", e.id, values.join(" | ")).unwrap();
    writeln!(s).unwrap();
}

fn render_type_stub(s: &mut String, t: &CimType, spec: &CimSpecification) {
    let base = if t.super_type.is_empty() || !spec.types.contains_key(&t.super_type) {
        "TypedDict".to_string()
    } else {
        t.super_type.clone()
    };

    if !t.comment.is_empty() {
        writeln!(s, "# {}", t.comment).unwrap();
    }
    writeln!(s, "class {}({}, total=False):", t.id, base).unwrap();

    let fields: Vec<&CimAttribute> = t
        .attributes
        .iter()
        .filter(|a| a.is_association_used)
        .collect();

    if t.super_type.is_empty() {
        writeln!(s, "    id: str").unwrap();
    }

    if fields.is_empty() && t.super_type.is_empty() {
        // id field already emitted above
    } else if fields.is_empty() {
        writeln!(s, "    pass").unwrap();
    } else {
        for attr in fields {
            let fname = to_snake_case(&attr.label);
            let ftype = py_field_type(attr);
            writeln!(s, "    {fname}: {ftype}").unwrap();
        }
    }

    writeln!(s).unwrap();
}

fn py_field_type(attr: &CimAttribute) -> &'static str {
    if attr.is_primitive || attr.is_cim_datatype {
        if attr.is_list {
            if attr.lang_type == "Vec<f64>" {
                "list[float]"
            } else if attr.lang_type == "Vec<i64>" {
                "list[int]"
            } else if attr.lang_type == "Vec<bool>" {
                "list[bool]"
            } else {
                "list[str]"
            }
        } else {
            match attr.lang_type.as_str() {
                "f64" => "float | None",
                "i64" => "int | None",
                "bool" => "bool | None",
                _ => "str",
            }
        }
    } else if attr.is_enum_value {
        "str | None"
    } else if attr.is_list {
        "list[str]"
    } else {
        "str | None"
    }
}

fn topological_sort(types: &HashMap<String, CimType>) -> Vec<String> {
    let mut sorted = Vec::new();
    let mut visited = HashSet::new();
    let mut ids: Vec<String> = types.keys().cloned().collect();
    ids.sort();
    for id in ids {
        topo_visit(&id, types, &mut visited, &mut sorted);
    }
    sorted
}

fn topo_visit(
    id: &str,
    types: &HashMap<String, CimType>,
    visited: &mut HashSet<String>,
    sorted: &mut Vec<String>,
) {
    if visited.contains(id) {
        return;
    }
    visited.insert(id.to_string());
    if let Some(t) = types.get(id) {
        if !t.super_type.is_empty() && types.contains_key(&t.super_type) {
            topo_visit(&t.super_type, types, visited, sorted);
        }
    }
    sorted.push(id.to_string());
}

fn to_snake_case(s: &str) -> String {
    let mut result = String::new();
    let chars: Vec<char> = s.chars().collect();
    for (i, &c) in chars.iter().enumerate() {
        if c.is_uppercase() {
            if i > 0 && !chars[i - 1].is_uppercase() {
                result.push('_');
            } else if i > 0
                && chars[i - 1].is_uppercase()
                && i + 1 < chars.len()
                && chars[i + 1].is_lowercase()
            {
                result.push('_');
            }
            result.push(c.to_lowercase().next().unwrap());
        } else {
            result.push(c);
        }
    }
    result
}