"""Type factories and the :class:`SchemaRegistry` for building NER schemas.

Each factory method on :class:`SchemaRegistry` returns a Pydantic
``Annotated`` type that:

1. Validates and normalises the raw LLM output for that field.
2. Registers a human-readable extraction rule that is later injected into
   the LLM prompt.

All extracted values are wrapped in :class:`EvidenceField`, which pairs
the extracted *value* with a short *evidence* quote from the source text.

Usage::

    from llmner import SchemaRegistry, EvidenceField

    registry = SchemaRegistry()

    GenderType = registry.categorical(
        "gender",
        options=["male", "female"],
        instruction="Extract the subject's gender.",
    )
    AgeType = registry.int_range("age", "Age in years or range, e.g. '25-30'.")

    # Pass registry.rules to NERExtractor
"""

from __future__ import annotations

import re
from datetime import datetime
from typing import Annotated, Any

from pydantic import BaseModel, BeforeValidator, ValidationInfo

__all__ = ["EvidenceField", "SchemaRegistry", "is_null_value"]

# ---------------------------------------------------------------------------
# EvidenceField
# ---------------------------------------------------------------------------


class EvidenceField(BaseModel):
    """Container for a single extracted entity and its supporting evidence.

    Attributes:
        value:    The normalised extracted value, or ``None`` if not found.
        evidence: A verbatim short quote from the source text that justifies
                  *value*.  ``None`` if no supporting quote was provided or
                  the quote could not be verified in the source.
    """

    value: str | None = None
    evidence: str | None = None


# ---------------------------------------------------------------------------
# Module-level helpers (not part of the public API)
# ---------------------------------------------------------------------------


def _coerce_null(v: Any) -> Any:
    """Convert common null-like strings to Python ``None``."""
    if isinstance(v, str) and v.strip().lower() in ("null", "", "none", "nan"):
        return None
    return v


def _normalize_ws(text: str) -> str:
    """Collapse all whitespace (newlines, tabs, etc.) into single spaces."""
    return " ".join(text.split())


def _validate_evidence(evidence: str | None, info: ValidationInfo) -> str | None:
    """Return *evidence* only when it can be found in the source text.

    The comparison is case-insensitive and whitespace-normalised: newlines,
    tabs, and runs of multiple spaces are collapsed to single spaces before
    matching.  This accounts for text that has been line-wrapped or
    reformatted by the LLM.

    If no ``input_text`` was supplied via Pydantic context the evidence is
    returned as-is (validation is skipped).
    """
    if not evidence:
        return None
    evidence = evidence.strip().strip("'\"")
    if not evidence:
        return None
    ctx = info.context or {}
    input_text: str = ctx.get("input_text", "")
    if (
        input_text
        and _normalize_ws(evidence).lower() not in _normalize_ws(input_text).lower()
    ):
        return None
    return evidence


def _extract_ev(v: Any) -> tuple[Any, str | None]:
    """Unpack ``{value, evidence}`` dicts or :class:`EvidenceField` objects."""
    evidence: str | None = None
    if isinstance(v, dict):
        evidence = v.get("evidence")
        v = v.get("value")
    elif hasattr(v, "value"):
        evidence = getattr(v, "evidence", None)
        v = v.value
    return v, evidence


# ---------------------------------------------------------------------------
# is_null_value helper
# ---------------------------------------------------------------------------


def is_null_value(val: object) -> bool:
    """Return ``True`` if *val* should be treated as a missing / null value.

    Covers Python ``None`` as well as common string representations such as
    ``"null"``, ``"none"``, ``"nan"``, and empty strings.
    """
    return val is None or str(val).strip().lower() in {
        "null",
        "",
        "none",
        "nan",
    }


# ---------------------------------------------------------------------------
# SchemaRegistry
# ---------------------------------------------------------------------------


