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

from __future__ import annotations

import json
from collections.abc import Callable, Iterable, Sequence
from pathlib import Path
from typing import Any, Literal, override

import datasets
import dspy
import pydantic
import sklearn

from sieves.data import Doc
from sieves.model_wrappers import ModelType, gliner_
from sieves.model_wrappers.types import ModelSettings
from sieves.serialization import Config
from sieves.tasks.distillation.distillation_import import model2vec, setfit
from sieves.tasks.distillation.types import DistillationFramework
from sieves.tasks.predictive.classification.bridges import (
    DSPyClassification,
    HuggingFaceClassification,
    LangChainClassification,
    OutlinesClassification,
)
from sieves.tasks.predictive.core import PredictiveTask
from sieves.tasks.predictive.gliner_bridge import GliNERBridge
from sieves.tasks.predictive.schemas.classification import (
    FewshotExampleMultiLabel,
    FewshotExampleSingleLabel,
    ResultMultiLabel,
    ResultSingleLabel,
    TaskModel,
    TaskPromptSignature,
    TaskResult,
)

_TaskBridge = (
    DSPyClassification | GliNERBridge | LangChainClassification | HuggingFaceClassification | OutlinesClassification
)

FewshotExample = FewshotExampleMultiLabel | FewshotExampleSingleLabel


