//! Python type definitions output backend.
//!
//! Generates Python dataclasses with type hints (3.10+).

use crate::ir::{EnumKind, Field, Schema, Type, TypeDef, TypeDefKind};
use crate::traits::{Backend, BackendCategory};

/// Options for Python code generation.
#[derive(Debug, Clone, Default)]
pub struct PythonOptions {
    /// Style: "dataclass" (default), "typeddict", or "namedtuple".
    pub style: PythonStyle,
    /// Whether to use `from __future__ import annotations` for forward refs.
    pub future_annotations: bool,
    /// Whether to make dataclass fields frozen (immutable).
    pub frozen: bool,
}

#[derive(Debug, Clone, Default, PartialEq)]
pub enum PythonStyle {
    #[default]
    Dataclass,
    TypedDict,
}

/// Generate Python type definitions from an IR schema.
pub fn generate_python_types(schema: &Schema, options: &PythonOptions) -> String {
    let mut out = String::new();
    out.push_str("# Auto-generated by normalize-typegen\n");

    // Imports
    if options.future_annotations {
        out.push_str("from __future__ import annotations\n");
    }
    out.push('\n');

    match options.style {
        PythonStyle::Dataclass => {
            out.push_str("from dataclasses import dataclass\n");
        }
        PythonStyle::TypedDict => {
            out.push_str("from typing import TypedDict\n");
        }
    }

    // Collect needed typing imports
    let mut typing_imports = Vec::new();
    for def in &schema.definitions {
        collect_typing_imports(&mut typing_imports, def);
    }
    typing_imports.sort();
    typing_imports.dedup();

    if !typing_imports.is_empty() {
        if options.style == PythonStyle::TypedDict {
            // TypedDict already imported
            typing_imports.retain(|i| *i != "TypedDict");
        }
        if !typing_imports.is_empty() {
            out.push_str("from typing import ");
            out.push_str(&typing_imports.join(", "));
            out.push('\n');
        }
    }

    // Check if we need Literal for enums
    let has_enums = schema
        .definitions
        .iter()
        .any(|d| matches!(d.kind, TypeDefKind::Enum(_)));
    if has_enums && !typing_imports.contains(&"Literal") {
        out.push_str("from typing import Literal\n");
    }

    out.push('\n');

    for (i, def) in schema.definitions.iter().enumerate() {
        if i > 0 {
            out.push('\n');
        }
        generate_typedef(&mut out, def, options);
    }

    out
}

fn collect_typing_imports(imports: &mut Vec<&'static str>, def: &TypeDef) {
    match &def.kind {
        TypeDefKind::Struct(s) => {
            for field in &s.fields {
                collect_type_imports(imports, &field.ty);
                if !field.required {
                    imports.push("Optional");
                }
            }
        }
        TypeDefKind::Enum(e) => {
            if let EnumKind::Tagged(t) = &e.kind {
                imports.push("Union");
                for variant in &t.variants {
                    for field in &variant.fields {
                        collect_type_imports(imports, &field.ty);
                        if !field.required {
                            imports.push("Optional");
                        }
                    }
                }
            }
        }
        TypeDefKind::Alias(ty) => {
            collect_type_imports(imports, ty);
        }
    }
}

fn collect_type_imports(imports: &mut Vec<&'static str>, ty: &Type) {
    match ty {
        Type::Array(_) => imports.push("list"),
        Type::Map { .. } => imports.push("dict"),
        Type::Optional(inner) => {
            imports.push("Optional");
            collect_type_imports(imports, inner);
        }
        Type::Union(types) => {
            imports.push("Union");
            for t in types {
                collect_type_imports(imports, t);
            }
        }
        Type::Any => imports.push("Any"),
        _ => {}
    }
}

