//! Pydantic model output backend.
//!
//! Generates Pydantic v2 models for runtime validation.

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

/// Options for Pydantic code generation.
#[derive(Debug, Clone, Default)]
pub struct PydanticOptions {
    /// Pydantic version: V1 or V2 (default).
    pub version: PydanticVersion,
    /// Whether to use strict mode.
    pub strict: bool,
    /// Whether to generate ConfigDict.
    pub frozen: bool,
}

#[derive(Debug, Clone, Default, PartialEq)]
pub enum PydanticVersion {
    V1,
    #[default]
    V2,
}

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

    // Imports
    match options.version {
        PydanticVersion::V1 => {
            out.push_str("from pydantic import BaseModel, Field\n");
        }
        PydanticVersion::V2 => {
            out.push_str("from pydantic import BaseModel, Field, ConfigDict\n");
        }
    }

    // Typing imports
    let mut typing_imports = vec!["Optional"];
    let has_enums = schema
        .definitions
        .iter()
        .any(|d| matches!(d.kind, TypeDefKind::Enum(_)));
    if has_enums {
        typing_imports.push("Literal");
    }
    let has_unions = schema.definitions.iter().any(|d| {
        matches!(
            &d.kind,
            TypeDefKind::Enum(e) if matches!(e.kind, EnumKind::Tagged(_))
        )
    });
    if has_unions {
        typing_imports.push("Union");
        if options.version == PydanticVersion::V2 {
            typing_imports.push("Annotated");
        }
    }

    out.push_str("from typing import ");
    out.push_str(&typing_imports.join(", "));
    out.push('\n');

    if has_unions && options.version == PydanticVersion::V2 {
        out.push_str("from pydantic import Discriminator\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 generate_typedef(out: &mut String, def: &TypeDef, options: &PydanticOptions) {
    match &def.kind {
        TypeDefKind::Struct(s) => {
            // Doc comment
            if let Some(docs) = &def.docs {
                out.push_str(&format!("\"\"\"{}.\"\"\"\n", docs));
            }

            out.push_str("class ");
            out.push_str(&def.name);
            out.push_str("(BaseModel):\n");

            // Config
            if options.frozen || options.strict {
                if options.version == PydanticVersion::V2 {
                    out.push_str("    model_config = ConfigDict(");
                    let mut config_parts = Vec::new();
                    if options.frozen {
                        config_parts.push("frozen=True");
                    }
                    if options.strict {
                        config_parts.push("strict=True");
                    }
                    out.push_str(&config_parts.join(", "));
                    out.push_str(")\n\n");
                } else {
                    out.push_str("    class Config:\n");
                    if options.frozen {
                        out.push_str("        frozen = True\n");
                    }
                    if options.strict {
                        out.push_str("        strict = True\n");
                    }
                    out.push('\n');
                }
            }

            if s.fields.is_empty() {
                out.push_str("    pass\n");
            } else {
                for field in &s.fields {
                    generate_field(out, field);
                }
            }
        }
        TypeDefKind::Enum(e) => match &e.kind {
            EnumKind::StringLiteral(variants) => {
                if let Some(docs) = &def.docs {
                    out.push_str(&format!("\"\"\"{}.\"\"\"\n", docs));
                }
                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) => {
                if let Some(docs) = &def.docs {
                    out.push_str(&format!("\"\"\"{}.\"\"\"\n", docs));
                }
                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 models
                for variant in &tagged.variants {
                    if let Some(docs) = &variant.docs {
                        out.push_str(&format!("\"\"\"{}.\"\"\"\n", docs));
                    }
                    out.push_str("class ");
                    out.push_str(&def.name);
                    out.push_str(&to_pascal_case(&variant.tag));
                    out.push_str("(BaseModel):\n");

                    out.push_str("    ");
                    out.push_str(&to_snake_case(&tagged.discriminator));
                    out.push_str(": Literal[\"");
                    out.push_str(&variant.tag);
                    out.push_str("\"] = \"");
                    out.push_str(&variant.tag);
                    out.push_str("\"\n");

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

                // Generate discriminated union
                if let Some(docs) = &def.docs {
                    out.push_str(&format!("\"\"\"{}.\"\"\"\n", docs));
                }

                if options.version == PydanticVersion::V2 {
                    out.push_str(&def.name);
                    out.push_str(" = Annotated[\n    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    Discriminator(\"");
                    out.push_str(&to_snake_case(&tagged.discriminator));
                    out.push_str("\")\n]\n");
                } else {
                    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) => {
            if let Some(docs) = &def.docs {
                out.push_str(&format!("\"\"\"{}.\"\"\"\n", docs));
            }
            out.push_str(&def.name);
            out.push_str(" = ");
            out.push_str(&type_to_python(ty));
            out.push('\n');
        }
    }
}

fn generate_field(out: &mut String, field: &Field) {
    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(']');
    }

    // Add Field() with description if present
    if field.docs.is_some() || !field.required {
        out.push_str(" = Field(");
        let mut parts = Vec::new();
        if !field.required {
            parts.push("default=None".to_string());
        }
        if let Some(docs) = &field.docs {
            parts.push(format!("description=\"{}\"", docs.replace('"', "\\\"")));
        }
        out.push_str(&parts.join(", "));
        out.push(')');
    }
    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) => format!("\"{}\"", name), // Forward reference
        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 PYDANTIC_BACKEND: PydanticBackend = PydanticBackend {
    options: PydanticOptions {
        version: PydanticVersion::V2,
        strict: false,
        frozen: false,
    },
};

/// Pydantic backend with configurable options.
pub struct PydanticBackend {
    /// Generation options.
    pub options: PydanticOptions,
}

impl PydanticBackend {
    /// Create a new Pydantic backend with the given options.
    pub fn new(options: PydanticOptions) -> Self {
        Self { options }
    }
}

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

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

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

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

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

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

    #[test]
    fn generate_simple_model() {
        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_pydantic(&schema, &PydanticOptions::default());

        assert!(output.contains("from pydantic import BaseModel"));
        assert!(output.contains("class User(BaseModel):"));
        assert!(output.contains("id: str"));
        assert!(output.contains("name: str"));
        assert!(output.contains("email: Optional[str] = Field(default=None)"));
    }

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

        let output = generate_pydantic(&schema, &PydanticOptions::default());

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

    #[test]
    fn generate_frozen_model() {
        let mut schema = Schema::default();
        schema.add(TypeDef::structure("User", vec![]));

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

        assert!(output.contains("model_config = ConfigDict(frozen=True)"));
    }
}
