"""Named‑Entity Recognition (NER) predictive task."""

from __future__ import annotations

from collections.abc import Callable, Iterable, Sequence
from pathlib import Path
from typing import Any, Literal, override

import datasets
import dspy
import pydantic

from sieves.data import Doc
from sieves.model_wrappers import ModelType
from sieves.model_wrappers.types import ModelSettings
from sieves.serialization import Config
from sieves.tasks.distillation.types import DistillationFramework
from sieves.tasks.predictive.core import PredictiveTask
from sieves.tasks.predictive.gliner_bridge import GliNERBridge
from sieves.tasks.predictive.ner.bridges import DSPyNER, PydanticNER
from sieves.tasks.predictive.schemas.ner import (
    FewshotExample,
    TaskModel,
    TaskPromptSignature,
    TaskResult,
)

_TaskBridge = DSPyNER | GliNERBridge | PydanticNER


class NER(PredictiveTask[TaskPromptSignature, TaskResult, _TaskBridge]):
    """Extract named entities from text using various model wrappers."""

    def __init__(
        self,
        entities: Sequence[str] | dict[str, str] | None = None,
        model: TaskModel | None = None,
        task_id: str | None = None,
        include_meta: bool = True,
        batch_size: int = -1,
        prompt_instructions: str | None = None,
        fewshot_examples: Sequence[FewshotExample] = (),
        model_settings: ModelSettings = ModelSettings(),
        condition: Callable[[Doc], bool] | None = None,
    ) -> None:
        """Initialize NER task.

        :param entities: Entity types to extract. Supports two formats:
            - List format: `["PERSON", "LOCATION", "ORGANIZATION"]`
            - Dict format: `{"PERSON": "Names of people", "LOCATION": "Geographic locations"}`
            The dict format allows you to provide descriptions that help the model better identify entities.
            When using GliNER, descriptions are passed directly to the model for improved recognition.
            If None, defaults to `["PERSON", "LOCATION", "ORGANIZATION"]`.
        :param model: Model to use.
        :param task_id: Task ID.
        :param include_meta: Whether to include meta information generated by the task.
        :param batch_size: Batch size to use for inference. Use -1 to process all documents at once.
        :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
        :param fewshot_examples: Few-shot examples.
        :param model_settings: Settings for structured generation.
        :param condition: Optional callable that determines whether to process each document.
        """
        if entities is None:
            entities = ["PERSON", "LOCATION", "ORGANIZATION"]

        if isinstance(entities, dict):
            self._entities = list(entities.keys())
            self._entity_descriptions = entities
        else:
            self._entities = list(entities)
            self._entity_descriptions = {}

        self._entities_param = entities

        super().__init__(
            model=model,
            task_id=task_id,
            include_meta=include_meta,
            batch_size=batch_size,
            overwrite=False,
            prompt_instructions=prompt_instructions,
            fewshot_examples=fewshot_examples,
            model_settings=model_settings,
            condition=condition,
        )

    @property
    @override
    def fewshot_example_type(self) -> type[FewshotExample]:
        return FewshotExample

    @property
    @override
    def prompt_signature(self) -> type[pydantic.BaseModel]:
        """Return the unified Pydantic prompt signature for this task.

        :return: Unified Pydantic prompt signature.
        """
        # Create a dynamic entity model with Literal for the entity types.
        EntityTypes = Literal[*(tuple(self._entities))] if self._entities else str  # type: ignore[valid-type]

        DynamicEntity = pydantic.create_model(
            "NEREntity",
            __doc__="Extracted named entity with its context and type.",
            text=(str, pydantic.Field(description="The specific text segment identified as an entity.")),
            context=(str, pydantic.Field(description="The surrounding text providing context for the entity.")),
            entity_type=(
                EntityTypes,
                pydantic.Field(description="The category or type of the entity (e.g., PERSON, ORGANIZATION)."),
            ),
            score=(
                float | None,
                pydantic.Field(
                    default=None,
                    description="Provide a confidence score for the entity identification, between 0 and 1.",
                ),
            ),
            __base__=pydantic.BaseModel,
        )

        return pydantic.create_model(
            "NEROutput",
            __doc__="Result of named-entity recognition. Contains a list of extracted entities.",
            entities=(
                list[DynamicEntity],  # type: ignore[valid-type]
                pydantic.Field(..., description="List of extracted named entities."),
            ),
        )

    @property
    @override
    def metric(self) -> str:
        return "F1"

    @override
    def _compute_metrics(self, truths: list[Any], preds: list[Any], judge: dspy.LM | None = None) -> dict[str, float]:
        """Compute corpus-level metrics.

        :param truths: List of ground truths.
        :param preds: List of predictions.
        :param judge: Optional DSPy LM instance to use as judge for generative tasks.
        :return: Dictionary of metrics.
        """
        tp = 0
        fp = 0
        fn = 0

        for gold, pred in zip(truths, preds):
            if gold is None:
                true_entities = set()
            else:
                # Convert to DSPy format first to reuse logic or do it directly.
                assert isinstance(gold, TaskResult)
                true_entities = {(e.text, e.entity_type) for e in gold.entities}

            if pred is None:
                pred_entities = set()
            else:
                if hasattr(pred, "entities"):
                    pred_entities = {(e.text, e.entity_type) for e in pred.entities}
                elif isinstance(pred, dict) and "entities" in pred:
                    pred_entities = {(e["text"], e["entity_type"]) for e in pred.get("entities", [])}
                else:
                    pred_entities = set()

            tp += len(true_entities & pred_entities)
            fp += len(pred_entities - true_entities)
            fn += len(true_entities - pred_entities)

        precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
        recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0
        f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0

        return {self.metric: f1}

    @override
    def _init_bridge(self, model_type: ModelType) -> _TaskBridge:
        if model_type == ModelType.gliner:
            from sieves.model_wrappers import gliner_

            return GliNERBridge(
                task_id=self._task_id,
                prompt_instructions=self._custom_prompt_instructions,
                prompt_signature=self.prompt_signature,
                model_settings=self._model_settings,
                inference_mode=gliner_.InferenceMode.entities,
            )

        bridge_types = {
            ModelType.langchain: PydanticNER,
            ModelType.outlines: PydanticNER,
            ModelType.dspy: DSPyNER,
        }

        try:
            bridge_class = bridge_types[model_type]
            result = bridge_class(
                task_id=self._task_id,
                prompt_instructions=self._custom_prompt_instructions,
                entities=self._entities_param,
                model_settings=self._model_settings,
                prompt_signature=self.prompt_signature,
                model_type=model_type,
                fewshot_examples=self._fewshot_examples,
            )
            return result  # type: ignore[return-value]
        except KeyError as err:
            raise KeyError(f"Model type {model_type} is not supported by {self.__class__.__name__}.") from err

    @staticmethod
    @override
    def supports() -> set[ModelType]:
        return {
            ModelType.langchain,
            ModelType.dspy,
            ModelType.outlines,
            ModelType.gliner,
        }

    @override
    def _validate_fewshot_examples(self) -> None:
        for fs_example in self._fewshot_examples or []:
            assert isinstance(fs_example, FewshotExample)
            for entity in fs_example.entities:
                if entity.entity_type not in self._entities:
                    raise ValueError(f"Entity {entity.entity_type} not in {self._entities}.")

    @override
    @property
    def _state(self) -> dict[str, Any]:
        # Store entities as dict if descriptions exist, else as list
        entities_state = self._entity_descriptions if self._entity_descriptions else self._entities
        return {
            **super()._state,
            "entities": entities_state,
        }

    @override
    def to_hf_dataset(self, docs: Iterable[Doc], threshold: float | None = None) -> datasets.Dataset:
        # Define metadata and features for the dataset
        features = datasets.Features(
            {
                "text": datasets.Value("string"),
                "entities": datasets.Sequence(
                    datasets.Features(
                        {
                            "text": datasets.Value("string"),
                            "start": datasets.Value("int32"),
                            "end": datasets.Value("int32"),
                            "entity_type": datasets.Value("string"),
                            "score": datasets.Value("float32"),
                        }
                    )
                ),
            }
        )

        info = datasets.DatasetInfo(
            description=f"Named Entity Recognition dataset with entity types {self._entities}. Generated with sieves "
            f"v{Config.get_version()}.",
            features=features,
        )

        # Fetch data used for generating dataset
        try:
            data: list[tuple[str, list[dict[str, Any]]]] = []
            for doc in docs:
                if self._task_id not in doc.results:
                    raise KeyError(f"Document does not have results for task ID {self._task_id}")

                # Get the entities from the document results
                result = doc.results[self._task_id].entities
                entities: list[dict[str, Any]] = []

                # List format (could be list of dictionaries or other entities)
                for entity in result:
                    assert hasattr(entity, "text")
                    assert hasattr(entity, "start")
                    assert hasattr(entity, "end")
                    assert hasattr(entity, "entity_type")

                    entities.append(
                        {
                            "text": entity.text,
                            "start": entity.start,
                            "end": entity.end,
                            "entity_type": entity.entity_type,
                            "score": getattr(entity, "score", None),
                        }
                    )

                data.append((doc.text or "", entities))

        except KeyError as err:
            raise KeyError(f"Not all documents have results for this task with ID {self._task_id}") from err

        def generate_data() -> Iterable[dict[str, Any]]:
            """Yield results as dicts.

            :return: Results as dicts.
            """
            for text, entities in data:
                yield {"text": text, "entities": entities}

        # Create dataset
        return datasets.Dataset.from_generator(generate_data, features=features, info=info)

    @override
    def distill(
        self,
        base_model_id: str,
        framework: DistillationFramework,
        data: datasets.Dataset | Sequence[Doc],
        output_path: Path | str,
        val_frac: float,
        init_kwargs: dict[str, Any] | None = None,
        train_kwargs: dict[str, Any] | None = None,
        seed: int | None = None,
    ) -> None:
        raise NotImplementedError

    @override
    def _evaluate_dspy_example(self, truth: dspy.Example, pred: dspy.Prediction, trace: Any, model: dspy.LM) -> float:
        # Compute entity-level F1 score based on (text, entity_type) pairs
        true_entities = {(e["text"], e["entity_type"]) for e in truth["entities"]}
        pred_entities = {(e["text"], e["entity_type"]) for e in pred.get("entities", [])}

        if not true_entities:
            return 1.0 if not pred_entities else 0.0

        precision = len(true_entities & pred_entities) / len(pred_entities) if pred_entities else 0
        recall = len(true_entities & pred_entities) / len(true_entities)
        return 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0
