"""Information extraction."""

from __future__ import annotations

import types
import typing
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, gliner_
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.information_extraction.bridges import (
    DSPyInformationExtraction,
    PydanticInformationExtraction,
)
from sieves.tasks.predictive.schemas.information_extraction import (
    FewshotExampleMulti,
    FewshotExampleSingle,
    TaskModel,
    TaskPromptSignature,
    TaskResult,
)
from sieves.tasks.utils import PydanticToHFDatasets

_TaskBridge = GliNERBridge | DSPyInformationExtraction | PydanticInformationExtraction

FewshotExample = FewshotExampleMulti | FewshotExampleSingle


def _wrap_with_score(entity_type: type[pydantic.BaseModel]) -> type[pydantic.BaseModel]:
    """Create a subclass of the provided entity type that includes a score field.

    :param entity_type: The Pydantic model class to wrap.
    :return: A new Pydantic model class with a score field.
    """
    if hasattr(entity_type, "score"):
        return entity_type

    class ScoredEntity(entity_type):  # type: ignore[valid-type, misc]
        """Extracted entity with an additional confidence score."""

        score: float | None = pydantic.Field(
            default=None, description="Provide a confidence score for the entity extraction, between 0 and 1."
        )

    ScoredEntity.__name__ = entity_type.__name__

    return ScoredEntity


