"""Unit tests for :class:`~llmner.factories.SchemaRegistry` type factories."""

from __future__ import annotations

from llmner import NERBaseModel, SchemaRegistry

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_schema(registry: SchemaRegistry, **field_types: type) -> type[NERBaseModel]:
    """Dynamically create a NERBaseModel subclass with the given fields."""
    annotations = {name: ft | None for name, ft in field_types.items()}
    defaults = {name: None for name in field_types}
    ns: dict = {"__annotations__": annotations, **defaults}
    return type("DynSchema", (NERBaseModel,), ns)


def _parse(schema_cls: type[NERBaseModel], data: dict, text: str = "") -> NERBaseModel:
    """Shortcut: validate *data* through safe_parse."""
    result = schema_cls.safe_parse(data, input_text=text)
    assert result is not None
    return result


# ---------------------------------------------------------------------------
# categorical
# ---------------------------------------------------------------------------


class TestCategorical:
    def setup_method(self) -> None:
        self.reg = SchemaRegistry()
        self.Gender = self.reg.categorical(
            "gender", options=["male", "female"], instruction="Gender."
        )
        self.Schema = _make_schema(self.reg, gender=self.Gender)

    def test_valid_value(self) -> None:
        res = _parse(
            self.Schema,
            {"gender": {"value": "male", "evidence": "male"}},
            "He is male.",
        )
        assert res.gender.value == "male"

    def test_case_insensitive(self) -> None:
        res = _parse(self.Schema, {"gender": {"value": "Male", "evidence": None}})
        assert res.gender.value == "male"

    def test_invalid_value_becomes_none(self) -> None:
        res = _parse(self.Schema, {"gender": {"value": "robot", "evidence": None}})
        assert res.gender.value is None

    def test_null_string_becomes_none(self) -> None:
        res = _parse(self.Schema, {"gender": {"value": "null", "evidence": None}})
        assert res.gender.value is None

    def test_registers_rule(self) -> None:
        assert len(self.reg.rules) == 1
        assert "GENDER" in self.reg.rules[0]
        assert "male" in self.reg.rules[0]

    def test_dict_options_with_descriptions(self) -> None:
        reg = SchemaRegistry()
        Status = reg.categorical(
            "status",
            options={"active": "currently working", "inactive": "retired"},
            instruction="Status.",
        )
        Schema = _make_schema(reg, status=Status)
        res = _parse(Schema, {"status": {"value": "active", "evidence": None}})
        assert res.status.value == "active"
        assert "active" in reg.rules[0]
        assert "currently working" in reg.rules[0]

    def test_replacements(self) -> None:
        reg = SchemaRegistry()
        Hair = reg.categorical(
            "hair",
            options=["dark_brown", "black"],
            instruction="Hair colour.",
            replacements={" ": "_"},
        )
        Schema = _make_schema(reg, hair=Hair)
        res = _parse(Schema, {"hair": {"value": "dark brown", "evidence": None}})
        assert res.hair.value == "dark_brown"

    def test_allow_multiple(self) -> None:
        reg = SchemaRegistry()
        Weapon = reg.categorical(
            "weapon",
            options=["knife", "gun", "fists"],
            instruction="Weapon.",
            allow_multiple=True,
        )
        Schema = _make_schema(reg, weapon=Weapon)
        res = _parse(Schema, {"weapon": {"value": "knife and fists", "evidence": None}})
        assert "knife" in res.weapon.value
        assert "fists" in res.weapon.value


# ---------------------------------------------------------------------------
# int_range
# ---------------------------------------------------------------------------


class TestIntRange:
    def setup_method(self) -> None:
        self.reg = SchemaRegistry()
        self.Age = self.reg.int_range("age", "Age in years.")
        self.Schema = _make_schema(self.reg, age=self.Age)

    def test_single_integer(self) -> None:
        res = _parse(self.Schema, {"age": {"value": "25", "evidence": None}})
        assert res.age.value == "25"

    def test_range(self) -> None:
        res = _parse(self.Schema, {"age": {"value": "25-30", "evidence": None}})
        assert res.age.value == "25-30"

    def test_range_with_spaces(self) -> None:
        res = _parse(self.Schema, {"age": {"value": "25 - 30", "evidence": None}})
        assert res.age.value == "25-30"

    def test_metre_conversion(self) -> None:
        reg = SchemaRegistry()
        Height = reg.int_range("height", "Height in cm.")
        Schema = _make_schema(reg, height=Height)
        res = _parse(Schema, {"height": {"value": "1.75", "evidence": None}})
        assert res.height.value == "175"

    def test_null_string(self) -> None:
        res = _parse(self.Schema, {"age": {"value": "null", "evidence": None}})
        assert res.age.value is None

    def test_non_numeric_becomes_none(self) -> None:
        res = _parse(self.Schema, {"age": {"value": "unknown", "evidence": None}})
        assert res.age.value is None

    def test_registers_rule(self) -> None:
        assert len(self.reg.rules) == 1
        assert "AGE" in self.reg.rules[0]


# ---------------------------------------------------------------------------
# generic
# ---------------------------------------------------------------------------


