//! Build script that generates Python bindings from Rust function metadata.
//!
//! This runs during `cargo build` and generates:
//! - r4pm/bindings/__init__.py (main bindings module)
//! - r4pm/bindings/__init__.pyi (type stubs)
//! - r4pm/bindings/<module>/*.py (per-module wrappers)
//! - r4pm/bindings/<module>/*.pyi (per-module stubs)

use std::collections::{BTreeMap, BTreeSet};
use std::fs;
use std::path::Path;

use process_mining::bindings::BindingMeta;

fn main() {
    // Tell Cargo to rerun if these change
    println!("cargo:rerun-if-changed=build.rs");
    
    let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap();
    let bindings_dir = Path::new(&manifest_dir).join("r4pm").join("bindings");
    
    // Clean and recreate bindings directory
    if bindings_dir.exists() {
        fs::remove_dir_all(&bindings_dir).expect("Failed to remove old bindings");
    }
    fs::create_dir_all(&bindings_dir).expect("Failed to create bindings directory");
    
    // Get function metadata from process_mining
    let functions = process_mining::bindings::list_functions_meta();
    
    // Group functions by module
    let mut module_groups: BTreeMap<String, Vec<BindingMeta>> = BTreeMap::new();
    for func in functions {
        let py_module = rust_module_to_python(&func.module);
        module_groups.entry(py_module).or_default().push(func);
    }
    
    // println!("cargo:warning=Generating bindings for {} modules", module_groups.len());
    
    // Generate module files
    for (module_path, funcs) in &module_groups {
        generate_module_files(&bindings_dir, module_path, funcs);
    }
    
    // Generate intermediate __init__.py files
    generate_submodule_inits(&bindings_dir, &module_groups);
    
    // Generate main bindings/__init__.py
    generate_bindings_init(&bindings_dir, &module_groups);
}

/// Convert Rust module path to Python module path
fn rust_module_to_python(rust_module: &str) -> String {
    rust_module
        .replace("process_mining::", "")
        .replace("::", ".")
}

/// Convert JSON schema to Python type hint
fn schema_to_python_type(schema: &serde_json::Value) -> String {
    if schema.get("x-registry-ref").is_some() {
        return "str".to_string();
    }
    
    match schema.get("type").and_then(|t| t.as_str()) {
        Some("string") => "str".to_string(),
        Some("integer") => "int".to_string(),
        Some("number") => "float".to_string(),
        Some("boolean") => "bool".to_string(),
        Some("null") => "None".to_string(),
        Some("array") => {
            let items_type = schema
                .get("items")
                .map(|i| schema_to_python_type(i))
                .unwrap_or_else(|| "Any".to_string());
            format!("List[{}]", items_type)
        }
        Some("object") => "Dict[str, Any]".to_string(),
        _ => "Any".to_string(),
    }
}

/// Sanitize function name for Python
fn sanitize_name(name: &str) -> String {
    name.replace('-', "_").replace('+', "plus")
}

