//! Zod schema output backend.
//!
//! Generates Zod schemas that can be used for runtime validation.
//! Types can be inferred from Zod schemas using `z.infer<typeof schema>`.

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

/// Options for Zod code generation.
#[derive(Debug, Clone, Default)]
pub struct ZodOptions {
    /// Whether to export all schemas.
    pub export: bool,
    /// Whether to also generate inferred type exports.
    pub infer_types: bool,
}

/// Generate Zod schemas from an IR schema.
pub fn generate_zod(schema: &Schema, options: &ZodOptions) -> String {
    let mut out = String::new();
    out.push_str("// Auto-generated by normalize-typegen\n");
    out.push_str("import { z } from \"zod\";\n\n");

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

        if options.infer_types {
            out.push('\n');
            generate_inferred_type(&mut out, def, options);
        }
    }

    out
}

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

    let export = if options.export { "export " } else { "" };

    match &def.kind {
        TypeDefKind::Struct(s) => {
            out.push_str(export);
            out.push_str("const ");
            out.push_str(&schema_name(&def.name));
            out.push_str(" = z.object({\n");

            for field in &s.fields {
                generate_field(out, field);
            }

            out.push_str("});\n");
        }
        TypeDefKind::Enum(e) => match &e.kind {
            EnumKind::StringLiteral(variants) => {
                out.push_str(export);
                out.push_str("const ");
                out.push_str(&schema_name(&def.name));
                out.push_str(" = z.enum([");

                for (i, variant) in variants.iter().enumerate() {
                    if i > 0 {
                        out.push_str(", ");
                    }
                    out.push('"');
                    out.push_str(&variant.value);
                    out.push('"');
                }

                out.push_str("]);\n");
            }
            EnumKind::IntLiteral(variants) => {
                out.push_str(export);
                out.push_str("const ");
                out.push_str(&schema_name(&def.name));
                out.push_str(" = z.union([");

                for (i, variant) in variants.iter().enumerate() {
                    if i > 0 {
                        out.push_str(", ");
                    }
                    out.push_str("z.literal(");
                    out.push_str(&variant.value.to_string());
                    out.push(')');
                }

                out.push_str("]);\n");
            }
            EnumKind::Tagged(tagged) => {
                out.push_str(export);
                out.push_str("const ");
                out.push_str(&schema_name(&def.name));
                out.push_str(" = z.discriminatedUnion(\"");
                out.push_str(&tagged.discriminator);
                out.push_str("\", [\n");

                for variant in &tagged.variants {
                    out.push_str("  z.object({\n");
                    out.push_str("    ");
                    out.push_str(&tagged.discriminator);
                    out.push_str(": z.literal(\"");
                    out.push_str(&variant.tag);
                    out.push_str("\"),\n");

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

                    out.push_str("  }),\n");
                }

                out.push_str("]);\n");
            }
        },
        TypeDefKind::Alias(ty) => {
            out.push_str(export);
            out.push_str("const ");
            out.push_str(&schema_name(&def.name));
            out.push_str(" = ");
            out.push_str(&type_to_zod(ty));
            out.push_str(";\n");
        }
    }
}

fn generate_inferred_type(out: &mut String, def: &TypeDef, options: &ZodOptions) {
    let export = if options.export { "export " } else { "" };
    out.push_str(export);
    out.push_str("type ");
    out.push_str(&def.name);
    out.push_str(" = z.infer<typeof ");
    out.push_str(&schema_name(&def.name));
    out.push_str(">;\n");
}

fn generate_field(out: &mut String, field: &Field) {
    // Doc comment
    if let Some(docs) = &field.docs {
        out.push_str("  /** ");
        out.push_str(docs);
        out.push_str(" */\n");
    }

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

    let zod_type = type_to_zod(&field.ty);
    out.push_str(&zod_type);

    if !field.required {
        out.push_str(".optional()");
    }

    out.push_str(",\n");
}

fn generate_field_inline(out: &mut String, field: &Field) {
    out.push_str(&field.name);
    out.push_str(": ");
    out.push_str(&type_to_zod(&field.ty));
    if !field.required {
        out.push_str(".optional()");
    }
    out.push(',');
}

fn type_to_zod(ty: &Type) -> String {
    match ty {
        Type::String => "z.string()".to_string(),
        Type::Integer { .. } => "z.number().int()".to_string(),
        Type::Float { .. } => "z.number()".to_string(),
        Type::Boolean => "z.boolean()".to_string(),
        Type::Null => "z.null()".to_string(),
        Type::Array(inner) => format!("z.array({})", type_to_zod(inner)),
        Type::Map { value, .. } => format!("z.record(z.string(), {})", type_to_zod(value)),
        Type::Optional(inner) => format!("{}.optional()", type_to_zod(inner)),
        Type::Ref(name) => schema_name(name),
        Type::Union(types) => {
            if types.len() == 1 {
                return type_to_zod(&types[0]);
            }
            let parts: Vec<_> = types.iter().map(type_to_zod).collect();
            format!("z.union([{}])", parts.join(", "))
        }
        Type::StringLiteral(s) => format!("z.literal(\"{}\")", s),
        Type::IntLiteral(i) => format!("z.literal({})", i),
        Type::BoolLiteral(b) => format!("z.literal({})", b),
        Type::Any => "z.unknown()".to_string(),
    }
}

fn schema_name(type_name: &str) -> String {
    // Convention: UserSchema for User type
    format!("{}Schema", type_name)
}

/// Static backend instance with default options.
pub static ZOD_BACKEND: ZodBackend = ZodBackend {
    options: ZodOptions {
        export: true,
        infer_types: true,
    },
};

/// Zod backend with configurable options.
pub struct ZodBackend {
    /// Generation options.
    pub options: ZodOptions,
}

impl ZodBackend {
    /// Create a new Zod backend with the given options.
    pub fn new(options: ZodOptions) -> Self {
        Self { options }
    }
}

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

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

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

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

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

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

    #[test]
    fn generate_simple_object() {
        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_zod(&schema, &ZodOptions::default());

        assert!(output.contains("import { z } from \"zod\""));
        assert!(output.contains("const UserSchema = z.object("));
        assert!(output.contains("id: z.string(),"));
        assert!(output.contains("name: z.string(),"));
        assert!(output.contains("email: z.string().optional(),"));
    }

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

        let output = generate_zod(&schema, &ZodOptions::default());

        assert!(
            output.contains("const StatusSchema = z.enum([\"pending\", \"active\", \"done\"])")
        );
    }

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

        let output = generate_zod(
            &schema,
            &ZodOptions {
                export: true,
                infer_types: true,
            },
        );

        assert!(output.contains("export const UserSchema = z.object("));
        assert!(output.contains("export type User = z.infer<typeof UserSchema>"));
    }
}