class Classification(PredictiveTask[TaskPromptSignature, TaskResult, _TaskBridge]):
    """Predictive task for text classification."""

    def __init__(
        self,
        labels: Sequence[str] | dict[str, 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] = (),
        mode: Literal["single", "multi"] = "multi",
        model_settings: ModelSettings = ModelSettings(),
        condition: Callable[[Doc], bool] | None = None,
    ) -> None:
        """Initialize new Classification task.

        :param labels: Labels to predict. Supports two formats:
            - List format: `["science", "politics", "sports"]`
            - Dict format: `{"science": "Scientific topics", "politics": "Political topics"}`
            The dict format allows you to provide descriptions that help the model better understand each label.
        :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 mode: If 'multi', task returns confidence scores for all specified labels. If 'single', task returns
            most likely class label.
        :param model_settings: Model settings.
        :param condition: Optional callable that determines whether to process each document.
        """
        if isinstance(labels, dict):
            self._labels = list(labels.keys())
            self._label_descriptions = labels
        else:
            self._labels = list(labels)
            self._label_descriptions = {}
        self._mode = mode

        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,
            model_settings=model_settings,
            condition=condition,
        )

    @property
    @override
    def fewshot_example_type(self) -> type[FewshotExample]:
        """Return few-shot example type.

        :return: Few-shot example type.
        """
        if self._mode == "multi":
            return FewshotExampleMultiLabel

        return FewshotExampleSingleLabel

    @property
    @override
    def prompt_signature(self) -> type[pydantic.BaseModel]:
        if self._mode == "single":
            labels = self._labels
            LabelType = Literal[*labels]  # type: ignore[valid-type]

            class SingleLabelClassification(pydantic.BaseModel):
                """Result of single-label classification. Contains the most likely label and its confidence score."""

                label: LabelType = pydantic.Field(description="The predicted label from the specified set of labels.")
                score: float = pydantic.Field(
                    default=None, description="Provide a confidence score for the predicted label, between 0 and 1."
                )

            return SingleLabelClassification

        # For multi-label, create a model with fields for each label.
        fields: dict[str, tuple[type, Any]] = {}
        for label in self._labels:
            assert isinstance(self._label_descriptions, dict)
            description = self._label_descriptions.get(
                label, f"Confidence score for the '{label}' category, between 0 and 1."
            )  # type: ignore[no-matching-overload]
            fields[label] = (float, pydantic.Field(description=description))

        return pydantic.create_model(  # type: ignore[no-matching-overload]
            "MultiLabelClassification",
            __base__=pydantic.BaseModel,
            __doc__="Result of multi-label classification. Contains confidence scores for each potential label.",
            **fields,
        )

    @property
    @override
    def metric(self) -> str:
        return "F1 (Macro)"

    @override
    def _compute_metrics(self, truths: list[Any], preds: list[Any], judge: dspy.LM | None = None) -> dict[str, float]:
        """Compute corpus-level metrics.

        :param truths: List of ground truths.
        :param preds: List of predictions.
        :param judge: Optional DSPy LM instance to use as judge for generative tasks.
        :return: Dictionary of metrics.
        """
        # Prepare labels and mappings.
        labels_list = list(self._labels)
        label_to_idx = {label: i for i, label in enumerate(labels_list)}

        y_true: list[float | list[float] | int] = []
        y_pred: list[float | list[float] | int] = []

        for gold, pred in zip(truths, preds):
            if gold is None or pred is None:
                # If either is None, we handle it as a failure for this example.
                # In multi-label, this is an all-zero vector. In single-label, it's -1 (no class).
                if self._mode == "multi":
                    y_true.append([0.0] * len(labels_list))
                    y_pred.append([0.0] * len(labels_list))
                else:
                    y_true.append(-1)
                    y_pred.append(-1)
                continue

            # Convert to normalized representation using the existing logic.
            # We reuse `_task_result_to_dspy_dict` indirectly or just use `_result_to_scores`.
            gold_scores = self._result_to_scores(self._task_result_to_pydantic(gold))
            pred_scores = self._result_to_scores(self._task_result_to_pydantic(pred))

            if self._mode == "multi":
                # Binary multi-hot vectors.
                y_true.append([1.0 if gold_scores.get(label, 0.0) >= self.THRESHOLD else 0.0 for label in labels_list])
                y_pred.append([1.0 if pred_scores.get(label, 0.0) >= self.THRESHOLD else 0.0 for label in labels_list])
            else:
                # Single label indices.
                # Get the label with the highest score.
                gold_label = max(gold_scores.items(), key=lambda x: x[1])[0]
                pred_label = max(pred_scores.items(), key=lambda x: x[1])[0]
                y_true.append(label_to_idx.get(gold_label, -1))
                y_pred.append(label_to_idx.get(pred_label, -1))

        # Filter out examples where gold or pred were None (represented by -1 or all-zero if that's what we chose).
        score = sklearn.metrics.f1_score(y_true, y_pred, average="macro", zero_division=0)

        return {self.metric: float(score)}

    def _init_bridge(self, model_type: ModelType) -> _TaskBridge:
        """Initialize bridge.

        :return: ModelWrapper task.
        :raises ValueError: If model type is not supported.
        """
        # Reconstruct labels parameter (as dict if descriptions exist, else as list)
        labels = self._label_descriptions if self._label_descriptions else self._labels

        if model_type == ModelType.gliner:
            return GliNERBridge(
                task_id=self._task_id,
                prompt_instructions=self._custom_prompt_instructions,
                prompt_signature=self.prompt_signature,
                model_settings=self._model_settings,
                inference_mode=gliner_.InferenceMode.classification,
                mode=self._mode,
            )

        bridge_types: dict[ModelType, type[_TaskBridge]] = {
            ModelType.dspy: DSPyClassification,
            ModelType.huggingface: HuggingFaceClassification,
            ModelType.outlines: OutlinesClassification,
            ModelType.langchain: LangChainClassification,
        }

        try:
            bridge_type = bridge_types[model_type]
            assert not issubclass(bridge_type, GliNERBridge)

            return bridge_type(
                task_id=self._task_id,
                prompt_instructions=self._custom_prompt_instructions,
                labels=labels,
                mode=self._mode,
                model_settings=self._model_settings,
                prompt_signature=self.prompt_signature,
                model_type=model_type,
                fewshot_examples=self._fewshot_examples,
            )
        except KeyError as err:
            raise KeyError(f"Model type {model_type} is not supported by {self.__class__.__name__}.") from err

    @staticmethod
    @override
    def supports() -> set[ModelType]:
        return {
            ModelType.dspy,
            ModelType.gliner,
            ModelType.huggingface,
            ModelType.langchain,
            ModelType.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: mode = {mode} requires {example_type}."

        for fs_example in self._fewshot_examples or []:
            if self._mode == "multi":
                assert isinstance(fs_example, FewshotExampleMultiLabel), TypeError(
                    example_type_error_text.format(example_type=FewshotExampleMultiLabel, mode=self._mode)
                )
                if any([label not in self._labels for label in fs_example.score_per_label]) or not all(
                    [label in fs_example.score_per_label for label in self._labels]
                ):
                    raise ValueError(
                        label_error_text.format(
                            task_id=self.id, labels=self._labels, example_labels=fs_example.score_per_label.keys()
                        )
                    )
            else:
                assert isinstance(fs_example, FewshotExampleSingleLabel), TypeError(
                    example_type_error_text.format(example_type=FewshotExampleSingleLabel, mode=self._mode)
                )
                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]:
        # Store labels as dict if descriptions exist, otherwise as list
        labels = self._label_descriptions if self._label_descriptions else self._labels
        return {
            **super()._state,
            "labels": labels,
        }

    @staticmethod
    def _result_to_scores(result: ResultMultiLabel | ResultSingleLabel) -> dict[str, float]:
        """Normalize a single result to a mapping of label → 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, ResultMultiLabel):
            return {str(label): float(score) for label, score in result.label_scores}

        if isinstance(result, ResultSingleLabel):
            return {result.label: result.score}

        raise TypeError(f"Unsupported result type in _result_to_scores: {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._mode == "multi":
                    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 [str(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 | None = None) -> 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 model wrappers 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. Defaults to `THRESHOLD`.

        :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.

        """
        threshold = threshold or self.THRESHOLD

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

        # Define metadata and features (multi-hot across declared labels for multi-label).
        if self._mode == "multi":
            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._mode == 'multi' 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._mode == "multi":
                    result_normalized = [int(scores.get(label, 0.0) >= threshold) for label in self._labels]  # type: ignore[no-matching-overload]
                # 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)

    def _task_result_to_pydantic(self, result: Any) -> ResultSingleLabel | ResultMultiLabel:
        """Convert a result to the appropriate Pydantic model for this task.

        :param result: Input result.
        :return: Normalized Pydantic result.
        """
        if isinstance(result, ResultSingleLabel | ResultMultiLabel):
            return result

        # Use the same logic as `_task_result_to_dspy_dict` but return the Pydantic model.
        if isinstance(result, str):
            return ResultSingleLabel(label=result, score=1.0)
        elif isinstance(result, list) and all(isinstance(item, str) for item in result):
            if self._mode == "single":
                if len(result) == 1:
                    return ResultSingleLabel(label=result[0], score=1.0)
                else:
                    raise ValueError(f"Got list of {len(result)} labels for single-label task.")
            else:
                label_scores: list[tuple[str, float]] = []
                active_set = set(result)
                for label in self._labels:
                    s = 1.0 if label in active_set else 0.0
                    label_scores.append((label, s))
                return ResultMultiLabel(label_scores=label_scores)

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

    @override
    def _task_result_to_dspy_dict(self, result: Any) -> dict[str, Any]:
        # We accept `str` and `list[str]` for better UX.
        if isinstance(result, str):
            result = ResultSingleLabel(label=result, score=1.0)
        elif isinstance(result, list) and all(isinstance(item, str) for item in result):
            if self._mode == "single":
                if len(result) == 1:
                    result = ResultSingleLabel(label=result[0], score=1.0)
                else:
                    raise ValueError(f"Got list of {len(result)} labels for single-label task.")
            else:
                label_scores: list[tuple[str, float]] = []
                active_set = set(result)
                for label in self._labels:
                    s = 1.0 if label in active_set else 0.0
                    label_scores.append((label, s))
                result = ResultMultiLabel(label_scores=label_scores)

        assert isinstance(result, ResultSingleLabel | ResultMultiLabel)

        scores = self._result_to_scores(result)
        if self._mode == "multi":
            return {"score_per_label": scores}
        else:
            # For single label, it returns a dict with one key.
            label = list(scores.keys())[0]
            score = scores[label]
            return {"label": label, "score": score}

    @override
    def _evaluate_dspy_example(self, truth: dspy.Example, pred: dspy.Prediction, trace: Any, model: dspy.LM) -> float:
        if self._mode == "single":
            return 1.0 if truth["label"] == pred["label"] else 0.0

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

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