"""Classification predictive task and few‑shot example schemas."""

from __future__ import annotations

import json
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_, glix_, huggingface_, instructor_, langchain_, ollama_, outlines_, vllm_
from sieves.engines.types import GenerationSettings
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

_TaskModel = (
    dspy_.Model
    | glix_.Model
    | instructor_.Model
    | langchain_.Model
    | huggingface_.Model
    | ollama_.Model
    | huggingface_.Model
    | outlines_.Model
    | vllm_.Model
)
_TaskPromptSignature = glix_.PromptSignature | pydantic.BaseModel | dspy_.PromptSignature | vllm_.PromptSignature
_TaskResult = str | pydantic.BaseModel | dspy_.Result | huggingface_.Result | glix_.Result | vllm_.Result
_TaskBridge = (
    DSPyClassification
    | GliXBridge
    | InstructorClassification
    | LangChainClassification
    | HuggingFaceClassification
    | OllamaClassification
    | OutlinesClassification
    | VLLMClassification
)

# TODO
#   - model signature must be _TaskModel, not engine wrapper.Model!
#   - use GenerationSettings consistently across all engines
#   - adjust tests


class FewshotExampleMultiLabel(pydantic.BaseModel):
    """Few‑shot example for multi‑label classification with per‑label confidences."""

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

    @pydantic.model_validator(mode="after")
    def check_confidence(self) -> FewshotExampleMultiLabel:
        """Validate that confidences lie within [0, 1]."""
        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):
    """Few‑shot example for single‑label classification with a global confidence."""

    text: str
    reasoning: str
    label: str
    confidence: float

    @pydantic.model_validator(mode="after")
    def check_confidence(self) -> FewshotExampleSingleLabel:
        """Check confidence value.

        Return:
            FewshotExampleSingleLabel instance.

        """
        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]):
    """Predictive task for text classification across multiple engine backends."""

    def __init__(
        self,
        labels: list[str],
        model: _TaskModel,
        task_id: str | None = None,
        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,
        generation_settings: GenerationSettings = GenerationSettings(),
    ) -> None:
        """Initialize new PredictiveTask.

        :param labels: Labels to predict.
        :param model: Model to use.
        :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.
        :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.
        :param generation_settings: Generation settings.
        """
        self._labels = labels
        self._label_descriptions = label_descriptions or {}
        self._validate_label_descriptions()
        self._multi_label = multi_label

        super().__init__(
            model=model,
            task_id=task_id,
            include_meta=include_meta,
            overwrite=False,
            prompt_template=prompt_template,
            prompt_signature_desc=prompt_signature_desc,
            fewshot_examples=fewshot_examples,
            generation_settings=generation_settings,
        )
        self._fewshot_examples: Iterable[FewshotExample]

    def _validate_label_descriptions(self) -> None:
        """Validate 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

    @override
    @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,
        }

    @staticmethod
    def _result_to_scores(result: Any) -> dict[str, float]:
        """Normalize a single result to a mapping of label → score.

        Supports lists of pairs, a single (label, score) pair, a plain
        string label (assumes score 1.0), or a Pydantic model with
        attributes ``label`` and optional ``score``.

        :params result: One result value from ``doc.results``.

        :return: Mapping from label to score.

        :raises TypeError: If the result has an unsupported type or shape.

        """
        if isinstance(result, list) and all(isinstance(item, list | tuple) and len(item) == 2 for item in result):
            return {str(label): float(score) for label, score in result}

        if isinstance(result, tuple) and len(result) == 2:
            label, score = result
            return {str(label): float(score)}

        if isinstance(result, str):
            return {result: 1.0}

        if isinstance(result, pydantic.BaseModel) or hasattr(result, "model_dump"):
            try:
                label = getattr(result, "label")
                score = getattr(result, "score", 1.0)
                return {str(label): float(score)}
            except Exception as exc:
                raise TypeError(f"Unsupported pydantic result shape: {type(result)}") from exc

        raise TypeError(f"Unsupported result type in to_hf_dataset: {type(result)}")

    def to_hf_dataset(self, docs: Iterable[Doc], threshold: float = 0.5) -> datasets.Dataset:
        """Convert results to a Hugging Face dataset with multi-hot labels.

        The emitted dataset contains a ``text`` column and a ``labels`` column which is a multi-hot list aligned to
        ``self._labels``. This method is robust to different result shapes produced by various engines and bridges in
        both single-label and multi-label configurations:
        - ``list[tuple[str, float]]`` for multi-label results
        - ``tuple[str, float]`` for single-label results
        - ``str`` for single-label results (assumes score ``1.0``)
        - ``pydantic.BaseModel`` exposing ``label`` and optional ``score``

        :param docs: Documents whose ``results`` contain outputs for this task id.
        :param threshold: Threshold to convert scores into multi-hot indicators.

        :return: A ``datasets.Dataset`` with ``text`` and multi-hot ``labels``.

        :raises KeyError: If any document is missing this task's results.
        :raises TypeError: If a result cannot be interpreted.

        """
        # Define metadata and features (multi-hot across declared labels for multi-label).
        if self._multi_label:
            features = datasets.Features(
                {"text": datasets.Value("string"), "labels": datasets.Sequence(datasets.Value("bool"))}
            )
        else:
            features = datasets.Features(
                {"text": datasets.Value("string"), "labels": datasets.ClassLabel(names=self._labels)}
            )

        info = datasets.DatasetInfo(
            description=(
                f"{'Multi-label' if self._multi_label else 'Single-label'} classification dataset with labels "
                f"{self._labels}. Generated with sieves v{Config.get_version()}."
            ),
            features=features,
        )

        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]]:
            """Yield rows with text and multi-hot labels for HF datasets.

            Iterates over the prepared ``data`` list, converts result shapes to
            a label→score mapping, and emits the multi-hot vector using the
            provided ``threshold``.

            Yields:
                Dicts with keys ``text`` and ``labels``.

            """
            for text, result in data:
                # Normalize result types.
                scores = Classification._result_to_scores(result)

                # If multi-label: store one-hot representation..
                if self._multi_label:
                    result_normalized = [int(scores.get(label, 0.0) >= threshold) for label in labels]
                # If single-label: get single-label result as is
                else:
                    keys = list(scores.keys())
                    assert len(keys) == 1
                    result_normalized = keys[0]

                yield {"text": text, "labels": result_normalized}

        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:
        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]:
                    """Convert 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}"
                )
