use super::gen_enum_stub;
use crate::core::ir::{CoreWrapper, EnumDef, EnumVariant, FieldDef, MethodDef, PrimitiveType, TypeRef};
use ahash::AHashSet;

/// No dataclass-backed config DTOs — factory params map exactly as `python_type` would.
fn no_dtos() -> AHashSet<&'static str> {
    AHashSet::new()
}

fn field(name: &str, ty: TypeRef) -> FieldDef {
    FieldDef {
        name: name.to_string(),
        ty,
        optional: false,
        default: None,
        doc: String::new(),
        sanitized: false,
        is_boxed: false,
        type_rust_path: None,
        cfg: None,
        typed_default: None,
        core_wrapper: CoreWrapper::None,
        vec_inner_core_wrapper: CoreWrapper::None,
        newtype_wrapper: None,
        serde_rename: None,
        serde_flatten: false,
        binding_excluded: false,
        binding_exclusion_reason: None,
        original_type: None,
    }
}

fn optional_field(name: &str, ty: TypeRef) -> FieldDef {
    FieldDef {
        optional: true,
        ..field(name, ty)
    }
}

fn variant(name: &str, fields: Vec<FieldDef>) -> EnumVariant {
    EnumVariant {
        name: name.to_string(),
        fields,
        doc: String::new(),
        is_default: false,
        serde_rename: None,
        binding_excluded: false,
        binding_exclusion_reason: None,
        is_tuple: false,
        originally_had_data_fields: false,
        cfg: None,
        version: Default::default(),
    }
}

fn enum_def(name: &str, variants: Vec<EnumVariant>) -> EnumDef {
    EnumDef {
        name: name.to_string(),
        rust_path: format!("test_lib::{name}"),
        original_rust_path: String::new(),
        variants,
        methods: vec![],
        doc: String::new(),
        cfg: None,
        is_copy: false,
        has_serde: true,
        has_default: false,
        serde_tag: Some("type".to_string()),
        serde_untagged: false,
        serde_rename_all: None,
        binding_excluded: false,
        binding_exclusion_reason: None,
        excluded_variants: vec![],
        version: Default::default(),
    }
}

fn shape_enum() -> EnumDef {
    enum_def(
        "Shape",
        vec![
            variant("Circle", vec![field("radius", TypeRef::Primitive(PrimitiveType::F64))]),
            variant(
                "Rect",
                vec![
                    field("width", TypeRef::Primitive(PrimitiveType::U32)),
                    field("height", TypeRef::Primitive(PrimitiveType::U32)),
                ],
            ),
        ],
    )
}

#[test]
fn emits_staticmethod_constructor_per_struct_variant() {
    let stub = gen_enum_stub(&shape_enum(), false, &no_dtos());

    assert!(stub.contains("class Shape:"), "{stub}");
    assert!(stub.contains("    type: str"), "{stub}");
    assert!(
        stub.contains("    @staticmethod\n    def circle(radius: float) -> Shape: ..."),
        "{stub}"
    );
    assert!(
        stub.contains("    @staticmethod\n    def rect(width: int, height: int) -> Shape: ..."),
        "{stub}"
    );
    // Constructors precede the dunder stubs.
    let circle_at = stub.find("def circle").unwrap();
    let str_at = stub.find("def __str__").unwrap();
    assert!(circle_at < str_at, "constructors must precede dunders: {stub}");
}

#[test]
fn maps_named_dto_field_to_its_type() {
    let def = enum_def(
        "Source",
        vec![variant(
            "Llm",
            vec![field("config", TypeRef::Named("LlmConfig".to_string()))],
        )],
    );

    let stub = gen_enum_stub(&def, false, &no_dtos());

    assert!(
        stub.contains("    @staticmethod\n    def llm(config: LlmConfig) -> Source: ..."),
        "{stub}"
    );
}

#[test]
fn widens_dataclass_backed_config_dto_factory_param() {
    // When the payload type is a dataclass-backed config DTO, the factory param must accept the
    // public `options` dataclass or a dict (what the runtime coercion accepts) — not the compiled
    // class the bare name resolves to. A primitive sibling param is unaffected.
    let def = enum_def(
        "EmbeddingModelType",
        vec![
            variant("Llm", vec![field("llm", TypeRef::Named("LlmConfig".to_string()))]),
            variant("Preset", vec![field("name", TypeRef::String)]),
        ],
    );
    let coercible: AHashSet<&str> = ["LlmConfig"].into_iter().collect();

    let stub = gen_enum_stub(&def, false, &coercible);

    assert!(
        stub.contains(
            "    @staticmethod\n    def llm(llm: options.LlmConfig | dict[str, Any]) -> EmbeddingModelType: ..."
        ),
        "coercible DTO factory param must accept the public dataclass or a dict: {stub}"
    );
    assert!(
        stub.contains("    @staticmethod\n    def preset(name: str) -> EmbeddingModelType: ..."),
        "primitive factory param must be unchanged: {stub}"
    );
}

