cimoxide-gen 0.4.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();
}

/// One class as a functional `TypedDict`: its keys are the attributes'
/// `Class.attr` ids, which are not Python identifiers, and a functional
/// `TypedDict` cannot extend another, so inherited attributes are listed too.
/// Values are strings as written; a list-valued attribute is a list of them.
fn render_type_stub(s: &mut String, t: &CimType, spec: &CimSpecification) {
    if !t.comment.is_empty() {
        writeln!(s, "# {}", t.comment.replace('\n', " ")).unwrap();
    }
    let mut keys: Vec<(String, &'static str)> = vec![
        ("id".to_string(), "str"),
        ("_type".to_string(), "str"),
    ];
    let mut chain = vec![t];
    while let Some(parent) = spec.types.get(&chain.last().unwrap().super_type) {
        chain.push(parent);
    }
    for class in chain.iter().rev() {
        for attr in class.attributes.iter().filter(|a| a.is_association_used) {
            let ty = if attr.is_list { "list[str]" } else { "str" };
            keys.push((attr.id.clone(), ty));
        }
    }
    let fields: Vec<String> = keys.iter().map(|(k, ty)| format!("\"{k}\": {ty}")).collect();
    writeln!(s, "{id} = TypedDict(\"{id}\", {{{}}}, total=False)", fields.join(", "), id = t.id).unwrap();
    writeln!(s).unwrap();
}

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());
}