from __future__ import annotations

import json
from collections.abc import Iterable
from pathlib import Path
from typing import Any, TypeAlias

import datasets
import pydantic

from sieves.data import Doc
from sieves.engines import Engine, EngineType, dspy_, glix_, huggingface_, vllm_
from sieves.serialization import Config
from sieves.tasks.postprocessing.distillation.distillation_import import model2vec, setfit
from sieves.tasks.postprocessing.distillation.types import DistillationFramework
from sieves.tasks.predictive.bridges import GliXBridge
from sieves.tasks.predictive.classification.bridges import (
    DSPyClassification,
    HuggingFaceClassification,
    InstructorClassification,
    LangChainClassification,
    OllamaClassification,
    OutlinesClassification,
    VLLMClassification,
)
from sieves.tasks.predictive.core import PredictiveTask

_TaskPromptSignature: TypeAlias = (
    glix_.PromptSignature | pydantic.BaseModel | dspy_.PromptSignature | vllm_.PromptSignature
)
_TaskResult: TypeAlias = str | pydantic.BaseModel | dspy_.Result | huggingface_.Result | glix_.Result | vllm_.Result
_TaskBridge: TypeAlias = (
    DSPyClassification
    | GliXBridge
    | InstructorClassification
    | LangChainClassification
    | HuggingFaceClassification
    | OllamaClassification
    | OutlinesClassification
    | VLLMClassification
)


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

    @pydantic.model_validator(mode="after")
    def check_confidence(self) -> FewshotExampleMultiLabel:
        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 FewshotExampleSingleLabel(pydantic.BaseModel):
    text: str
    reasoning: str
    label: str
    confidence: float

    @pydantic.model_validator(mode="after")
    def check_confidence(self) -> FewshotExampleSingleLabel:
        if not (0 <= self.confidence <= 1):
            raise ValueError("Confidence has to be between 0 and 1.")
        return self


FewshotExample = FewshotExampleMultiLabel | FewshotExampleSingleLabel