class SchemaRegistry:
    """Collects extraction rules while producing typed Pydantic annotations.

    Create one registry per schema and pass it to
    :class:`~llmner.extractor.NERExtractor` via ``rules_registry``::

        registry = SchemaRegistry()

        GenderType = registry.categorical("gender", ["male", "female"], "...")
        AgeType    = registry.int_range("age", "Age in years.")

        extractor = NERExtractor(
            schema_class=MySchema,
            ...,
            rules_registry=registry.rules,
        )

    The registry is intentionally **not** shared between schemas; create a
    fresh instance for each.
    """

    def __init__(self) -> None:
        self._rules: list[str] = []

    @property
    def rules(self) -> list[str]:
        """Extraction rules in insertion order (one per registered type)."""
        return list(self._rules)

    def _register(self, rule: str) -> None:
        self._rules.append(rule)

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

    def categorical(
        self,
        name: str,
        options: list[str] | dict[str, str],
        instruction: str,
        *,
        replacements: dict[str, str] | None = None,
        allow_multiple: bool = False,
    ) -> Any:
        """Annotated type for categorical (enum-like) fields.

        Args:
            name:           Label used in the prompt rule.
            options:        Allowed values.  Use a ``dict`` to attach a
                            description to each value shown in the prompt.
            instruction:    Extraction instruction sent to the LLM.
            replacements:   Optional ``{old: new}`` string replacements applied
                            to the raw value before matching against *options*.
            allow_multiple: When ``True`` every matching option is collected
                            and joined with ``", "``.

        Returns:
            ``Annotated[EvidenceField, BeforeValidator(...)]``
        """
        _replacements = replacements or {}
        options_list: list[str] = (
            list(options.keys()) if isinstance(options, dict) else options
        )

        def validate(v: Any, info: ValidationInfo) -> EvidenceField:
            raw, evidence = _extract_ev(v)
            raw = _coerce_null(raw)
            evidence = _validate_evidence(evidence, info)

            if raw is None:
                return EvidenceField(value=None, evidence=evidence)

            s = str(raw).strip().lower()
            for old, new in _replacements.items():
                s = s.replace(old, new)

            if allow_multiple:
                found = [opt for opt in options_list if opt in s]
                return EvidenceField(
                    value=", ".join(found) if found else None,
                    evidence=evidence,
                )

            return EvidenceField(
                value=s if s in options_list else None,
                evidence=evidence,
            )

        if isinstance(options, dict):
            preview_dict = list(options.items())[:15]
            allowed_str = ", ".join(f"'{k}' ({v})" for k, v in preview_dict)
        else:
            preview_list = options_list[:15]
            allowed_str = ", ".join(preview_list)
        suffix = "..." if len(options_list) > 15 else ""
        self._register(
            f"{name.upper()}: {instruction}. Allowed values: {allowed_str}{suffix}"
        )

        return Annotated[EvidenceField, BeforeValidator(validate)]

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

    def int_range(self, name: str, instruction: str) -> Any:
        """Annotated type for integer or ``MIN-MAX`` range fields.

        Handles common LLM quirks such as metric notation (``"1.75 m"``
        → ``"175"`` cm) and comma/dash separators in ranges.

        Args:
            name:        Label used in the prompt rule.
            instruction: Extraction instruction sent to the LLM.

        Returns:
            ``Annotated[EvidenceField, BeforeValidator(...)]``
        """

        def validate(v: Any, info: ValidationInfo) -> EvidenceField:
            raw, evidence = _extract_ev(v)
            raw = _coerce_null(raw)
            evidence = _validate_evidence(evidence, info)

            if raw is None:
                return EvidenceField(value=None, evidence=evidence)

            s = str(raw).strip().lower().replace("cm", "").replace(", null", "").strip()
            # Convert metre notation "1.75" → "175"
            if len(s) >= 4 and re.match(r"^1[,.]\d{2}", s):
                s = s[0] + s[2:4]

            match = re.search(r"(\d+)(?:\s*[,y\-]\s*(\d+))?", s)
            if match:
                val = (
                    f"{match.group(1)}-{match.group(2)}"
                    if match.group(2)
                    else match.group(1)
                )
                return EvidenceField(value=val, evidence=evidence)

            return EvidenceField(value=None, evidence=evidence)

        self._register(f"{name.upper()}: {instruction}")
        return Annotated[EvidenceField, BeforeValidator(validate)]

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

    def generic(self, name: str, instruction: str, python_type: Any = str) -> Any:
        """Annotated type for free-text or weakly-typed fields.

        Args:
            name:        Label used in the prompt rule.
            instruction: Extraction instruction sent to the LLM.
            python_type: Callable used to coerce the raw string value
                         (default: ``str``).

        Returns:
            ``Annotated[EvidenceField, BeforeValidator(...)]``
        """

        def validate(v: Any, info: ValidationInfo) -> EvidenceField:
            raw, evidence = _extract_ev(v)
            raw = _coerce_null(raw)
            evidence = _validate_evidence(evidence, info)

            if raw is None:
                return EvidenceField(value=None, evidence=evidence)

            try:
                coerced = python_type(raw)
                return EvidenceField(value=str(coerced), evidence=evidence)
            except Exception:
                return EvidenceField(value=None, evidence=evidence)

        self._register(f"{name.upper()}: {instruction}")
        return Annotated[EvidenceField, BeforeValidator(validate)]

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

    def datetime_format(self, name: str, instruction: str) -> Any:
        """Annotated type that normalises dates to ``YYYY-MM-DD HH:MM:SS``.

        Accepts a range of common formats and converts them all to the
        canonical ISO-like representation for consistency.

        Args:
            name:        Label used in the prompt rule.
            instruction: Extraction instruction sent to the LLM.

        Returns:
            ``Annotated[EvidenceField, BeforeValidator(...)]``
        """
        _FORMATS = (
            "%Y-%m-%d %H:%M:%S",
            "%Y-%m-%d %H:%M",
            "%Y-%m-%d",
            "%d/%m/%Y %H:%M:%S",
            "%d/%m/%Y",
        )

        def validate(v: Any, info: ValidationInfo) -> EvidenceField:
            raw, evidence = _extract_ev(v)
            raw = _coerce_null(raw)
            evidence = _validate_evidence(evidence, info)

            if raw is None:
                return EvidenceField(value=None, evidence=evidence)

            s = str(raw).strip()
            for fmt in _FORMATS:
                try:
                    dt = datetime.strptime(s, fmt)
                    return EvidenceField(
                        value=dt.strftime("%Y-%m-%d %H:%M:%S"),
                        evidence=evidence,
                    )
                except ValueError:
                    continue

            return EvidenceField(value=None, evidence=evidence)

        self._register(f"{name.upper()}: {instruction}")
        return Annotated[EvidenceField, BeforeValidator(validate)]
