"""Information extraction."""

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 gliner2.inference.engine
import pydantic

from sieves.data import Doc
from sieves.model_wrappers import ModelType, dspy_, gliner_, langchain_, outlines_
from sieves.model_wrappers.types import ModelSettings
from sieves.serialization import Config
from sieves.tasks.distillation.types import DistillationFramework
from sieves.tasks.predictive.bridges import GliNERBridge
from sieves.tasks.predictive.core import FewshotExample as BaseFewshotExample
from sieves.tasks.predictive.core import PredictiveTask
from sieves.tasks.predictive.information_extraction.bridges import (
    DSPyInformationExtraction,
    LangChainInformationExtraction,
    OutlinesInformationExtraction,
)
from sieves.tasks.utils import PydanticToHFDatasets

_TaskModel = dspy_.Model | gliner_.Model | langchain_.Model | outlines_.Model
_TaskPromptSignature = pydantic.BaseModel | dspy_.PromptSignature | gliner_.PromptSignature
_TaskResult = gliner_.Result | outlines_.Result | dspy_.Result
_TaskBridge = GliNERBridge | DSPyInformationExtraction | LangChainInformationExtraction | OutlinesInformationExtraction


class FewshotExampleMulti(BaseFewshotExample):
    """Few-shot example for multi-entity extraction."""

    entities: list[pydantic.BaseModel]

    @override
    @property
    def target_fields(self) -> Sequence[str]:
        return ("entities",)


class FewshotExampleSingle(BaseFewshotExample):
    """Few-shot example for single-entity extraction."""

    entity: pydantic.BaseModel | None

    @override
    @property
    def target_fields(self) -> Sequence[str]:
        return ("entity",)


FewshotExample = FewshotExampleMulti | FewshotExampleSingle


class InformationExtraction(PredictiveTask[_TaskPromptSignature, _TaskResult, _TaskBridge]):
    """Information extraction task."""

    def __init__(
        self,
        entity_type: type[pydantic.BaseModel] | gliner2.inference.engine.StructureBuilder,
        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: Object type to extract. When using this task with a GLiNER2 model, this has to be of type
            `gliner2.inference.engine.StructureBuilder`; otherwise a `pydantic.BaseModel`.
        :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

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

    @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:
            if not isinstance(self._entity_type, gliner2.inference.engine.StructureBuilder):
                raise TypeError(
                    "You need to specify `entity_type` as a `gliner2.inference.engine.StructureBuilder` when running "
                    "this task with a GLiNER2 model."
                )

            return GliNERBridge(
                task_id=self._task_id,
                prompt_instructions=self._custom_prompt_instructions,
                prompt_signature=self._entity_type,
                model_settings=self._model_settings,
                inference_mode=gliner_.InferenceMode.structure,
                mode=self._mode,
            )

        if not issubclass(self._entity_type, pydantic.BaseModel):
            raise TypeError(
                "You need to use specify `entity_type` as a `pydantic.BaseModel` when not running with a GLiNER2 model."
            )

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

        try:
            bridge = bridge_types[model_type](
                task_id=self._task_id,
                prompt_instructions=self._custom_prompt_instructions,
                entity_type=self._entity_type,
                model_settings=self._model_settings,
                mode=self._mode,
            )
        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 fs_example in 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 = 0.5) -> datasets.Dataset:
        # If we're using GLiNER2: fetch Pydantic prompt signature representation.
        if isinstance(self._bridge, GliNERBridge):
            entity_type = self._bridge.prompt_signature_pydantic
            assert entity_type is not None
        else:
            entity_type = self._entity_type

        # 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})

        info = datasets.DatasetInfo(
            description=f"Information extraction dataset for entity type {entity_type.__class__.__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]])
                    for doc in docs
                ]
            else:
                data = [
                    (
                        doc.text,
                        PydanticToHFDatasets.model_to_dict(doc.results[self._task_id])
                        if doc.results[self._task_id]
                        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_optimization_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")
            return 1.0 if true_entity == pred_entity else 0.0

        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)
        return 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0