#[test]
fn qualifies_builtin_shadowed_by_a_variant_factory_name() {
    // A `List` variant emits `def list(...)`, which shadows the builtin `list` within the class
    // body — so a sibling factory annotated `list[MetadataEntry]` would resolve to the factory and
    // mypy rejects it (`Function ... is not valid as a type`). The shadowed builtin container must be
    // qualified as `builtins.list[...]` in the factory annotations.
    let def = enum_def(
        "NodeContent",
        vec![
            variant("List", vec![field("ordered", TypeRef::String)]),
            variant(
                "MetadataBlock",
                vec![field(
                    "entries",
                    TypeRef::Vec(Box::new(TypeRef::Named("MetadataEntry".to_string()))),
                )],
            ),
        ],
    );

    let stub = gen_enum_stub(&def, false, &no_dtos());

    assert!(
        stub.contains("    @staticmethod\n    def list(ordered: str) -> NodeContent: ..."),
        "{stub}"
    );
    assert!(
        stub.contains("def metadata_block(entries: builtins.list[MetadataEntry]) -> NodeContent: ..."),
        "builtin shadowed by the `list` factory must be qualified as builtins.list: {stub}"
    );
}

#[test]
fn skips_unit_tuple_excluded_and_sanitized_variants() {
    let mut tuple_variant = variant("Pair", vec![field("_0", TypeRef::String)]);
    tuple_variant.is_tuple = true;
    let mut excluded = variant("Hidden", vec![field("value", TypeRef::String)]);
    excluded.binding_excluded = true;
    let mut sanitized_field = field("raw", TypeRef::String);
    sanitized_field.sanitized = true;
    let sanitized_variant = variant("Raw", vec![sanitized_field]);

    let def = enum_def(
        "Shape",
        vec![
            variant("Empty", vec![]),
            tuple_variant,
            excluded,
            sanitized_variant,
            variant("Real", vec![field("value", TypeRef::String)]),
        ],
    );

    let stub = gen_enum_stub(&def, false, &no_dtos());

    assert!(!stub.contains("def empty("), "{stub}");
    assert!(!stub.contains("def pair("), "{stub}");
    assert!(!stub.contains("def hidden("), "{stub}");
    assert!(!stub.contains("def raw("), "{stub}");
    assert!(
        stub.contains("    @staticmethod\n    def real(value: str) -> Shape: ..."),
        "{stub}"
    );
}

#[test]
fn optional_field_is_nilable_with_default() {
    let def = enum_def(
        "Source",
        vec![variant("Tag", vec![optional_field("label", TypeRef::String)])],
    );

    let stub = gen_enum_stub(&def, false, &no_dtos());

    assert!(
        stub.contains("    @staticmethod\n    def tag(label: str | None = None) -> Source: ..."),
        "{stub}"
    );
}

#[test]
fn param_after_optional_is_promoted_to_nilable() {
    // `width` is not optional in the IR, but it follows the optional `radius`, so the runtime PyO3
    // signature widens it to `Optional[int] = None`. The stub must match (no required param may
    // follow a defaulted one).
    let def = enum_def(
        "Shape",
        vec![variant(
            "Ring",
            vec![
                optional_field("radius", TypeRef::Primitive(PrimitiveType::F64)),
                field("width", TypeRef::Primitive(PrimitiveType::U32)),
            ],
        )],
    );

    let stub = gen_enum_stub(&def, false, &no_dtos());

    assert!(
        stub.contains(
            "    @staticmethod\n    def ring(radius: float | None = None, width: int | None = None) -> Shape: ..."
        ),
        "{stub}"
    );
}

#[test]
fn yields_to_hand_written_method_of_same_name() {
    let def = EnumDef {
        methods: vec![MethodDef {
            name: "circle".to_string(),
            is_static: true,
            ..Default::default()
        }],
        ..shape_enum()
    };

    let stub = gen_enum_stub(&def, false, &no_dtos());

    assert!(!stub.contains("def circle("), "hand-written method wins: {stub}");
    assert!(
        stub.contains("    @staticmethod\n    def rect(width: int, height: int) -> Shape: ..."),
        "{stub}"
    );
}
