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

from __future__ import annotations

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

import datasets
import pydantic

from sieves.data import Doc
from sieves.engines import EngineType, dspy_, instructor_, langchain_, ollama_, outlines_, vllm_
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 PredictiveTask
from sieves.tasks.predictive.sentiment_analysis.bridges import (
    DSPySentimentAnalysis,
    InstructorSentimentAnalysis,
    LangChainSentimentAnalysis,
    OllamaSentimentAnalysis,
    OutlinesSentimentAnalysis,
    VLLMSentimentAnalysis,
)

_TaskModel = dspy_.Model | instructor_.Model | langchain_.Model | ollama_.Model | outlines_.Model | vllm_.Model
_TaskPromptSignature = pydantic.BaseModel | dspy_.PromptSignature
_TaskResult = str | pydantic.BaseModel | dspy_.Result
_TaskBridge = (
    DSPySentimentAnalysis
    | InstructorSentimentAnalysis
    | LangChainSentimentAnalysis
    | OllamaSentimentAnalysis
    | OutlinesSentimentAnalysis
    | VLLMSentimentAnalysis
)


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

    text: str
    reasoning: str
    sentiment_per_aspect: dict[str, float]

    @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,
        prompt_template: str | None = None,
        prompt_signature_desc: str | None = None,
        fewshot_examples: Iterable[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 prompt_template: Custom prompt template. If None, task's default template is being used.
        :param prompt_signature_desc: Custom prompt signature description. If None, default will be 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,
            overwrite=False,
            prompt_template=prompt_template,
            prompt_signature_desc=prompt_signature_desc,
            fewshot_examples=fewshot_examples,
        )
        self._fewshot_examples: Iterable[FewshotExample]

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

        try:
            bridge_type = bridge_types[engine_type]

            return bridge_type(
                task_id=self._task_id,
                prompt_template=self._custom_prompt_template,
                prompt_signature_desc=self._custom_prompt_signature_desc,
                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.instructor,
            EngineType.langchain,
            EngineType.ollama,
            EngineType.outlines,
            EngineType.vllm,
        }

    @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,
        distillation_framework: DistillationFramework,
        hf_dataset: datasets.Dataset,
        init_kwargs: dict[str, Any],
        train_kwargs: dict[str, Any],
        output_path: Path | str,
        train_frac: float,
        val_frac: float,
        seed: int | None = None,
    ) -> None:
        raise NotImplementedError
