"""Unit tests for :class:`~llmner.base_model.NERBaseModel`."""

from __future__ import annotations

from pydantic import Field

from llmner import NERBaseModel, SchemaRegistry
from llmner.factories import EvidenceField

# ---------------------------------------------------------------------------
# Helpers – reusable test schemas
# ---------------------------------------------------------------------------

_reg = SchemaRegistry()
_NameType = _reg.generic("name", "Full name.")
_AgeType = _reg.int_range("age", "Age in years.")


class SimpleSchema(NERBaseModel):
    name: _NameType | None = None  # type: ignore[valid-type]
    age: _AgeType | None = None  # type: ignore[valid-type]


_reg2 = SchemaRegistry()
_InnerName = _reg2.generic("inner_name", "Name.")


class InnerModel(NERBaseModel):
    inner_name: _InnerName | None = None  # type: ignore[valid-type]


class NestedSchema(NERBaseModel):
    child: InnerModel | None = None
    items: list[SimpleSchema] = Field(default_factory=list)


# ---------------------------------------------------------------------------
# prompt_schema
# ---------------------------------------------------------------------------


class TestPromptSchema:
    def test_flat_schema_has_all_keys(self) -> None:
        schema = SimpleSchema.prompt_schema()
        assert '"name"' in schema
        assert '"age"' in schema
        assert '"value"' in schema
        assert '"evidence"' in schema

    def test_nested_schema(self) -> None:
        schema = NestedSchema.prompt_schema()
        assert '"child"' in schema
        assert '"items"' in schema
        assert '"inner_name"' in schema

    def test_returns_valid_json(self) -> None:
        import json

        data = json.loads(SimpleSchema.prompt_schema())
        assert isinstance(data, dict)
        assert "name" in data
        assert "age" in data


# ---------------------------------------------------------------------------
# has_missing_fields
# ---------------------------------------------------------------------------


class TestHasMissingFields:
    def test_all_filled(self) -> None:
        obj = SimpleSchema(
            name=EvidenceField(value="John", evidence="John"),
            age=EvidenceField(value="30", evidence="30"),
        )
        assert obj.has_missing_fields() is False

    def test_value_none_triggers_true(self) -> None:
        obj = SimpleSchema(
            name=EvidenceField(value=None, evidence=None),
            age=EvidenceField(value="30", evidence="30"),
        )
        assert obj.has_missing_fields() is True

    def test_evidence_none_does_not_trigger(self) -> None:
        """Missing evidence should NOT trigger a retry."""
        obj = SimpleSchema(
            name=EvidenceField(value="John", evidence=None),
            age=EvidenceField(value="30", evidence=None),
        )
        assert obj.has_missing_fields() is False

    def test_top_level_none_is_skipped(self) -> None:
        """A field that is entirely None is not considered missing."""
        obj = SimpleSchema(name=None, age=EvidenceField(value="30", evidence="30"))
        assert obj.has_missing_fields() is False

    def test_nested_model_missing_value(self) -> None:
        inner = InnerModel(inner_name=EvidenceField(value=None, evidence=None))
        obj = NestedSchema(child=inner)
        assert obj.has_missing_fields() is True

    def test_list_item_missing_value(self) -> None:
        item = SimpleSchema(
            name=EvidenceField(value="John", evidence=None),
            age=EvidenceField(value=None, evidence=None),
        )
        obj = NestedSchema(items=[item])
        assert obj.has_missing_fields() is True

    def test_list_item_all_filled(self) -> None:
        item = SimpleSchema(
            name=EvidenceField(value="John", evidence=None),
            age=EvidenceField(value="30", evidence=None),
        )
        obj = NestedSchema(items=[item])
        assert obj.has_missing_fields() is False


# ---------------------------------------------------------------------------
# merge
# ---------------------------------------------------------------------------


