"""Relation extraction predictive task."""

from __future__ import annotations

import warnings
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.relation_extraction.bridges import DSPyRelationExtraction, PydanticRelationExtraction
from sieves.tasks.predictive.schemas.relation_extraction import (
    FewshotExample,
    TaskModel,
    TaskPromptSignature,
    TaskResult,
)

_TaskBridge = GliNERBridge | DSPyRelationExtraction | PydanticRelationExtraction


class RelationExtraction(PredictiveTask[TaskPromptSignature, TaskResult, _TaskBridge]):
    """Extract relations between entities in text."""

    def __init__(
        self,
        relations: Sequence[str] | dict[str, str],
        model: TaskModel,
        entity_types: Sequence[str] | dict[str, str] | 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 RelationExtraction task.

        :param relations: Relations to extract. Can be a list of relation types or a dict mapping types to descriptions.
        :param model: Model to use.
        :param entity_types: Optional constraints on entity types involved in relations.
        :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.
        """
        self._relations = relations
        self._entity_types = entity_types

        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 few-shot example type.

        :return: Few-shot example type.
        """
        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.
        """
        # Define relation type literal if relations are provided.
        if isinstance(self._relations, dict):
            relation_names = list(self._relations.keys())
        else:
            relation_names = list(self._relations)

        RelationType = Literal[*(tuple(relation_names))] if relation_names else str  # type: ignore[invalid-type-form]

        # Define entity type literal if entity types are provided.
        entity_type_names: list[str] = []
        if self._entity_types:
            if isinstance(self._entity_types, dict):
                entity_type_names = list(self._entity_types.keys())
            else:
                entity_type_names = list(self._entity_types)

        EntityType = Literal[*(tuple(entity_type_names))] if entity_type_names else str  # type: ignore[invalid-type-form]

        # Create dynamic models.
        DynamicEntity = pydantic.create_model(
            "RelationEntity",
            text=(str, pydantic.Field(..., description="Surface text of the entity as it appears in the document.")),
            entity_type=(
                EntityType,
                pydantic.Field(..., description="The category or type of the entity."),
            ),
            __doc__="An entity involved in a relation.",
            __base__=pydantic.BaseModel,
        )

        DynamicTriplet = pydantic.create_model(
            "RelationTriplet",
            head=(DynamicEntity, pydantic.Field(..., description="The subject entity (head) of the relation.")),
            relation=(RelationType, pydantic.Field(..., description="The type of relation between the head and tail.")),
            tail=(DynamicEntity, pydantic.Field(..., description="The object entity (tail) of the relation.")),
            score=(
                float | None,
                pydantic.Field(
                    default=None, description="Provide a confidence score for this relation triplet, between 0 and 1."
                ),
            ),
            __doc__="A relation triplet consisting of a head entity, a relation type, and a tail entity.",
            __base__=pydantic.BaseModel,
        )

        return pydantic.create_model(
            "RelationExtractionOutput",
            __doc__="Result of relation extraction. Contains a list of extracted relation triplets.",
            triplets=(
                list[DynamicTriplet],  # type: ignore[invalid-type-form]
                pydantic.Field(..., description="List of extracted relation triplets."),
            ),
        )

    @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 not None:
                assert isinstance(gold, TaskResult)
                true_triplets = {(t.head.text.lower(), t.relation.lower(), t.tail.text.lower()) for t in gold.triplets}
            else:
                true_triplets = set()

            if pred is not None:
                assert isinstance(pred, TaskResult)
                pred_triplets = {(t.head.text.lower(), t.relation.lower(), t.tail.text.lower()) for t in pred.triplets}
            else:
                pred_triplets = set()

            tp += len(true_triplets & pred_triplets)
            fp += len(pred_triplets - true_triplets)
            fn += len(true_triplets - pred_triplets)

        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:
            if self._entity_types is not None:
                warnings.warn(
                    "GliNER2 backend does not support entity type constraints for relation extraction. "
                    "The `entity_types` parameter will be ignored.",
                )

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

        bridge_types: dict[ModelType, type[DSPyRelationExtraction | PydanticRelationExtraction]] = {
            ModelType.dspy: DSPyRelationExtraction,
            ModelType.langchain: PydanticRelationExtraction,
            ModelType.outlines: PydanticRelationExtraction,
        }

        try:
            return bridge_types[model_type](
                task_id=self._task_id,
                relations=self._relations,
                entity_types=self._entity_types,
                prompt_instructions=self._custom_prompt_instructions,
                model_settings=self._model_settings,
                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

    @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,
            "relations": self._relations,
            "entity_types": self._entity_types,
        }

    @override
    def to_hf_dataset(self, docs: Iterable[Doc], threshold: float | None = None) -> datasets.Dataset:
        # Define metadata and features.
        entity_feature = datasets.Features(
            {
                "text": datasets.Value("string"),
                "entity_type": datasets.Value("string"),
            }
        )
        triplet_feature = datasets.Features(
            {
                "head": entity_feature,
                "relation": datasets.Value("string"),
                "tail": entity_feature,
                "score": datasets.Value("float32"),
            }
        )
        features = datasets.Features({"text": datasets.Value("string"), "triplets": datasets.Sequence(triplet_feature)})
        info = datasets.DatasetInfo(
            description=f"Relation extraction dataset. Generated with sieves v{Config.get_version()}.",
            features=features,
        )

        # Fetch data used for generating dataset.
        try:
            data: list[dict[str, Any]] = []
            for doc in docs:
                triplets: list[dict[str, Any]] = []
                for triplet in doc.results[self._task_id].triplets:
                    triplets.append(
                        {
                            "head": triplet.head.model_dump(),
                            "relation": triplet.relation,
                            "tail": triplet.tail.model_dump(),
                            "score": triplet.score,
                        }
                    )
                data.append({"text": doc.text, "triplets": triplets})
        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(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 triplet-level F1 score based on (head_text, relation, tail_text) triples.
        # Use lowercase for robust matching.
        true_triplets = {
            (t["head"]["text"].lower(), t["relation"].lower(), t["tail"]["text"].lower()) for t in truth["triplets"]
        }
        pred_triplets = {
            (t["head"]["text"].lower(), t["relation"].lower(), t["tail"]["text"].lower())
            for t in pred.get("triplets", [])
        }

        if not true_triplets:
            base_f1 = 1.0 if not pred_triplets else 0.0
        else:
            precision = len(true_triplets & pred_triplets) / len(pred_triplets) if pred_triplets else 0
            recall = len(true_triplets & pred_triplets) / len(true_triplets)
            base_f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0

        # If score is available, incorporate it.
        if "triplets" in truth and "triplets" in pred:
            true_scores = [t.get("score") for t in truth["triplets"] if t.get("score") is not None]
            pred_scores = [t.get("score") for t in pred["triplets"] if t.get("score") is not None]

            if true_scores and pred_scores:
                true_avg = sum(true_scores) / len(true_scores)
                pred_avg = sum(pred_scores) / len(pred_scores)
                score_accuracy = 1 - abs(true_avg - max(min(pred_avg, 1), 0))
                return (base_f1 + score_accuracy) / 2

        return base_f1
