from __future__ import annotations

from collections.abc import Iterable
from typing import Any, TypeAlias

import datasets
import pydantic

from sieves.data import Doc
from sieves.engines import Engine, EngineType, dspy_, glix_, huggingface_, instructor_, langchain_, ollama_, outlines_
from sieves.engines.core import EngineInferenceMode, EngineModel, EnginePromptSignature, EngineResult
from sieves.serialization import Config
from sieves.tasks.predictive.bridges import GliXBridge
from sieves.tasks.predictive.classification.bridges import (
    DSPyClassification,
    HuggingFaceClassification,
    InstructorClassification,
    LangChainClassification,
    OllamaClassification,
    OutlinesClassification,
)
from sieves.tasks.predictive.core import PredictiveTask

_TaskPromptSignature: TypeAlias = glix_.PromptSignature | pydantic.BaseModel | dspy_.PromptSignature
_TaskInferenceMode: TypeAlias = (
    outlines_.InferenceMode
    | dspy_.InferenceMode
    | huggingface_.InferenceMode
    | instructor_.InferenceMode
    | glix_.InferenceMode
    | langchain_.InferenceMode
    | ollama_.InferenceMode
)
_TaskResult: TypeAlias = str | pydantic.BaseModel | dspy_.Result | huggingface_.Result | glix_.Result
_TaskBridge: TypeAlias = (
    DSPyClassification
    | GliXBridge
    | InstructorClassification
    | LangChainClassification
    | HuggingFaceClassification
    | OllamaClassification
    | OutlinesClassification
)


class TaskFewshotExample(pydantic.BaseModel):
    text: str
    reasoning: str
    confidence_per_label: dict[str, float]

    @pydantic.model_validator(mode="after")
    def check_confidence(self) -> TaskFewshotExample:
        if any([conf for conf in self.confidence_per_label.values() if not 0 <= conf <= 1]):
            raise ValueError("Confidence has to be between 0 and 1.")
        return self


class Classification(PredictiveTask[_TaskPromptSignature, _TaskResult, _TaskBridge]):
    def __init__(
        self,
        labels: list[str],
        engine: Engine[EnginePromptSignature, EngineResult, EngineModel, EngineInferenceMode],
        task_id: str | None = None,
        show_progress: bool = True,
        include_meta: bool = True,
        prompt_template: str | None = None,
        prompt_signature_desc: str | None = None,
        fewshot_examples: Iterable[TaskFewshotExample] = (),
    ) -> None:
        """
        Initializes new PredictiveTask.
        :param labels: Labels to predict.
        :param task_id: Task ID.
        :param show_progress: Whether to show progress bar for processed documents.
        :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._labels = labels
        super().__init__(
            engine=engine,
            task_id=task_id,
            show_progress=show_progress,
            include_meta=include_meta,
            overwrite=False,
            prompt_template=prompt_template,
            prompt_signature_desc=prompt_signature_desc,
            fewshot_examples=fewshot_examples,
        )
        self._fewshot_examples: Iterable[TaskFewshotExample]

    def _init_bridge(self, engine_type: EngineType) -> _TaskBridge:
        """Initialize bridge.
        :return: Engine task.
        :raises ValueError: If engine type is not supported.
        """
        if engine_type == EngineType.glix:
            # GliXBridge needs different arguments than other bridges, hence we instantiate it differently.
            return GliXBridge(
                task_id=self._task_id,
                prompt_template=self._custom_prompt_template,
                prompt_signature_desc=self._custom_prompt_signature_desc,
                prompt_signature=self._labels,
                inference_mode=glix_.InferenceMode.classification,
                label_whitelist=tuple(self._labels),
            )

        bridge_types: dict[EngineType, type[_TaskBridge]] = {
            EngineType.dspy: DSPyClassification,
            EngineType.instructor: InstructorClassification,
            EngineType.huggingface: HuggingFaceClassification,
            EngineType.outlines: OutlinesClassification,
            EngineType.ollama: OllamaClassification,
            EngineType.langchain: LangChainClassification,
        }

        try:
            bridge_type = bridge_types[engine_type]
            assert not issubclass(bridge_type, GliXBridge)

            return bridge_type(
                task_id=self._task_id,
                prompt_template=self._custom_prompt_template,
                prompt_signature_desc=self._custom_prompt_signature_desc,
                labels=self._labels,
            )
        except KeyError as err:
            raise KeyError(f"Engine type {engine_type} is not supported by {self.__class__.__name__}.") from err

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

    def _validate_fewshot_examples(self) -> None:
        for fs_example in self._fewshot_examples or []:
            if any([label not in self._labels for label in fs_example.confidence_per_label]) or not all(
                [label in fs_example.confidence_per_label for label in self._labels]
            ):
                raise ValueError(
                    f"Label mismatch: {self._task_id} has labels {self._labels}. Few-shot examples has "
                    f"labels {fs_example.confidence_per_label.keys()}."
                )

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

    def to_dataset(self, docs: Iterable[Doc]) -> datasets.Dataset:
        # Define metadata.
        features = datasets.Features(
            {"text": datasets.Value("string"), "label": datasets.Sequence(datasets.Value("float32"))}
        )
        info = datasets.DatasetInfo(
            description=f"Multi-label classification dataset with labels {self._labels}. Generated with sieves "
            f"v{Config.get_version()}.",
            features=features,
        )

        # Fetch data used for generating dataset.
        labels = self._labels
        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]]:
            """Yields results as dicts.
            :return: Results as dicts.
            """
            for text, result in data:
                scores = {label_score[0]: label_score[1] for label_score in result}
                yield {"text": text, "label": [scores[label] for label in labels]}

        # Create dataset.
        return datasets.Dataset.from_generator(generate_data, features=features, info=info)