class Classification(PredictiveTask[_TaskPromptSignature, _TaskResult, _TaskBridge]):
    def __init__(
        self,
        labels: list[str],
        engine: Engine,
        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[FewshotExample] = (),
        label_descriptions: dict[str, str] | None = None,
        multi_label: bool = True,
    ) -> None:
        """
        Initializes new PredictiveTask.

        :param labels: Labels to predict.
        :param engine: Engine to use for prediction.
        :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.
        :param label_descriptions: Optional descriptions for each label. If provided, the keys must match the labels.
        :param multi_label: If True, task returns confidence scores for all specified labels. If False, task returns
            most likely class label. In the latter case label forcing mechanisms are utilized, which can lead to higher
            accuracy.
        """
        self._labels = labels
        self._label_descriptions = label_descriptions or {}
        self._validate_label_descriptions()
        self._multi_label = multi_label

        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[FewshotExample]

    def _validate_label_descriptions(self) -> None:
        """
        Validates that all label descriptions correspond to valid labels.

        :raises ValueError: If any label description key is not present in the labels list.
        """
        if not self._label_descriptions:
            return

        invalid_labels = set(self._label_descriptions.keys()) - set(self._labels)
        if invalid_labels:
            raise ValueError(f"Label descriptions contain invalid labels: {invalid_labels}")

    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),
                only_keep_best=not self._multi_label,
            )

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

        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,
                label_descriptions=self._label_descriptions,
                multi_label=self._multi_label,
            )
        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,
            EngineType.vllm,
        }

    def _validate_fewshot_examples(self) -> None:
        label_error_text = (
            "Label mismatch: {task_id} has labels {labels}. Few-shot examples have labels {example_labels}."
        )
        example_type_error_text = "Fewshot example type mismatch: multi_label = {multi_label} requires {example_type}."

        for fs_example in self._fewshot_examples or []:
            if self._multi_label:
                assert isinstance(fs_example, FewshotExampleMultiLabel), TypeError(
                    example_type_error_text.format(example_type=FewshotExampleMultiLabel, multi_label=self._multi_label)
                )
                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(
                        label_error_text.format(
                            task_id=self.id, labels=self._labels, example_labels=fs_example.confidence_per_label.keys()
                        )
                    )
            else:
                assert isinstance(fs_example, FewshotExampleSingleLabel), TypeError(
                    example_type_error_text.format(
                        example_type=FewshotExampleSingleLabel, multi_label=self._multi_label
                    )
                )
                if fs_example.label not in self._labels:
                    raise ValueError(
                        label_error_text.format(task_id=self.id, labels=self._labels, example_labels=(fs_example.label))
                    )

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

    def to_hf_dataset(self, docs: Iterable[Doc], threshold: float = 0.5) -> datasets.Dataset:
        # Define metadata.
        features = datasets.Features(
            {"text": datasets.Value("string"), "labels": datasets.Sequence(datasets.ClassLabel(names=self._labels))}
        )
        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,
                    # One-hot encode labels.
                    "labels": [int(scores.get(label, 0) >= threshold) for label in labels],
                }

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

    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:
        output_path = Path(output_path)
        output_path.mkdir(parents=True, exist_ok=True)

        required_columns = {"text", "labels"}
        if not required_columns.issubset(hf_dataset.column_names):
            raise ValueError(f"Dataset must contain columns: {required_columns}. Found: {hf_dataset.column_names}")

        dataset_splits = self._split_dataset(hf_dataset, train_frac, val_frac, seed)
        dataset_splits.save_to_disk(output_path / "data")

        # Framework-specific distillation/fine-tuning logic using match-case
        match distillation_framework:
            case DistillationFramework.setfit:
                default_init_kwargs: dict[str, Any] = {}
                metric_kwargs: dict[str, Any] = {}
                if self._multi_label:
                    default_init_kwargs["multi_target_strategy"] = "multi-output"
                    metric_kwargs = {"average": "macro"}

                model = setfit.SetFitModel.from_pretrained(base_model_id, **(default_init_kwargs | init_kwargs))

                args = setfit.TrainingArguments(
                    output_dir=str(output_path),
                    eval_strategy="epoch",
                    save_strategy="epoch",
                    load_best_model_at_end=True,
                    **train_kwargs,
                )

                trainer = setfit.Trainer(
                    model=model,
                    args=args,
                    train_dataset=dataset_splits["train"],
                    eval_dataset=dataset_splits["val"],
                    metric="f1",
                    column_mapping={"text": "text", "labels": "label"},
                    metric_kwargs=metric_kwargs,
                )
                trainer.train()
                trainer.model.save_pretrained(output_path)

                metrics = trainer.evaluate()
                with open(output_path / "metrics.json", "w") as f:
                    json.dump(metrics, f, indent=4)

            case DistillationFramework.model2vec:

                def one_hot_to_label(label_indices: list[int]) -> list[str]:
                    """Converts list of label indices into list of labels.
                    :param label_indices: List of label indices.
                    :return: List of labels.
                    """
                    return [self._labels[i] for i, is_label in enumerate(label_indices) if is_label]

                classifier = model2vec.train.StaticModelForClassification.from_pretrained(
                    model_name=base_model_id, **init_kwargs
                )
                classifier.fit(
                    dataset_splits["train"]["text"],
                    [one_hot_to_label(encoded_labels) for encoded_labels in dataset_splits["train"]["labels"]],
                    **train_kwargs,
                )
                classifier.to_pipeline().save_pretrained(output_path)

                metrics = classifier.evaluate(
                    dataset_splits["val"]["text"],
                    [one_hot_to_label(encoded_labels) for encoded_labels in dataset_splits["val"]["labels"]],
                )
                with open(output_path / "metrics.json", "w") as f:
                    json.dump(metrics, f, indent=4)

            case _:
                raise NotImplementedError(
                    f"Unsupported distillation framework for this task: {distillation_framework}. "
                    f"Please choose one of {DistillationFramework.setfit, DistillationFramework.model2vec}"
                )
