"""Core base model for LLM-powered Named Entity Recognition.

All user-defined extraction schemas must subclass :class:`NERBaseModel`.
It provides four utilities built entirely on Pydantic reflection, so they
work on any field layout without requiring manual maintenance:

* :meth:`~NERBaseModel.prompt_schema`  – JSON skeleton for the LLM prompt.
* :meth:`~NERBaseModel.has_missing_fields` – detect incomplete extractions.
* :meth:`~NERBaseModel.merge` – combine two partial results, keeping non-null
  values.
* :meth:`~NERBaseModel.safe_parse` – tolerantly validate LLM output,
  converting individual bad values to ``None`` instead of raising.
"""

from __future__ import annotations

import json
import types
from typing import Annotated, Any, Union, get_args, get_origin

from pydantic import BaseModel, ConfigDict

__all__ = ["NERBaseModel"]


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _is_union_origin(origin: Any) -> bool:
    """Return ``True`` for both ``typing.Union`` and ``types.UnionType``.

    Python 3.10 introduced ``X | Y`` syntax which creates a
    ``types.UnionType`` instead of ``typing.Union``.  Both must be handled.
    """
    return origin is Union or origin is types.UnionType


def _unwrap_optional(annotation: Any) -> Any:
    """Return the inner type of ``Optional[X]`` / ``X | None``.

    Handles both ``typing.Union[X, None]`` and the Python 3.10+
    ``X | None`` (``types.UnionType``) syntax.
    Any other annotation is returned unchanged.
    """
    origin = get_origin(annotation)
    if _is_union_origin(origin):
        non_none = [a for a in get_args(annotation) if a is not type(None)]
        if len(non_none) == 1:
            return non_none[0]
    return annotation


def _base_type(annotation: Any) -> Any:
    """Return the first argument of ``Annotated[X, ...]``, else *annotation*.

    This unwraps factory-produced types such as
    ``Annotated[EvidenceField, BeforeValidator(...)]`` to reveal the
    underlying Pydantic model class.
    """
    if get_origin(annotation) is Annotated:
        args = get_args(annotation)
        if args:
            return args[0]
    return annotation


def _field_schema_entry() -> dict[str, str]:
    return {
        "value": "string or null",
        "evidence": (
            "string or null "
            "(shortest verbatim quote from the text that unequivocally "
            "justifies this specific value — ideally 1-5 words)"
        ),
    }


# ---------------------------------------------------------------------------
# Base model
# ---------------------------------------------------------------------------