fn generate_typedef(out: &mut String, def: &TypeDef, options: &PythonOptions) {
    // Doc comment
    if let Some(docs) = &def.docs {
        out.push_str("\"\"\"");
        out.push_str(docs);
        out.push_str("\"\"\"\n");
    }

    match &def.kind {
        TypeDefKind::Struct(s) => match options.style {
            PythonStyle::Dataclass => {
                if options.frozen {
                    out.push_str("@dataclass(frozen=True)\n");
                } else {
                    out.push_str("@dataclass\n");
                }
                out.push_str("class ");
                out.push_str(&def.name);
                out.push_str(":\n");

                if s.fields.is_empty() {
                    out.push_str("    pass\n");
                } else {
                    // Required fields first, then optional
                    let (required, optional): (Vec<_>, Vec<_>) =
                        s.fields.iter().partition(|f| f.required);

                    for field in required.iter().chain(optional.iter()) {
                        generate_dataclass_field(out, field);
                    }
                }
            }
            PythonStyle::TypedDict => {
                // Check if all fields are required
                let all_required = s.fields.iter().all(|f| f.required);

                out.push_str("class ");
                out.push_str(&def.name);
                if all_required {
                    out.push_str("(TypedDict):\n");
                } else {
                    out.push_str("(TypedDict, total=False):\n");
                }

                if s.fields.is_empty() {
                    out.push_str("    pass\n");
                } else {
                    for field in &s.fields {
                        generate_typeddict_field(out, field);
                    }
                }
            }
        },
        TypeDefKind::Enum(e) => match &e.kind {
            EnumKind::StringLiteral(variants) => {
                out.push_str(&def.name);
                out.push_str(" = Literal[\n");
                for (i, variant) in variants.iter().enumerate() {
                    if i > 0 {
                        out.push_str(",\n");
                    }
                    out.push_str("    \"");
                    out.push_str(&variant.value);
                    out.push('"');
                }
                out.push_str("\n]\n");
            }
            EnumKind::IntLiteral(variants) => {
                out.push_str(&def.name);
                out.push_str(" = Literal[\n");
                for (i, variant) in variants.iter().enumerate() {
                    if i > 0 {
                        out.push_str(",\n");
                    }
                    out.push_str("    ");
                    out.push_str(&variant.value.to_string());
                }
                out.push_str("\n]\n");
            }
            EnumKind::Tagged(tagged) => {
                // Generate individual variant classes
                for variant in &tagged.variants {
                    if options.frozen {
                        out.push_str("@dataclass(frozen=True)\n");
                    } else {
                        out.push_str("@dataclass\n");
                    }
                    out.push_str("class ");
                    out.push_str(&def.name);
                    out.push_str(&to_pascal_case(&variant.tag));
                    out.push_str(":\n");
                    out.push_str("    ");
                    out.push_str(&to_snake_case(&tagged.discriminator));
                    out.push_str(": Literal[\"");
                    out.push_str(&variant.tag);
                    out.push_str("\"]\n");

                    for field in &variant.fields {
                        generate_dataclass_field(out, field);
                    }
                    out.push('\n');
                }

                // Generate union type
                out.push_str(&def.name);
                out.push_str(" = Union[\n");
                for (i, variant) in tagged.variants.iter().enumerate() {
                    if i > 0 {
                        out.push_str(",\n");
                    }
                    out.push_str("    ");
                    out.push_str(&def.name);
                    out.push_str(&to_pascal_case(&variant.tag));
                }
                out.push_str("\n]\n");
            }
        },
        TypeDefKind::Alias(ty) => {
            out.push_str(&def.name);
            out.push_str(" = ");
            out.push_str(&type_to_python(ty));
            out.push('\n');
        }
    }
}

fn generate_dataclass_field(out: &mut String, field: &Field) {
    if let Some(docs) = &field.docs {
        out.push_str("    # ");
        out.push_str(docs);
        out.push('\n');
    }

    out.push_str("    ");
    out.push_str(&to_snake_case(&field.name));
    out.push_str(": ");

    if field.required {
        out.push_str(&type_to_python(&field.ty));
    } else {
        out.push_str("Optional[");
        out.push_str(&type_to_python(&field.ty));
        out.push_str("] = None");
    }
    out.push('\n');
}

fn generate_typeddict_field(out: &mut String, field: &Field) {
    if let Some(docs) = &field.docs {
        out.push_str("    # ");
        out.push_str(docs);
        out.push('\n');
    }

    out.push_str("    ");
    out.push_str(&to_snake_case(&field.name));
    out.push_str(": ");
    out.push_str(&type_to_python(&field.ty));
    out.push('\n');
}

