"""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.

Every factory accepts an optional ``fallback_parser`` callback.  When the
LLM returns a null *value* but a non-null *evidence* quote, the callback
receives the evidence string and may return a recovered raw value that is
then fed through the normal validation pipeline.  This makes extraction
more robust without changing the LLM prompt or the validation rules.

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 collections.abc import Callable
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


_MIN_EVIDENCE_LEN = 3  # Minimum character length for a value to serve as evidence


def _descore(text: str) -> str:
    """Replace underscores with spaces for fuzzy matching against source text.

    Many categorical schemas use underscores as word separators
    (e.g. ``"physical_assault"``), but the original text normally contains
    spaces (``"physical assault"``).  This helper produces the space-separated
    form so that evidence resolution can still anchor the value.
    """
    return text.replace("_", " ")


def _resolve_evidence(
    value: str | None,
    evidence: str | None,
    info: ValidationInfo,
    *,
    raw_value: str | None = None,
) -> str | None:
    """Return the best evidence anchor for *value* within the source text.

    Priority rules applied in order:

    1. **Full value match**: if the complete *value* string is found verbatim
       in the source text (and is at least ``_MIN_EVIDENCE_LEN`` characters
       long), it is returned as the evidence.  This supersedes any model-
       provided *evidence* because the canonical value is the most precise and
       self-consistent anchor possible.

    1a. **Full value match (descored)**: if the value contains underscores
        and the underscore-to-space form appears in the text, that form is
        returned as evidence.  This covers categorical values like
        ``"physical_assault"`` whose natural text form is ``"physical assault"``.

    1b. **Full raw value match**: if the normalised value is *not* found but
        the pre-transformation form (e.g. ``"1.75"`` before metre→cm
        conversion) appears in the text, the raw value is returned as
        evidence.

    1c. **Full raw value match (descored)**: same as 1a but applied to the
        raw (pre-transformation) value.

    2. **Existing model evidence**: if neither value nor raw value is found
       but the model already provided a validated *evidence* quote, that
       quote is kept unchanged.

    3. **Partial value match** (null-evidence fallback): only when *evidence*
       is ``None`` and the full value is not found, the longest token-prefix
       of *value* that appears in the source text is used as a fallback
       evidence.  The descored (underscore→space) form is also tried.

    3b. **Partial raw value prefix**: same as rule 3 but applied to the raw
        (pre-transformation) value.

    4. ``None`` — no usable evidence is available.

    Args:
        value:     The normalised (post-transformation) value.
        evidence:  Model-provided evidence quote (already validated).
        info:      Pydantic ``ValidationInfo`` carrying ``input_text``.
        raw_value: Optional pre-transformation form of the value.  Factories
                   that transform the LLM output (e.g. ``int_range`` metre→cm,
                   ``datetime_format`` format normalisation) pass the original
                   string so that evidence can still be anchored in the source
                   text.

    When no ``input_text`` is available in the Pydantic context the function
    returns *evidence* unchanged.
    """
    if not value:
        return evidence
    ctx = info.context or {}
    input_text: str = ctx.get("input_text", "")
    if not input_text:
        return evidence

    norm_text = _normalize_ws(input_text).lower()
    norm_value = _normalize_ws(value)
    norm_raw = _normalize_ws(str(raw_value)) if raw_value is not None else None

    # Rule 1 — full value present in text: always prefer value as evidence
    if len(norm_value) >= _MIN_EVIDENCE_LEN and norm_value.lower() in norm_text:
        return norm_value

    # Rule 1a — descored value (underscores → spaces) present in text
    descored_value = _descore(norm_value)
    if (
        "_" in norm_value
        and len(descored_value) >= _MIN_EVIDENCE_LEN
        and descored_value.lower() in norm_text
    ):
        return descored_value

    # Rule 1b — full raw value present in text (pre-transformation form)
    if (
        norm_raw is not None
        and len(norm_raw) >= _MIN_EVIDENCE_LEN
        and norm_raw.lower() in norm_text
    ):
        return norm_raw

    # Rule 1c — descored raw value present in text
    if norm_raw is not None and "_" in norm_raw:
        descored_raw = _descore(norm_raw)
        if (
            len(descored_raw) >= _MIN_EVIDENCE_LEN
            and descored_raw.lower() in norm_text
        ):
            return descored_raw

    # Rule 2 — full value absent but model evidence exists: keep it
    if evidence is not None:
        return evidence

    # Rule 3 — no evidence at all: try a partial token-prefix as fallback
    # Try both the original value and the descored form
    for candidate in (norm_value, descored_value):
        tokens = candidate.split()
        for n in range(len(tokens) - 1, 0, -1):
            prefix = " ".join(tokens[:n])
            if len(prefix) >= _MIN_EVIDENCE_LEN and prefix.lower() in norm_text:
                return prefix

    # Rule 3b — partial token-prefix of raw value (and descored form)
    if norm_raw is not None:
        raw_candidates = [norm_raw]
        if "_" in norm_raw:
            raw_candidates.append(_descore(norm_raw))
        for candidate in raw_candidates:
            raw_tokens = candidate.split()
            for n in range(len(raw_tokens) - 1, 0, -1):
                prefix = " ".join(raw_tokens[:n])
                if len(prefix) >= _MIN_EVIDENCE_LEN and prefix.lower() in norm_text:
                    return prefix

    return None


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


def _try_fallback(
    raw: Any,
    evidence: str | None,
    fallback_parser: Callable[[str], str | None] | None,
) -> Any:
    """Attempt to recover a raw value from the evidence string.

    Called when the LLM returned a null value but provided an evidence
    quote.  If a *fallback_parser* is configured and returns a non-``None``
    result, that result replaces the original null *raw* and is fed back
    into the factory's normal validation logic.

    Returns the (possibly recovered) *raw* value.
    """
    if raw is not None or not evidence or fallback_parser is None:
        return raw
    return fallback_parser(evidence)


def _apply_evidence_required(
    field: EvidenceField,
    required: bool,
    info: ValidationInfo,
) -> EvidenceField:
    """Nullify *value* when evidence is required but was not found.

    When ``required`` is ``True`` and the field has a non-null *value* but
    a null *evidence*, the value cannot be trusted because there is no
    supporting quote from the source text.  In that case both *value* and
    *evidence* are set to ``None``.

    Only enforced when ``input_text`` is available in the Pydantic context —
    without source text, evidence verification was never attempted, so
    discarding would be incorrect.
    """
    if not required:
        return field
    ctx = info.context or {}
    if not ctx.get("input_text", ""):
        return field
    if field.value is not None and field.evidence is None:
        return EvidenceField(value=None, evidence=None)
    return field


# ---------------------------------------------------------------------------
# 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,
        fallback_parser: Callable[[str], str | None] | None = None,
        evidence_required: 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 ``", "``.
            fallback_parser: Optional callback ``(evidence: str) -> str | None``
                             invoked when the LLM returns a null value but
                             provides an evidence quote.  Should return a raw
                             value candidate or ``None``.
            evidence_required: When ``True``, a non-null *value* is set back
                             to ``None`` if no supporting *evidence* quote
                             could be found or verified in the source text.

        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)
            raw = _try_fallback(raw, evidence, fallback_parser)

            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]
                value = ", ".join(found) if found else None
                result = EvidenceField(
                    value=value,
                    evidence=_resolve_evidence(value, evidence, info),
                )
                return _apply_evidence_required(result, evidence_required, info)

            value = s if s in options_list else None
            result = EvidenceField(
                value=value,
                evidence=_resolve_evidence(value, evidence, info),
            )
            return _apply_evidence_required(result, evidence_required, info)

        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,
        *,
        fallback_parser: Callable[[str], str | None] | None = None,
        evidence_required: bool = False,
    ) -> 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.
            fallback_parser: Optional callback ``(evidence: str) -> str | None``
                             invoked when the LLM returns a null value but
                             provides an evidence quote.  Should return a raw
                             value candidate or ``None``.
            evidence_required: When ``True``, a non-null *value* is set back
                             to ``None`` if no supporting *evidence* quote
                             could be found or verified in the source text.

        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)
            raw = _try_fallback(raw, evidence, fallback_parser)

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

            raw_str = str(raw).strip()
            s = raw_str.lower().replace("cm", "").replace(", null", "").strip()
            # Convert metre notation to centimetres: range first ("1,65-1,80" → "165-180"),
            # then single value ("1.75" → "175").
            _m_range = re.match(r"^(1[,.]\d{2})\s*(?:[-y]|a)\s*(1[,.]\d{2})", s)
            if _m_range:
                a, b = _m_range.group(1), _m_range.group(2)
                s = (a[0] + a[2:4]) + "-" + (b[0] + b[2:4])
            elif re.fullmatch(r"1[,.]\d{2}\s*m?\w*", 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)
                )
                result = EvidenceField(
                    value=val,
                    evidence=_resolve_evidence(
                        val, evidence, info, raw_value=raw_str
                    ),
                )
                return _apply_evidence_required(result, evidence_required, info)

            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,
        *,
        fallback_parser: Callable[[str], str | None] | None = None,
        evidence_required: bool = False,
    ) -> 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``).
            fallback_parser: Optional callback ``(evidence: str) -> str | None``
                             invoked when the LLM returns a null value but
                             provides an evidence quote.  Should return a raw
                             value candidate or ``None``.
            evidence_required: When ``True``, a non-null *value* is set back
                             to ``None`` if no supporting *evidence* quote
                             could be found or verified in the source text.

        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)
            raw = _try_fallback(raw, evidence, fallback_parser)

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

            try:
                coerced = python_type(raw)
                value = str(coerced)
                result = EvidenceField(
                    value=value, evidence=_resolve_evidence(value, evidence, info)
                )
                return _apply_evidence_required(result, evidence_required, info)
            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,
        *,
        extra_formats: tuple[str, ...] = (),
        fallback_parser: Callable[[str], str | None] | None = None,
        evidence_required: bool = False,
    ) -> 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.
            extra_formats:   Additional :func:`~datetime.datetime.strptime`
                             format strings appended after the built-in ones.
            fallback_parser: Optional callback ``(evidence: str) -> str | None``
                             invoked when the LLM returns a null value but
                             provides an evidence quote.  Should return a raw
                             value candidate or ``None``.
            evidence_required: When ``True``, a non-null *value* is set back
                             to ``None`` if no supporting *evidence* quote
                             could be found or verified in the source text.

        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",
        ) + tuple(extra_formats)

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

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

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

            return EvidenceField(value=None, evidence=evidence)

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