class NERBaseModel(BaseModel):
    """Base class for all LLM-NER extraction schemas.

    Subclass this and declare fields using the type factories from
    :mod:`llmner.factories` (or plain Pydantic types).  Example::

        from llmner import NERBaseModel, SchemaRegistry

        registry = SchemaRegistry()

        class PersonSchema(NERBaseModel):
            name: registry.generic("name", "Full name of the person") = None
            age: registry.int_range("age", "Age in years") = None
    """

    model_config = ConfigDict(populate_by_name=True)

    # ------------------------------------------------------------------
    # Schema generation
    # ------------------------------------------------------------------

    @classmethod
    def prompt_schema(cls) -> str:
        """Return a JSON string describing the expected LLM output.

        The structure mirrors the model hierarchy so the LLM knows exactly
        which keys to produce and in what ``{value, evidence}`` format.
        """
        schema: dict[str, Any] = {}

        for name, field_info in cls.model_fields.items():
            key = field_info.alias or name
            annotation = _unwrap_optional(field_info.annotation)
            origin = get_origin(annotation)

            if origin is list:
                args = get_args(annotation)
                if args:
                    item_model = args[0]
                    if isinstance(item_model, type) and issubclass(
                        item_model, BaseModel
                    ):
                        schema[key] = [
                            {
                                f.alias or n: _field_schema_entry()
                                for n, f in item_model.model_fields.items()
                            }
                        ]
                        continue

            elif isinstance(annotation, type) and issubclass(annotation, BaseModel):
                schema[key] = {
                    f.alias or n: _field_schema_entry()
                    for n, f in annotation.model_fields.items()
                }
                continue

            schema[key] = _field_schema_entry()

        return json.dumps(schema, indent=2, ensure_ascii=False)

    # ------------------------------------------------------------------
    # Missing-field detection
    # ------------------------------------------------------------------

    def has_missing_fields(self) -> bool:
        """Return ``True`` if any nested extracted *value* is ``None``.

        Top-level fields that are ``None`` themselves are skipped – those
        represent entirely absent concepts and do not warrant a retry.
        Only *partially* populated sub-structures (e.g. an
        :class:`~llmner.factories.EvidenceField` whose ``value`` is ``None``)
        trigger ``True``.

        Note: only the ``"value"`` key inside ``EvidenceField``-like dicts
        is checked.  Missing ``"evidence"`` does **not** trigger a retry.
        """

        def _has_null_value(obj: Any) -> bool:
            if isinstance(obj, dict):
                if "value" in obj:
                    return obj["value"] is None
                return any(_has_null_value(v) for v in obj.values())
            if isinstance(obj, list):
                return any(_has_null_value(item) for item in obj)
            return False

        for name in self.__class__.model_fields:
            val = getattr(self, name)
            if val is None:
                continue
            if isinstance(val, BaseModel):
                if _has_null_value(val.model_dump()):
                    return True
            elif isinstance(val, list):
                for item in val:
                    if isinstance(item, BaseModel) and _has_null_value(
                        item.model_dump()
                    ):
                        return True
        return False

    # ------------------------------------------------------------------
    # Merging
    # ------------------------------------------------------------------

    @classmethod
    def merge(cls, first: NERBaseModel, second: NERBaseModel) -> NERBaseModel:
        """Return a new instance combining the two extractions.

        For every field whose value resolves to ``None`` or ``""`` in
        *first*, the corresponding value from *second* is used.  Nested
        models and lists of models are merged recursively.

        Args:
            first:  Primary extraction; its non-null values take precedence.
            second: Fallback extraction used to fill gaps in *first*.

        Returns:
            A new instance of the same class with gaps filled.
        """

        def _deep_merge(d1: dict[str, Any], d2: dict[str, Any]) -> dict[str, Any]:
            merged: dict[str, Any] = {}
            for k in d1.keys() | d2.keys():
                v1, v2 = d1.get(k), d2.get(k)
                if isinstance(v1, dict) and isinstance(v2, dict):
                    merged[k] = _deep_merge(v1, v2)
                else:
                    merged[k] = v1 if (v1 is not None and v1 != "") else v2
            return merged

        merged_args: dict[str, Any] = {}

        for name, field_info in cls.model_fields.items():
            val1 = getattr(first, name)
            val2 = getattr(second, name)
            annotation = _unwrap_optional(field_info.annotation)
            origin = get_origin(annotation)

            if origin is list:
                args = get_args(annotation)
                if args:
                    item_model = args[0]
                    if isinstance(item_model, type) and issubclass(
                        item_model, BaseModel
                    ):
                        l1: list[Any] = val1 or []
                        l2: list[Any] = val2 or []
                        merged_list: list[Any] = []
                        for i in range(max(len(l1), len(l2))):
                            i1 = l1[i] if i < len(l1) else None
                            i2 = l2[i] if i < len(l2) else None
                            if i1 is not None and i2 is not None:
                                merged_dict = _deep_merge(
                                    i1.model_dump(by_alias=True),
                                    i2.model_dump(by_alias=True),
                                )
                                merged_list.append(
                                    item_model.model_validate(merged_dict)
                                )
                            else:
                                merged_list.append(i1 or i2)
                        merged_args[name] = merged_list
                        continue

            if isinstance(annotation, type) and issubclass(annotation, BaseModel):
                d1 = val1.model_dump() if val1 is not None else {}
                d2 = val2.model_dump() if val2 is not None else {}
                merged_args[name] = annotation(**_deep_merge(d1, d2))
            else:
                # Handle Annotated[SomeBaseModel, ...] fields (e.g. EvidenceField
                # wrapped by a factory BeforeValidator).  A non-None val1 whose
                # inner value is null still needs to be merged with val2.
                base = _base_type(annotation)
                if isinstance(base, type) and issubclass(base, BaseModel):
                    d1 = val1.model_dump() if isinstance(val1, BaseModel) else {}
                    d2 = val2.model_dump() if isinstance(val2, BaseModel) else {}
                    merged_args[name] = base(**_deep_merge(d1, d2))
                else:
                    merged_args[name] = val1 if val1 is not None else val2

        return cls(**merged_args)

    # ------------------------------------------------------------------
    # Safe parsing
    # ------------------------------------------------------------------

    @classmethod
    def safe_parse(
        cls,
        data: dict[str, Any] | None,
        input_text: str = "",
    ) -> NERBaseModel | None:
        """Parse *data* tolerantly, isolating validation errors per field.

        Invalid enum values, malformed numbers, evidence not found in the
        source text, etc. are silently converted to ``None`` rather than
        raising a ``ValidationError``.

        Args:
            data:       Raw ``dict`` produced by the LLM (already JSON-decoded).
            input_text: Original source text passed as Pydantic context so
                        that evidence validators can verify quotes exist in it.

        Returns:
            Populated model instance, or ``None`` if *data* is ``None`` or
            not a ``dict``.
        """
        if not isinstance(data, dict):
            return None

        ctx = {"input_text": input_text}
        data_copy = dict(data)

        # ── Phase 1: validate list-of-model fields item by item ──────────
        for name, field_info in cls.model_fields.items():
            key = field_info.alias or name
            annotation = _unwrap_optional(field_info.annotation)
            origin = get_origin(annotation)

            raw = data_copy.get(key) or data_copy.get(name)

            if origin is list:
                args = get_args(annotation)
                if args:
                    item_model = args[0]
                    if isinstance(item_model, type) and issubclass(
                        item_model, BaseModel
                    ):
                        valid_items: list[Any] = []
                        if isinstance(raw, list):
                            for item_data in raw:
                                if isinstance(item_data, dict):
                                    try:
                                        obj = item_model.model_validate(
                                            item_data, context=ctx
                                        )
                                        valid_items.append(obj.model_dump())
                                    except Exception:
                                        safe: dict[str, Any] = {}
                                        for (
                                            f_name,
                                            f_info,
                                        ) in item_model.model_fields.items():
                                            c_key = f_info.alias or f_name
                                            i_val = item_data.get(
                                                c_key
                                            ) or item_data.get(f_name)
                                            try:
                                                tmp = item_model.model_validate(
                                                    {c_key: i_val},
                                                    context=ctx,
                                                )
                                                safe[f_name] = getattr(tmp, f_name)
                                            except Exception:
                                                safe[f_name] = None
                                        valid_items.append(safe)
                        data_copy[name] = valid_items
                        if key != name and key in data_copy:
                            del data_copy[key]

        # ── Phase 2: attempt full model validation ────────────────────────
        try:
            return cls.model_validate(data_copy, context=ctx)
        except Exception:
            pass

        # ── Phase 3: field-by-field fallback ─────────────────────────────
        safe_args: dict[str, Any] = {}

        for name, field_info in cls.model_fields.items():
            key = field_info.alias or name
            annotation = _unwrap_optional(field_info.annotation)
            origin = get_origin(annotation)

            raw = data_copy.get(key) or data_copy.get(name)

            if origin is list:
                args = get_args(annotation)
                if args:
                    item_model = args[0]
                    if isinstance(item_model, type) and issubclass(
                        item_model, BaseModel
                    ):
                        items = data_copy.get(name) or []
                        safe_args[name] = [
                            item_model.model_validate(s, context=ctx)
                            for s in items
                            if isinstance(s, dict)
                        ]
                        continue

            if isinstance(annotation, type) and issubclass(annotation, BaseModel):
                obj_args: dict[str, Any] = {}
                if isinstance(raw, dict):
                    for k, v in raw.items():
                        for f_name, f_info in annotation.model_fields.items():
                            c_key = f_info.alias or f_name
                            if c_key == k:
                                try:
                                    obj_args[f_name] = getattr(
                                        annotation.model_validate(
                                            {c_key: v}, context=ctx
                                        ),
                                        f_name,
                                    )
                                except Exception:
                                    obj_args[f_name] = None
                safe_args[name] = annotation(**obj_args)
            else:
                safe_args[name] = raw

        return cls(**safe_args)
