"""Aspect-based sentiment analysis predictive task."""

from __future__ import annotations

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

import datasets
import dspy
import pydantic

from sieves.data import Doc
from sieves.engines import EngineType, dspy_, langchain_, outlines_
from sieves.engines.types import GenerationSettings
from sieves.serialization import Config
from sieves.tasks.postprocessing.distillation.types import DistillationFramework
from sieves.tasks.predictive.core import FewshotExample as BaseFewshotExample
from sieves.tasks.predictive.core import PredictiveTask
from sieves.tasks.predictive.sentiment_analysis.bridges import (
    DSPySentimentAnalysis,
    LangChainSentimentAnalysis,
    OutlinesSentimentAnalysis,
)

_TaskModel = dspy_.Model | langchain_.Model | outlines_.Model
_TaskPromptSignature = pydantic.BaseModel | dspy_.PromptSignature
_TaskResult = str | pydantic.BaseModel | dspy_.Result
_TaskBridge = DSPySentimentAnalysis | LangChainSentimentAnalysis | OutlinesSentimentAnalysis


class FewshotExample(BaseFewshotExample):
    """Few-shot example with per-aspect sentiment scores."""

    reasoning: str
    sentiment_per_aspect: dict[str, float]

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

    @pydantic.model_validator(mode="after")
    def check_confidence(self) -> FewshotExample:
        """Validate that 'overall' exists and all scores are in [0, 1]."""
        assert "overall" in self.sentiment_per_aspect, ValueError(
            "'overall' score has to be given in `sentiment_per_aspect` dict."
        )
        if any([conf for conf in self.sentiment_per_aspect.values() if not 0 <= conf <= 1]):
            raise ValueError("Sentiment score has to be between 0 and 1.")
        return self


class SentimentAnalysis(PredictiveTask[_TaskPromptSignature, _TaskResult, _TaskBridge]):
    """Estimate per‑aspect and overall sentiment for a document."""

    def __init__(
        self,
        model: _TaskModel,
        generation_settings: GenerationSettings = GenerationSettings(),
        aspects: tuple[str, ...] = tuple(),
        task_id: str | None = None,
        include_meta: bool = True,
        batch_size: int = -1,
        prompt_instructions: str | None = None,
        fewshot_examples: Sequence[FewshotExample] = (),
    ) -> None:
        """
        Initialize SentimentAnalysis task.

        :param model: Model to use.
        :param generation_settings: Settings for structured generation.
        :param aspects: Aspects to consider in sentiment analysis. Overall sentiment will always be determined. If
            empty, only overall sentiment will be determined.
        :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.
        """
        self._aspects = tuple(sorted(set(aspects) | {"overall"}))
        super().__init__(
            model=model,
            generation_settings=generation_settings,
            task_id=task_id,
            include_meta=include_meta,
            batch_size=batch_size,
            overwrite=False,
            prompt_instructions=prompt_instructions,
            fewshot_examples=fewshot_examples,
        )
        self._fewshot_examples: Sequence[FewshotExample]

    @override
    def _init_bridge(self, engine_type: EngineType) -> _TaskBridge:
        bridge_types: dict[EngineType, type[_TaskBridge]] = {
            EngineType.dspy: DSPySentimentAnalysis,
            EngineType.outlines: OutlinesSentimentAnalysis,
            EngineType.langchain: LangChainSentimentAnalysis,
        }

        try:
            bridge_type = bridge_types[engine_type]

            return bridge_type(
                task_id=self._task_id,
                prompt_instructions=self._custom_prompt_instructions,
                aspects=self._aspects,
            )
        except KeyError as err:
            raise KeyError(f"Engine type {engine_type} is not supported by {self.__class__.__name__}.") from err

    @property
    @override
    def supports(self) -> set[EngineType]:
        return {
            EngineType.dspy,
            EngineType.langchain,
            EngineType.outlines,
        }

    @override
    def _validate_fewshot_examples(self) -> None:
        for fs_example in self._fewshot_examples or []:
            if any([aspect not in self._aspects for aspect in fs_example.sentiment_per_aspect]) or not all(
                [label in fs_example.sentiment_per_aspect for label in self._aspects]
            ):
                raise ValueError(
                    f"Aspect mismatch: {self._task_id} has aspects {self._aspects}. Few-shot examples have "
                    f"aspects {fs_example.sentiment_per_aspect.keys()}."
                )

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

    @override
    def to_hf_dataset(self, docs: Iterable[Doc], threshold: float = 0.5) -> datasets.Dataset:
        # Define metadata.
        features = datasets.Features(
            {"text": datasets.Value("string"), "aspect": datasets.Sequence(datasets.Value("float32"))}
        )
        info = datasets.DatasetInfo(
            description=f"Aspect-based sentiment analysis dataset with aspects {self._aspects}. Generated with sieves "
            f"v{Config.get_version()}.",
            features=features,
        )

        # Fetch data used for generating dataset.
        aspects = self._aspects
        try:
            data = [(doc.text, doc.results[self._task_id]) 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

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

            :return: Results as dicts.
            """
            for text, result in data:
                scores = {sent_score[0]: sent_score[1] for sent_score in result}
                yield {"text": text, "aspect": [scores[aspect] for aspect in aspects]}

        # 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_optimization_example(
        self, truth: dspy.Example, pred: dspy.Prediction, model: dspy.LM, trace: Any | None = None
    ) -> float:
        # Compute per-aspect accuracy as 1 - abs(true_sentiment - pred_sentiment)
        # Average across all aspects (same approach as multi-label classification)
        accuracy = 0
        for aspect, sentiment in truth["sentiment_per_aspect"].items():
            if aspect in pred["sentiment_per_aspect"]:
                pred_sentiment = max(min(pred["sentiment_per_aspect"][aspect], 1), 0)
                accuracy += 1 - abs(sentiment - pred_sentiment)

        return accuracy / len(truth["sentiment_per_aspect"])