class InformationExtraction(PredictiveTask[TaskPromptSignature, TaskResult, _TaskBridge]):
    """Information extraction task."""

    def __init__(
        self,
        entity_type: type[pydantic.BaseModel],
        model: TaskModel,
        task_id: str | None = None,
        include_meta: bool = True,
        batch_size: int = -1,
        prompt_instructions: str | None = None,
        fewshot_examples: Sequence[FewshotExample] = (),
        mode: Literal["multi", "single"] = "multi",
        model_settings: ModelSettings = ModelSettings(),
        condition: Callable[[Doc], bool] | None = None,
    ) -> None:
        """Initialize new PredictiveTask.

        :param entity_type: Pydantic model class representing the object type to extract.
        :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 mode: Extraction mode. If "multi", all occurrences of the entity are extracted. If "single", exactly one
            (or no) entity is extracted.
        :param model_settings: Settings for structured generation.
        :param condition: Optional callable that determines whether to process each document.
        """
        self._entity_type = entity_type
        self._mode = mode
        self._scored_entity_type = _wrap_with_score(entity_type)

        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,
        )

        if (
            isinstance(entity_type, type)
            and issubclass(entity_type, pydantic.BaseModel)
            and not entity_type.model_config.get("frozen", False)
        ):
            raise ValueError(
                f"Entity type provided to task {self._task_id} isn't frozen, which means that entities can't "
                f"be deduplicated and compared. Modify entity_type to be frozen=True."
            )

    @property
    @override
    def fewshot_example_type(self) -> type[FewshotExample]:
        """Return few-shot example type.

        :return: Few-shot example type.
        """
        if self._mode == "multi":
            return FewshotExampleMulti
        return FewshotExampleSingle

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

        :return: Unified Pydantic prompt signature.
        """
        scored_type = self._scored_entity_type

        if self._mode == "multi":
            return pydantic.create_model(
                f"{scored_type.__name__}Multi",
                __doc__=f"Result of multi-entity extraction for {scored_type.__name__}.",
                entities=(
                    list[scored_type],  # type: ignore[valid-type]
                    pydantic.Field(..., description=f"List of extracted {scored_type.__name__} entities."),
                ),
            )
        return pydantic.create_model(
            f"{scored_type.__name__}Single",
            __doc__=f"Result of single-entity extraction for {scored_type.__name__}.",
            entity=(
                scored_type | None,  # type: ignore[valid-type]
                pydantic.Field(..., description=f"The extracted {scored_type.__name__} entity."),
            ),
        )

    @property
    @override
    def metric(self) -> str:
        return "F1" if self._mode == "multi" else "Accuracy"

    @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.
        """
        # Single mode.
        if self._mode == "single":
            correct = 0
            total = 0
            for gold, pred in zip(truths, preds):
                # If gold is None, we assume it's a negative example (no entity).
                # If pred is None, it predicted no entity.
                # If both None -> Correct.
                if gold is not None:
                    assert isinstance(gold, TaskResult)
                    gold_entity = gold.entity
                else:
                    gold_entity = None

                if pred is not None:
                    assert isinstance(pred, TaskResult)
                    pred_entity = pred.entity
                else:
                    pred_entity = None

                if gold_entity == pred_entity:
                    correct += 1
                total += 1

            return {self.metric: correct / total if total > 0 else 0.0}

        # Multi mode.
        tp = 0
        fp = 0
        fn = 0

        def entity_to_tuple(entity: Any) -> tuple:
            # Handle Pydantic model
            d = entity.model_dump()

            # Remove score if present
            if "score" in d:
                del d["score"]

            items = sorted(d.items())
            return tuple((k, v if not (isinstance(v, list) or isinstance(v, dict)) else str(v)) for k, v in items)

        for gold, pred in zip(truths, preds):
            if gold is not None:
                assert isinstance(gold, TaskResult)
                true_entities = {entity_to_tuple(e) for e in gold.entities}
            else:
                true_entities = set()

            if pred is not None:
                assert isinstance(pred, TaskResult)
                pred_entities = {entity_to_tuple(e) for e in pred.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:
        """Initialize bridge.

        :param model_type: Type of model to initialize bridge for.
        :return _TaskBridge: ModelWrapper task bridge.
        :raises ValueError: If model type is not supported.
        :raises TypeError: On entity type and model type mismatch.
        """
        if model_type == ModelType.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.structure,
                mode=self._mode,
            )

        bridge_types: dict[ModelType, type[_TaskBridge]] = {
            ModelType.dspy: DSPyInformationExtraction,
            ModelType.langchain: PydanticInformationExtraction,
            ModelType.outlines: PydanticInformationExtraction,
        }

        try:
            bridge = bridge_types[model_type](
                task_id=self._task_id,
                prompt_instructions=self._custom_prompt_instructions,
                entity_type=self._scored_entity_type,
                model_settings=self._model_settings,
                mode=self._mode,
                prompt_signature=self.prompt_signature,
                model_type=model_type,
                fewshot_examples=self._fewshot_examples,
            )

        except KeyError as err:
            raise KeyError(f"Model type {model_type} is not supported by {self.__class__.__name__}.") from err

        return bridge

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

    @override
    @property
    def _state(self) -> dict[str, Any]:
        return {
            **super()._state,
            "entity_type": self._entity_type,
            "mode": self._mode,
        }

    @override
    def _validate_fewshot_examples(self) -> None:
        example_type_error_text = "Fewshot example type mismatch: mode = {mode} requires {example_type}."

        for i, fs_example in enumerate(self._fewshot_examples or []):
            if self._mode == "multi":
                assert isinstance(fs_example, FewshotExampleMulti), TypeError(
                    example_type_error_text.format(example_type=FewshotExampleMulti, mode=self._mode)
                )
            else:
                assert isinstance(fs_example, FewshotExampleSingle), TypeError(
                    example_type_error_text.format(example_type=FewshotExampleSingle, mode=self._mode)
                )

    @override
    def to_hf_dataset(self, docs: Iterable[Doc], threshold: float | None = None) -> datasets.Dataset:
        # Use the scored entity type for the dataset schema.
        # We need to unwrap the unified signature container if it exists.
        entity_type = self.prompt_signature
        if "entities" in entity_type.model_fields:
            entity_type = entity_type.model_fields["entities"].annotation
            if typing.get_origin(entity_type) is list:
                entity_type = typing.get_args(entity_type)[0]
        elif "entity" in entity_type.model_fields:
            entity_type = entity_type.model_fields["entity"].annotation
            if typing.get_origin(entity_type) in (typing.Union, types.UnionType):
                args = typing.get_args(entity_type)
                entity_type = [t for t in args if t is not type(None)][0]

        # Define metadata.
        hf_entity_type = PydanticToHFDatasets.model_cls_to_features(entity_type)
        if self._mode == "multi":
            features = datasets.Features(
                {"text": datasets.Value("string"), "entities": datasets.Sequence(hf_entity_type)}
            )
        else:
            features = datasets.Features({"text": datasets.Value("string"), "entity": hf_entity_type})

        entity_type_name = getattr(entity_type, "__name__", str(entity_type))
        info = datasets.DatasetInfo(
            description=f"Information extraction dataset for entity type {entity_type_name}. "
            f"Generated with sieves v{Config.get_version()}.",
            features=features,
        )

        # Fetch data used for generating dataset.
        try:
            if self._mode == "multi":
                data = [
                    (doc.text, [PydanticToHFDatasets.model_to_dict(res) for res in doc.results[self._task_id].entities])
                    for doc in docs
                ]
            else:
                data = [
                    (
                        doc.text,
                        PydanticToHFDatasets.model_to_dict(doc.results[self._task_id].entity)
                        if doc.results[self._task_id].entity
                        else None,
                    )
                    for doc in docs
                ]
        except KeyError as err:
            raise KeyError(f"Not all documents have results for this task with ID {self._task_id}") from err

        # Create dataset.
        return datasets.Dataset.from_list(
            [
                {
                    "text": text,
                    ("entities" if self._mode == "multi" else "entity"): res,
                }
                for text, res in 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:
        if self._mode == "single":
            true_entity = truth.get("entity")
            pred_entity = pred.get("entity")
            base_accuracy = 1.0 if true_entity == pred_entity else 0.0

            # If score is available in both truth and pred, incorporate it into the metric.
            if "score" in truth and truth["score"] is not None and "score" in pred and pred["score"] is not None:
                score_accuracy = 1 - abs(truth["score"] - max(min(pred["score"], 1), 0))
                return (base_accuracy + score_accuracy) / 2

            return base_accuracy

        def entity_to_tuple(entity: dict) -> tuple:
            """Convert entity dict to hashable tuple for comparison.

            Converts nested structures (lists, dicts) to strings to make them hashable.

            :param entity: Entity dictionary to convert.
            :return: Hashable tuple representation of the entity.
            """
            items = sorted(entity.items())
            return tuple((k, v if not (isinstance(v, list) or isinstance(v, dict)) else str(v)) for k, v in items)

        # Compute set-based F1 score for entity extraction
        true_entities = {entity_to_tuple(e) for e in truth["entities"]}
        pred_entities = {entity_to_tuple(e) 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)
        f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0

        # If scores are available, incorporate them.
        if "scores" in truth and truth["scores"] is not None and "scores" in pred and pred["scores"] is not None:
            # Simple averaging of score similarities for matched entities is complex here due to set-based matching.
            # We'll just average the overall entity f1 with a score similarity metric.
            true_avg_score = sum(truth["scores"]) / len(truth["scores"]) if truth["scores"] else 0
            pred_avg_score = sum(pred["scores"]) / len(pred["scores"]) if pred["scores"] else 0
            score_accuracy = 1 - abs(true_avg_score - max(min(pred_avg_score, 1), 0))
            return (f1 + score_accuracy) / 2

        return f1
