"""Unit tests for :class:`~llmner.extractor.NERExtractor`."""

from __future__ import annotations

from typing import Any

from llmner import NERBaseModel, NERExtractor, SchemaRegistry
from llmner.llm_client import BaseLLMClient

# ---------------------------------------------------------------------------
# Fake LLM client
# ---------------------------------------------------------------------------


class FakeLLMClient(BaseLLMClient):
    """Returns pre-configured responses in sequence."""

    def __init__(self, responses: list[dict[str, Any] | None]) -> None:
        self._responses = list(responses)
        self._call_count = 0

    def generate(self, prompt: str) -> dict[str, Any] | None:
        if self._call_count < len(self._responses):
            resp = self._responses[self._call_count]
            self._call_count += 1
            return resp
        return None

    @property
    def call_count(self) -> int:
        return self._call_count


# ---------------------------------------------------------------------------
# Schema
# ---------------------------------------------------------------------------

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


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


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


class TestNERExtractor:
    def _extractor(
        self,
        responses: list[dict[str, Any] | None],
        max_retries: int = 1,
    ) -> tuple[NERExtractor, FakeLLMClient]:
        client = FakeLLMClient(responses)
        ext = NERExtractor(
            schema_class=PersonSchema,
            system_role="Test role.",
            system_task="Test task.",
            rules_registry=_reg.rules,
            llm_client=client,
            max_retries=max_retries,
        )
        return ext, client

    def test_extract_one_success(self) -> None:
        ext, _ = self._extractor(
            [
                {
                    "name": {"value": "John", "evidence": "John"},
                    "age": {"value": "30", "evidence": "30"},
                },
            ]
        )
        result = ext.extract_one("John is 30.", retry_on_null=False)
        assert result is not None
        assert result.name.value == "John"
        assert result.age.value == "30"

    def test_extract_one_returns_none_on_failure(self) -> None:
        ext, _ = self._extractor([None])
        result = ext.extract_one("anything")
        assert result is None

    def test_retry_fills_missing_fields(self) -> None:
        ext, client = self._extractor(
            [
                {
                    "name": {"value": "John", "evidence": "John"},
                    "age": {"value": None, "evidence": None},
                },
                {
                    "name": {"value": None, "evidence": None},
                    "age": {"value": "30", "evidence": "30"},
                },
            ]
        )
        result = ext.extract_one("John is 30.", retry_on_null=True)
        assert result is not None
        assert result.name.value == "John"
        assert result.age.value == "30"
        assert client.call_count == 2

    def test_no_retry_when_complete(self) -> None:
        ext, client = self._extractor(
            [
                {
                    "name": {"value": "John", "evidence": "John"},
                    "age": {"value": "30", "evidence": "30"},
                },
                {
                    "name": {"value": "X", "evidence": None},
                    "age": {"value": "99", "evidence": None},
                },
            ]
        )
        result = ext.extract_one("John is 30.", retry_on_null=True)
        assert result is not None
        assert client.call_count == 1  # no retry needed

    def test_retry_on_null_false_skips_retries(self) -> None:
        ext, client = self._extractor(
            [
                {
                    "name": {"value": "John", "evidence": None},
                    "age": {"value": None, "evidence": None},
                },
                {
                    "name": {"value": None, "evidence": None},
                    "age": {"value": "30", "evidence": None},
                },
            ]
        )
        result = ext.extract_one("John is 30.", retry_on_null=False)
        assert result is not None
        assert result.age.value is None  # not retried
        assert client.call_count == 1

    def test_max_retries_respected(self) -> None:
        ext, client = self._extractor(
            [
                {
                    "name": {"value": None, "evidence": None},
                    "age": {"value": None, "evidence": None},
                },
                {
                    "name": {"value": None, "evidence": None},
                    "age": {"value": None, "evidence": None},
                },
                {
                    "name": {"value": None, "evidence": None},
                    "age": {"value": None, "evidence": None},
                },
                {
                    "name": {"value": "late", "evidence": None},
                    "age": {"value": "1", "evidence": None},
                },
            ],
            max_retries=2,
        )
        ext.extract_one("test", retry_on_null=True)
        assert client.call_count == 3  # 1 initial + 2 retries

    def test_retry_does_not_trigger_on_missing_evidence(self) -> None:
        """Retries should only trigger on missing values, not missing evidence."""
        ext, client = self._extractor(
            [
                {
                    "name": {"value": "John", "evidence": None},
                    "age": {"value": "30", "evidence": None},
                },
                {
                    "name": {"value": "X", "evidence": "X"},
                    "age": {"value": "X", "evidence": "X"},
                },
            ]
        )
        result = ext.extract_one("John is 30.", retry_on_null=True)
        assert result is not None
        assert client.call_count == 1  # no retry because values are present