class TestGeneric:
    def test_plain_string(self) -> None:
        reg = SchemaRegistry()
        Addr = reg.generic("address", "Full address.")
        Schema = _make_schema(reg, address=Addr)
        res = _parse(Schema, {"address": {"value": "123 Main St", "evidence": None}})
        assert res.address.value == "123 Main St"

    def test_null_becomes_none(self) -> None:
        reg = SchemaRegistry()
        Addr = reg.generic("address", "Full address.")
        Schema = _make_schema(reg, address=Addr)
        res = _parse(Schema, {"address": {"value": "null", "evidence": None}})
        assert res.address.value is None

    def test_custom_python_type(self) -> None:
        reg = SchemaRegistry()
        Count = reg.generic("count", "Number of items.", python_type=int)
        Schema = _make_schema(reg, count=Count)
        res = _parse(Schema, {"count": {"value": "42", "evidence": None}})
        # python_type validates but value is always stored as str
        assert res.count.value == "42"

    def test_custom_python_type_invalid(self) -> None:
        reg = SchemaRegistry()
        Count = reg.generic("count", "Number of items.", python_type=int)
        Schema = _make_schema(reg, count=Count)
        res = _parse(Schema, {"count": {"value": "not-a-number", "evidence": None}})
        assert res.count.value is None


# ---------------------------------------------------------------------------
# datetime_format
# ---------------------------------------------------------------------------


class TestDatetimeFormat:
    def setup_method(self) -> None:
        self.reg = SchemaRegistry()
        self.Dt = self.reg.datetime_format("datetime", "Date and time.")
        self.Schema = _make_schema(self.reg, datetime=self.Dt)

    def test_iso_format(self) -> None:
        res = _parse(
            self.Schema,
            {"datetime": {"value": "2024-03-15 14:30:00", "evidence": None}},
        )
        assert res.datetime.value == "2024-03-15 14:30:00"

    def test_short_iso(self) -> None:
        res = _parse(
            self.Schema, {"datetime": {"value": "2024-03-15 14:30", "evidence": None}}
        )
        assert res.datetime.value == "2024-03-15 14:30:00"

    def test_date_only(self) -> None:
        res = _parse(
            self.Schema, {"datetime": {"value": "2024-03-15", "evidence": None}}
        )
        assert res.datetime.value == "2024-03-15 00:00:00"

    def test_european_format(self) -> None:
        res = _parse(
            self.Schema, {"datetime": {"value": "15/03/2024", "evidence": None}}
        )
        assert res.datetime.value == "2024-03-15 00:00:00"

    def test_invalid_date_becomes_none(self) -> None:
        res = _parse(
            self.Schema, {"datetime": {"value": "not-a-date", "evidence": None}}
        )
        assert res.datetime.value is None

    def test_null_becomes_none(self) -> None:
        res = _parse(self.Schema, {"datetime": {"value": "null", "evidence": None}})
        assert res.datetime.value is None


# ---------------------------------------------------------------------------
# Evidence validation
# ---------------------------------------------------------------------------


class TestEvidenceValidation:
    def test_valid_evidence_kept(self) -> None:
        reg = SchemaRegistry()
        Name = reg.generic("name", "Name.")
        Schema = _make_schema(reg, name=Name)
        text = "The victim was John Smith, aged 34."
        res = _parse(
            Schema, {"name": {"value": "John Smith", "evidence": "John Smith"}}, text
        )
        assert res.name.evidence == "John Smith"

    def test_evidence_not_in_text_removed(self) -> None:
        reg = SchemaRegistry()
        Name = reg.generic("name", "Name.")
        Schema = _make_schema(reg, name=Name)
        text = "The victim was John Smith, aged 34."
        res = _parse(
            Schema, {"name": {"value": "John Smith", "evidence": "Jane Doe"}}, text
        )
        assert res.name.evidence is None

    def test_evidence_case_insensitive(self) -> None:
        reg = SchemaRegistry()
        Name = reg.generic("name", "Name.")
        Schema = _make_schema(reg, name=Name)
        text = "The victim was JOHN SMITH, aged 34."
        res = _parse(
            Schema, {"name": {"value": "John Smith", "evidence": "john smith"}}, text
        )
        assert res.name.evidence == "john smith"

    def test_evidence_whitespace_normalised(self) -> None:
        reg = SchemaRegistry()
        Name = reg.generic("name", "Name.")
        Schema = _make_schema(reg, name=Name)
        text = "The victim was John\n  Smith, aged 34."
        res = _parse(
            Schema, {"name": {"value": "John Smith", "evidence": "John Smith"}}, text
        )
        assert res.name.evidence == "John Smith"


# ---------------------------------------------------------------------------
# Registry rules accumulation
# ---------------------------------------------------------------------------


class TestRegistryRules:
    def test_rules_accumulate_in_order(self) -> None:
        reg = SchemaRegistry()
        reg.categorical("a", ["x"], "First.")
        reg.int_range("b", "Second.")
        reg.generic("c", "Third.")
        reg.datetime_format("d", "Fourth.")
        rules = reg.rules
        assert len(rules) == 4
        assert "A:" in rules[0]
        assert "B:" in rules[1]
        assert "C:" in rules[2]
        assert "D:" in rules[3]

    def test_rules_returns_copy(self) -> None:
        reg = SchemaRegistry()
        reg.generic("x", "Test.")
        rules = reg.rules
        rules.append("extra")
        assert len(reg.rules) == 1