fn type_to_python(ty: &Type) -> String {
    match ty {
        Type::String => "str".to_string(),
        Type::Integer { .. } => "int".to_string(),
        Type::Float { .. } => "float".to_string(),
        Type::Boolean => "bool".to_string(),
        Type::Null => "None".to_string(),
        Type::Array(inner) => format!("list[{}]", type_to_python(inner)),
        Type::Map { value, .. } => format!("dict[str, {}]", type_to_python(value)),
        Type::Optional(inner) => format!("Optional[{}]", type_to_python(inner)),
        Type::Ref(name) => name.clone(),
        Type::Union(types) => {
            let parts: Vec<_> = types.iter().map(type_to_python).collect();
            format!("Union[{}]", parts.join(", "))
        }
        Type::StringLiteral(s) => format!("Literal[\"{}\"]", s),
        Type::IntLiteral(i) => format!("Literal[{}]", i),
        Type::BoolLiteral(b) => format!("Literal[{}]", if *b { "True" } else { "False" }),
        Type::Any => "Any".to_string(),
    }
}

fn to_snake_case(s: &str) -> String {
    let mut result = String::new();
    for (i, c) in s.chars().enumerate() {
        if c.is_uppercase() {
            if i > 0 {
                result.push('_');
            }
            // normalize-syntax-allow: rust/unwrap-in-impl - to_lowercase() always yields at least one char
            result.push(c.to_lowercase().next().unwrap());
        } else {
            result.push(c);
        }
    }
    result
}

fn to_pascal_case(s: &str) -> String {
    let mut result = String::new();
    let mut capitalize_next = true;
    for c in s.chars() {
        if c == '_' || c == '-' {
            capitalize_next = true;
        } else if capitalize_next {
            // normalize-syntax-allow: rust/unwrap-in-impl - to_uppercase() always yields at least one char
            result.push(c.to_uppercase().next().unwrap());
            capitalize_next = false;
        } else {
            result.push(c);
        }
    }
    result
}

/// Static backend instance with default options.
pub static PYTHON_BACKEND: PythonBackend = PythonBackend {
    options: PythonOptions {
        style: PythonStyle::Dataclass,
        future_annotations: true,
        frozen: false,
    },
};

/// Python backend with configurable options.
pub struct PythonBackend {
    /// Generation options.
    pub options: PythonOptions,
}

impl PythonBackend {
    /// Create a new Python backend with the given options.
    pub fn new(options: PythonOptions) -> Self {
        Self { options }
    }
}

impl Backend for PythonBackend {
    fn name(&self) -> &'static str {
        "python"
    }

    fn language(&self) -> &'static str {
        "python"
    }

    fn extension(&self) -> &'static str {
        "py"
    }

    fn category(&self) -> BackendCategory {
        BackendCategory::Types
    }

    fn generate(&self, schema: &Schema) -> String {
        generate_python_types(schema, &self.options)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ir::{Field, TypeDef};

    #[test]
    fn generate_simple_dataclass() {
        let mut schema = Schema::default();
        schema.add(TypeDef::structure(
            "User",
            vec![
                Field::required("id", Type::String),
                Field::required("name", Type::String),
                Field::optional("email", Type::String),
            ],
        ));

        let output = generate_python_types(&schema, &PythonOptions::default());

        assert!(output.contains("@dataclass"));
        assert!(output.contains("class User:"));
        assert!(output.contains("id: str"));
        assert!(output.contains("name: str"));
        assert!(output.contains("email: Optional[str] = None"));
    }

    #[test]
    fn generate_string_enum() {
        let mut schema = Schema::default();
        schema.add(TypeDef::string_enum(
            "Status",
            vec!["pending", "active", "done"],
        ));

        let output = generate_python_types(&schema, &PythonOptions::default());

        assert!(output.contains("Status = Literal["));
        assert!(output.contains("\"pending\""));
        assert!(output.contains("\"active\""));
        assert!(output.contains("\"done\""));
    }

    #[test]
    fn generate_typeddict() {
        let mut schema = Schema::default();
        schema.add(TypeDef::structure(
            "User",
            vec![Field::required("id", Type::String)],
        ));

        let output = generate_python_types(
            &schema,
            &PythonOptions {
                style: PythonStyle::TypedDict,
                ..Default::default()
            },
        );

        assert!(output.contains("class User(TypedDict):"));
        assert!(output.contains("id: str"));
    }

    #[test]
    fn generate_frozen_dataclass() {
        let mut schema = Schema::default();
        schema.add(TypeDef::structure(
            "User",
            vec![Field::required("id", Type::String)],
        ));

        let output = generate_python_types(
            &schema,
            &PythonOptions {
                frozen: true,
                ..Default::default()
            },
        );

        assert!(output.contains("@dataclass(frozen=True)"));
    }
}