class TestMerge:
    def test_fills_none_from_second(self) -> None:
        first = SimpleSchema(
            name=EvidenceField(value="John", evidence="John"),
            age=EvidenceField(value=None, evidence=None),
        )
        second = SimpleSchema(
            name=EvidenceField(value="Jane", evidence="Jane"),
            age=EvidenceField(value="25", evidence="25"),
        )
        merged = SimpleSchema.merge(first, second)
        assert merged.name.value == "John"  # first takes priority
        assert merged.age.value == "25"  # filled from second

    def test_merge_preserves_existing(self) -> None:
        first = SimpleSchema(
            name=EvidenceField(value="John", evidence="John"),
            age=EvidenceField(value="30", evidence="30"),
        )
        second = SimpleSchema(
            name=EvidenceField(value="Jane", evidence="Jane"),
            age=EvidenceField(value="99", evidence="99"),
        )
        merged = SimpleSchema.merge(first, second)
        assert merged.name.value == "John"
        assert merged.age.value == "30"

    def test_merge_lists_by_index(self) -> None:
        item1 = SimpleSchema(
            name=EvidenceField(value="A", evidence=None),
            age=EvidenceField(value=None, evidence=None),
        )
        item2 = SimpleSchema(
            name=EvidenceField(value=None, evidence=None),
            age=EvidenceField(value="20", evidence=None),
        )
        first = NestedSchema(items=[item1])
        second = NestedSchema(items=[item2])
        merged = NestedSchema.merge(first, second)
        assert merged.items[0].name.value == "A"
        assert merged.items[0].age.value == "20"

    def test_merge_different_length_lists(self) -> None:
        item1 = SimpleSchema(
            name=EvidenceField(value="A", evidence=None),
            age=EvidenceField(value="1", evidence=None),
        )
        item2 = SimpleSchema(
            name=EvidenceField(value="B", evidence=None),
            age=EvidenceField(value="2", evidence=None),
        )
        first = NestedSchema(items=[item1])
        second = NestedSchema(items=[item1, item2])
        merged = NestedSchema.merge(first, second)
        assert len(merged.items) == 2


# ---------------------------------------------------------------------------
# safe_parse
# ---------------------------------------------------------------------------


class TestSafeParse:
    def test_valid_data(self) -> None:
        data = {
            "name": {"value": "John", "evidence": "John"},
            "age": {"value": "30", "evidence": "30"},
        }
        result = SimpleSchema.safe_parse(data, input_text="John is 30.")
        assert result is not None
        assert result.name.value == "John"
        assert result.age.value == "30"

    def test_none_input_returns_none(self) -> None:
        assert SimpleSchema.safe_parse(None) is None

    def test_non_dict_returns_none(self) -> None:
        assert SimpleSchema.safe_parse("not a dict") is None  # type: ignore[arg-type]

    def test_invalid_field_becomes_none(self) -> None:
        data = {
            "name": {"value": "John", "evidence": "John"},
            "age": {"value": "not-a-number", "evidence": None},
        }
        result = SimpleSchema.safe_parse(data, input_text="John")
        assert result is not None
        assert result.name.value == "John"
        assert result.age.value is None

    def test_nested_model_parsed(self) -> None:
        data = {
            "child": {"inner_name": {"value": "test", "evidence": None}},
            "items": [],
        }
        result = NestedSchema.safe_parse(data)
        assert result is not None
        assert result.child is not None
        assert result.child.inner_name.value == "test"

    def test_list_of_models_parsed(self) -> None:
        data = {
            "child": None,
            "items": [
                {
                    "name": {"value": "A", "evidence": None},
                    "age": {"value": "1", "evidence": None},
                },
                {
                    "name": {"value": "B", "evidence": None},
                    "age": {"value": "2", "evidence": None},
                },
            ],
        }
        result = NestedSchema.safe_parse(data)
        assert result is not None
        assert len(result.items) == 2
        assert result.items[0].name.value == "A"
        assert result.items[1].name.value == "B"

    def test_list_item_partial_invalid(self) -> None:
        """One bad field in a list item should not discard the entire item."""
        data = {
            "child": None,
            "items": [
                {
                    "name": {"value": "OK", "evidence": None},
                    "age": {"value": "bad", "evidence": None},
                },
            ],
        }
        result = NestedSchema.safe_parse(data)
        assert result is not None
        assert len(result.items) == 1
        assert result.items[0].name.value == "OK"
