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

from __future__ import annotations

import json
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_, glix_, huggingface_, langchain_, outlines_
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,
    LangChainClassification,
    OutlinesClassification,
)
from sieves.tasks.predictive.core import FewshotExample as BaseFewshotExample
from sieves.tasks.predictive.core import PredictiveTask

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


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

    reasoning: str
    confidence_per_label: dict[str, float]

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

    @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(BaseFewshotExample):
    """Few‑shot example for single‑label classification with a global confidence."""

    reasoning: str
    label: str
    confidence: float

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

    @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,
        batch_size: int = -1,
        prompt_instructions: str | None = None,
        fewshot_examples: Sequence[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 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.
        :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,
            batch_size=batch_size,
            overwrite=False,
            prompt_instructions=prompt_instructions,
            fewshot_examples=fewshot_examples,
            generation_settings=generation_settings,
        )
        self._fewshot_examples: Sequence[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_instructions=self._custom_prompt_instructions,
                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.huggingface: HuggingFaceClassification,
            EngineType.outlines: OutlinesClassification,
            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_instructions=self._custom_prompt_instructions,
                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.glix,
            EngineType.huggingface,
            EngineType.langchain,
            EngineType.outlines,
        }

    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)}")

    @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:
        init_kwargs = init_kwargs or {}
        train_kwargs = train_kwargs or {}
        output_path = Path(output_path)
        output_path.mkdir(parents=True, exist_ok=True)

        data = self.to_hf_dataset(data) if isinstance(data, Sequence) else data

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

        dataset_splits = self._split_dataset(data, 1 - val_frac, val_frac, seed)
        dataset_splits.save_to_disk(output_path / "data")

        match 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.get("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: {framework}. "
                    f"Please choose one of {DistillationFramework.setfit, DistillationFramework.model2vec}"
                )

    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.

        """
        data: list[dict[str, str | list[bool]]] = []

        # 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,
        )

        try:
            for doc in docs:
                scores = Classification._result_to_scores(doc.results[self._task_id])

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

                data.append({"text": doc.text, "labels": result_normalized})

        except KeyError as err:
            raise KeyError(f"Not all documents have results for this task with ID {self._task_id}") from err

        return datasets.Dataset.from_list(data, features=features, info=info)

    @override
    def _evaluate_optimization_example(
        self, truth: dspy.Example, pred: dspy.Prediction, model: dspy.LM, trace: Any | None = None
    ) -> float:
        if not self._multi_label:
            return 1 - abs(truth["confidence"] - pred["confidence"]) if truth["label"] == pred["label"] else 0

        # For multi-label: compute label-wise accuracy as
        # 1 - abs(true confidence for label - predicted confidence for label)
        # and normalize the sum of label-wise accuracies over all labels.
        accuracy = 0
        for label, confidence in truth["confidence_per_label"].items():
            if label in pred["confidence_per_label"]:
                pred_confidence = max(min(pred["confidence_per_label"][label], 1), 0)
                accuracy += 1 - abs(confidence - pred_confidence)

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