/// Generate .py and .pyi files for a module
fn generate_module_files(
    bindings_dir: &Path,
    module_path: &str,
    functions: &[BindingMeta],
) {
    let parts: Vec<&str> = module_path.split('.').collect();
    let import_depth = parts.len() + 1;
    
    // Create directory structure
    let module_dir = if parts.len() > 1 {
        bindings_dir.join(parts[..parts.len() - 1].join("/"))
    } else {
        bindings_dir.to_path_buf()
    };
    fs::create_dir_all(&module_dir).expect("Failed to create module directory");
    
    let file_name = parts.last().unwrap();
    
    // Generate .py file
    let mut py_content = format!(
        r#""""
Auto-generated bindings for {module_path}

DO NOT EDIT - Generated by build.rs
"""
from typing import Any, Dict, List, Optional

"#
    );
    
    let mut func_names = Vec::new();
    for func in functions {
        let py_name = sanitize_name(&func.name);
        func_names.push(py_name.clone());
        py_content.push_str(&generate_function_wrapper(func, import_depth));
        py_content.push_str("\n\n");
    }
    
    py_content.push_str("__all__ = [\n");
    for name in &func_names {
        py_content.push_str(&format!("    \"{}\",\n", name));
    }
    py_content.push_str("]\n");
    
    fs::write(module_dir.join(format!("{}.py", file_name)), py_content)
        .expect("Failed to write .py file");
    
    // Generate .pyi file
    let mut pyi_content = format!(
        r#""""
Type stubs for {module_path}

DO NOT EDIT - Generated by build.rs
"""
from typing import Any, Dict, List, Optional

"#
    );
    
    for func in functions {
        pyi_content.push_str(&generate_function_stub(func));
        pyi_content.push_str("\n\n");
    }
    
    pyi_content.push_str("__all__: List[str]\n");
    
    fs::write(module_dir.join(format!("{}.pyi", file_name)), pyi_content)
        .expect("Failed to write .pyi file");
}

/// Generate a Python wrapper function
fn generate_function_wrapper(func: &BindingMeta, import_depth: usize) -> String {
    let py_name = sanitize_name(&func.name);
    let required_args: BTreeSet<&str> = func.required_args.iter().map(|s| s.as_str()).collect();
    
    // Build parameters
    let mut params = Vec::new();
    let mut doc_args = Vec::new();
    let mut dict_lines = vec!["    args_dict = {}".to_string()];
    
    for (arg_name, arg_schema) in &func.args {
        let arg_type = schema_to_python_type(arg_schema);
        let is_required = required_args.contains(arg_name.as_str());
        
        if is_required {
            params.push(format!("{}: {}", arg_name, arg_type));
            dict_lines.push(format!("    args_dict[\"{}\"] = {}", arg_name, arg_name));
        } else {
            params.push(format!("{}: Optional[{}] = None", arg_name, arg_type));
            dict_lines.push(format!("    if {} is not None:", arg_name));
            dict_lines.push(format!("        args_dict[\"{}\"] = {}", arg_name, arg_name));
        }
        
        // Doc line
        if arg_schema.get("x-registry-ref").is_some() {
            let title = arg_schema.get("title").and_then(|t| t.as_str()).unwrap_or("item");
            doc_args.push(format!("        {}: Registry ID for {}", arg_name, title));
        } else if let Some(desc) = arg_schema.get("description").and_then(|d| d.as_str()) {
            doc_args.push(format!("        {}: {}", arg_name, desc.lines().next().unwrap_or("")));
        } else {
            doc_args.push(format!("        {}: {}", arg_name, arg_type));
        }
    }
    
    // Return type
    let return_type = schema_to_python_type(&func.return_type);
    
    // Build docstring
    let mut doc_lines: Vec<String> = func.docs.iter().map(|s| format!("    {}", s)).collect();
    if !doc_args.is_empty() {
        doc_lines.push("".to_string());
        doc_lines.push("    Args:".to_string());
        doc_lines.extend(doc_args);
    }
    doc_lines.push("".to_string());
    doc_lines.push("    Returns:".to_string());
    if func.return_type.get("x-registry-ref").is_some() {
        let title = func.return_type.get("title").and_then(|t| t.as_str()).unwrap_or("result");
        doc_lines.push(format!("        Registry ID for {}", title));
    } else {
        doc_lines.push(format!("        {}", return_type));
    }
    
    let import_prefix = ".".repeat(import_depth);
    
    format!(
        r#"def {py_name}({params}) -> {return_type}:
    """
{docstring}
    """
    import json
    from {import_prefix}r4pm import call_binding
{dict_lines}
    result = call_binding("{func_id}", json.dumps(args_dict))
    return json.loads(result)"#,
        py_name = py_name,
        params = params.join(", "),
        return_type = return_type,
        docstring = doc_lines.join("\n"),
        import_prefix = import_prefix,
        dict_lines = dict_lines.join("\n"),
        func_id = func.id,
    )
}

/// Generate a stub function for .pyi
fn generate_function_stub(func: &BindingMeta) -> String {
    let py_name = sanitize_name(&func.name);
    let required_args: BTreeSet<&str> = func.required_args.iter().map(|s| s.as_str()).collect();
    
    let mut params = Vec::new();
    let mut doc_args = Vec::new();
    
    for (arg_name, arg_schema) in &func.args {
        let arg_type = schema_to_python_type(arg_schema);
        let is_required = required_args.contains(arg_name.as_str());
        
        if is_required {
            params.push(format!("{}: {}", arg_name, arg_type));
        } else {
            params.push(format!("{}: Optional[{}] = None", arg_name, arg_type));
        }
        
        if arg_schema.get("x-registry-ref").is_some() {
            let title = arg_schema.get("title").and_then(|t| t.as_str()).unwrap_or("item");
            doc_args.push(format!("        {}: Registry ID for {}", arg_name, title));
        } else if let Some(desc) = arg_schema.get("description").and_then(|d| d.as_str()) {
            doc_args.push(format!("        {}: {}", arg_name, desc.lines().next().unwrap_or("")));
        } else {
            doc_args.push(format!("        {}: {}", arg_name, arg_type));
        }
    }
    
    let return_type = schema_to_python_type(&func.return_type);
    
    // Docstring
    let mut doc_lines: Vec<String> = func.docs.iter().map(|s| format!("    {}", s)).collect();
    if !doc_args.is_empty() {
        doc_lines.push("".to_string());
        doc_lines.push("    Args:".to_string());
        doc_lines.extend(doc_args);
    }
    doc_lines.push("".to_string());
    doc_lines.push("    Returns:".to_string());
    doc_lines.push(format!("        {}", return_type));
    
    format!(
        r#"def {py_name}({params}) -> {return_type}:
    """
{docstring}
    """
    ..."#,
        py_name = py_name,
        params = params.join(", "),
        return_type = return_type,
        docstring = doc_lines.join("\n"),
    )
}

/// Generate __init__.py files for intermediate directories
fn generate_submodule_inits(
    bindings_dir: &Path,
    module_groups: &BTreeMap<String, Vec<BindingMeta>>,
) {
    // Collect directory structure
    let mut dir_contents: BTreeMap<String, (BTreeSet<String>, BTreeSet<String>)> = BTreeMap::new();
    
    for module_path in module_groups.keys() {
        let parts: Vec<&str> = module_path.split('.').collect();
        
        // Track file at leaf level
        if parts.len() == 1 {
            dir_contents.entry(String::new()).or_default().1.insert(parts[0].to_string());
        } else {
            let parent = parts[..parts.len() - 1].join("/");
            dir_contents.entry(parent).or_default().1.insert(parts.last().unwrap().to_string());
        }
        
        // Track intermediate directories
        for i in 1..parts.len() {
            let parent = if i > 1 { parts[..i - 1].join("/") } else { String::new() };
            dir_contents.entry(parent).or_default().0.insert(parts[i - 1].to_string());
        }
    }
    
    // Generate __init__.py for each non-root directory
    for (dir_path, (submodules, files)) in &dir_contents {
        if dir_path.is_empty() {
            continue; // Root handled separately
        }
        
        let target_dir = bindings_dir.join(dir_path);
        fs::create_dir_all(&target_dir).expect("Failed to create submodule directory");
        
        let mut init_content = "\"\"\"Auto-generated module.\"\"\"\n\n".to_string();
        let mut all_items = Vec::new();
        
        for submod in submodules {
            init_content.push_str(&format!("from . import {}\n", submod));
            all_items.push(submod.clone());
        }
        
        for file_mod in files {
            init_content.push_str(&format!("from .{} import *\n", file_mod));
        }
        
        init_content.push_str(&format!("\n__all__ = {:?}\n", all_items));
        
        fs::write(target_dir.join("__init__.py"), &init_content)
            .expect("Failed to write submodule __init__.py");
        
        // Stub file
        let mut stub_content = "\"\"\"Type stubs.\"\"\"\nfrom typing import List\n\n".to_string();
        for submod in submodules {
            stub_content.push_str(&format!("from . import {} as {}\n", submod, submod));
        }
        for file_mod in files {
            stub_content.push_str(&format!("from .{} import *\n", file_mod));
        }
        stub_content.push_str("\n__all__: List[str]\n");
        
        fs::write(target_dir.join("__init__.pyi"), stub_content)
            .expect("Failed to write submodule __init__.pyi");
    }
}

/// Generate the main bindings/__init__.py
fn generate_bindings_init(
    bindings_dir: &Path,
    module_groups: &BTreeMap<String, Vec<BindingMeta>>,
) {
    let mut top_modules: BTreeSet<String> = BTreeSet::new();
    let mut all_funcs: Vec<(String, String)> = Vec::new();
    
    for (module_path, functions) in module_groups {
        let parts: Vec<&str> = module_path.split('.').collect();
        top_modules.insert(parts[0].to_string());
        
        for func in functions {
            let py_name = sanitize_name(&func.name);
            all_funcs.push((py_name, module_path.clone()));
        }
    }
    
    // Sort for deterministic output
    all_funcs.sort();
    
    let mut content = r#""""
Auto-generated bindings organized by module.

DO NOT EDIT - Generated by build.rs

Import functions directly or use submodules:
    from r4pm.bindings import discover_dfg
    from r4pm.bindings.discovery.case_centric import dfg
"""
from typing import Any, Dict, List, Optional

"#.to_string();
    
    // Import submodules
    for module in &top_modules {
        content.push_str(&format!("from . import {}\n", module));
    }
    
    content.push_str("\n# Re-export all functions at top level\n");
    for (func_name, module_path) in &all_funcs {
        let parts: Vec<&str> = module_path.split('.').collect();
        content.push_str(&format!("from .{} import {}\n", parts.join("."), func_name));
    }
    
    content.push_str("\n__all__ = [\n");
    for module in &top_modules {
        content.push_str(&format!("    \"{}\",\n", module));
    }
    for (func_name, _) in &all_funcs {
        content.push_str(&format!("    \"{}\",\n", func_name));
    }
    content.push_str("]\n");
    
    fs::write(bindings_dir.join("__init__.py"), &content)
        .expect("Failed to write bindings __init__.py");
    
    // Generate stub file
    let mut stub_content = r#""""
Type stubs for bindings module.

DO NOT EDIT - Generated by build.rs
"""
from typing import Any, Dict, List, Optional

"#.to_string();
    
    for module in &top_modules {
        stub_content.push_str(&format!("from . import {} as {}\n", module, module));
    }
    stub_content.push_str("\n");
    
    for functions in module_groups.values() {
        for func in functions {
            stub_content.push_str(&generate_function_stub(func));
            stub_content.push_str("\n\n");
        }
    }
    
    stub_content.push_str("__all__: List[str]\n");
    
    fs::write(bindings_dir.join("__init__.pyi"), stub_content)
        .expect("Failed to write bindings __init__.pyi");
}

