"""Text summarization predictive task."""

from __future__ import annotations

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

import datasets
import dspy
import pydantic

from sieves.data import Doc
from sieves.model_wrappers import ModelType
from sieves.model_wrappers.types import ModelSettings
from sieves.serialization import Config
from sieves.tasks.distillation.types import DistillationFramework
from sieves.tasks.predictive.core import PredictiveTask
from sieves.tasks.predictive.schemas.summarization import (
    FewshotExample,
    TaskModel,
    TaskPromptSignature,
    TaskResult,
)
from sieves.tasks.predictive.summarization.bridges import (
    DSPySummarization,
    PydanticSummarization,
)

_TaskBridge = DSPySummarization | PydanticSummarization


class Summarization(PredictiveTask[TaskPromptSignature, TaskResult, _TaskBridge]):
    """Summarize documents to a target length using structured model wrappers."""

    def __init__(
        self,
        n_words: int,
        model: TaskModel,
        task_id: str | None = None,
        include_meta: bool = True,
        batch_size: int = -1,
        overwrite: bool = False,
        prompt_instructions: str | None = None,
        fewshot_examples: Sequence[FewshotExample] = (),
        model_settings: ModelSettings = ModelSettings(),
        condition: Callable[[Doc], bool] | None = None,
    ) -> None:
        """Initialize new Summarization task.

        :param n_words: Maximal number of words (consider this a guideline, not a strict limit).
        :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 overwrite: Some tasks, e.g. anonymization or translation, output a modified version of the input text.
            If True, these tasks overwrite the original document text. If False, the result will just be stored in the
            documents' `.results` field.
        :param prompt_instructions: Custom prompt instructions. If None, default instructions are used.
        :param fewshot_examples: Few-shot examples.
        :param model_settings: Settings for structured generation. Use the `inference_mode` field to specify the
            inference mode for the model wrapper.
        :param condition: Optional callable that determines whether to process each document.
        """
        self._n_words = n_words

        super().__init__(
            model=model,
            task_id=task_id,
            include_meta=include_meta,
            batch_size=batch_size,
            overwrite=overwrite,
            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.
        """
        return FewshotExample

    @property
    @override
    def prompt_signature(self) -> type[pydantic.BaseModel]:
        # Dynamically create type with length restriction in description.
        class ResultWithLengthRestriction(TaskResult):
            pass

        ResultWithLengthRestriction.model_fields[
            "summary"
        ].description += f" This summary should be around {self._n_words} words."

        return ResultWithLengthRestriction

    @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.
        """
        for gold in truths:
            if gold is not None:
                assert isinstance(gold, TaskResult)
        for pred in preds:
            if pred is not None:
                assert isinstance(pred, TaskResult)

        return super()._compute_metrics(truths, preds, judge=judge)

    @override
    def _init_bridge(self, model_type: ModelType) -> _TaskBridge:
        bridge_types: dict[ModelType, type[DSPySummarization | PydanticSummarization]] = {
            ModelType.dspy: DSPySummarization,
            ModelType.langchain: PydanticSummarization,
            ModelType.outlines: PydanticSummarization,
        }

        try:
            return bridge_types[model_type](
                task_id=self._task_id,
                prompt_instructions=self._custom_prompt_instructions,
                n_words=self._n_words,
                overwrite=self._overwrite,
                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.langchain,
            ModelType.outlines,
        }

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

    @override
    def to_hf_dataset(self, docs: Iterable[Doc], threshold: float | None = None) -> datasets.Dataset:
        # Define metadata.
        features = datasets.Features(
            {"text": datasets.Value("string"), "summary": datasets.Value("string"), "score": datasets.Value("float32")}
        )
        info = datasets.DatasetInfo(
            description=f"Summarization dataset. Generated with sieves v{Config.get_version()}.",
            features=features,
        )

        # Fetch data used for generating dataset.
        try:
            data = [(doc.text, doc.results[self._task_id].summary, doc.results[self._task_id].score) 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 results as dicts.

            :return: Results as dicts.
            """
            for text, summary, score in data:
                yield {"text": text, "summary": summary, "score": score}

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

    @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:
        raise NotImplementedError
