"""Base class for predictive tasks composed with engines and bridges."""

from __future__ import annotations

import abc
import itertools
import sys
from collections.abc import Iterable, Sequence
from pathlib import Path
from typing import Any, Generic

import datasets
import pydantic

from sieves.data import Doc
from sieves.engines import EngineInferenceMode, EngineType  # noqa: F401
from sieves.engines.types import GenerationSettings
from sieves.engines.utils import init_engine
from sieves.serialization import Config
from sieves.tasks.core import Task
from sieves.tasks.postprocessing.distillation.types import DistillationFramework
from sieves.tasks.predictive.bridges import TaskBridge, TaskPromptSignature, TaskResult
from sieves.tasks.types import Model


class PredictiveTask(
    Generic[TaskPromptSignature, TaskResult, TaskBridge],
    Task,
    abc.ABC,
):
    """Base class for predictive tasks."""

    def __init__(
        self,
        model: Model,
        task_id: str | None,
        include_meta: bool,
        batch_size: int,
        overwrite: bool,
        prompt_template: str | None,
        prompt_signature_desc: str | None,
        fewshot_examples: Iterable[pydantic.BaseModel],
        generation_settings: GenerationSettings,
    ):
        """Initialize PredictiveTask.

        :param task_id: Task ID.
        :param model: Model to use.
        :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_template: Custom prompt template. If None, 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 generation_settings: Settings for structured generation.
        """
        super().__init__(task_id=task_id, include_meta=include_meta, batch_size=batch_size)

        self._engine = init_engine(model, generation_settings)
        self._overwrite = overwrite
        self._custom_prompt_template = prompt_template
        self._custom_prompt_signature_desc = prompt_signature_desc
        self._bridge = self._init_bridge(EngineType.get_engine_type(self._engine))
        self._fewshot_examples = fewshot_examples

        self._validate_fewshot_examples()

    def _validate_fewshot_examples(self) -> None:
        """Validate few‑shot examples.

        :raises ValueError: if fewshot examples don't pass validation.
        """
        pass

    @abc.abstractmethod
    def _init_bridge(self, engine_type: EngineType) -> TaskBridge:
        """Initialize bridge.

        :param engine_type: Type of engine to initialize bridge for.
        :return _TaskBridge: Engine task bridge.
        """

    @property
    @abc.abstractmethod
    def supports(self) -> set[EngineType]:
        """Return supported engine types.

        :return set[EngineType]: Supported engine types.
        """

    @property
    def prompt_template(self) -> str | None:
        """Return prompt template.

        :return str | None: Prompt template.
        """
        prompt_template = self._bridge.prompt_template
        assert prompt_template is None or isinstance(prompt_template, str)
        return prompt_template

    @property
    def prompt_signature_description(self) -> str | None:
        """Return prompt signature description.

        :return str | None: Prompt signature description.
        """
        sig_desc = self._bridge.prompt_signature_description
        assert sig_desc is None or isinstance(sig_desc, str)
        return sig_desc

    def __call__(self, docs: Iterable[Doc]) -> Iterable[Doc]:
        """Execute the task on a set of documents.

        :param docs: Documents to process.
        :return Iterable[Doc]: Processed documents.
        """
        # 1. Compile expected prompt signatures.
        signature = self._bridge.prompt_signature

        # 2. Build executable.
        executable = self._engine.build_executable(
            inference_mode=self._bridge.inference_mode,
            prompt_template=self.prompt_template,
            prompt_signature=signature,
            fewshot_examples=self._fewshot_examples,
        )

        # Compute batch-wise results.
        batch_size = self._batch_size if self._batch_size > 0 else sys.maxsize
        while docs_batch := [doc for doc in itertools.islice(docs, batch_size)]:
            if len(docs_batch) == 0:
                break

            # 3. Extract values from docs to inject/render those into prompt templates.
            docs_values = list(self._bridge.extract(docs_batch))
            assert len(docs_values) == len(docs_batch)

            # 4. Map extracted docs values onto chunks.
            docs_chunks_offsets: list[tuple[int, int]] = []
            docs_chunks: list[dict[str, Any]] = []
            for doc, doc_values in zip(docs_batch, docs_values):
                assert doc.text
                doc_chunks_values = [doc_values | {"text": chunk} for chunk in (doc.chunks or [doc.text])]
                docs_chunks_offsets.append((len(docs_chunks), len(docs_chunks) + len(doc_chunks_values)))
                docs_chunks.extend(doc_chunks_values)

            # 5. Execute prompts per chunk.
            results = list(executable(docs_chunks))
            assert len(results) == len(docs_chunks)

            # 6. Consolidate chunk results.
            results = list(self._bridge.consolidate(results, docs_chunks_offsets))
            assert len(results) == len(docs_batch)

            # 7. Integrate results into docs.
            docs_batch = self._bridge.integrate(results, docs_batch)

            yield from docs_batch

    @property
    def _state(self) -> dict[str, Any]:
        return {
            **super()._state,
            "model": self._engine.model,
            "generation_settings": self._engine.generation_settings.model_dump(),
            "prompt_template": self._custom_prompt_template,
            "prompt_signature_desc": self._custom_prompt_signature_desc,
            "fewshot_examples": self._fewshot_examples,
        }

    @classmethod
    def deserialize(
        cls, config: Config, **kwargs: dict[str, Any]
    ) -> PredictiveTask[TaskPromptSignature, TaskResult, TaskBridge]:
        """Generate PredictiveTask instance from config.

        :param config: Config to generate instance from.
        :param kwargs: Values to inject into loaded config.
        :return PredictiveTask[_TaskPromptSignature, _TaskResult, _TaskBridge]: Deserialized PredictiveTask instance.
        """
        init_dict = config.to_init_dict(cls, **kwargs)
        init_dict["generation_settings"] = GenerationSettings.model_validate(init_dict["generation_settings"])

        return cls(**init_dict)

    @abc.abstractmethod
    def to_hf_dataset(self, docs: Iterable[Doc], threshold: float = 0.5) -> datasets.Dataset:
        """Create Hugging Face datasets.Dataset from docs.

        :param docs: Docs to convert.
        :param threshold: Threshold to apply when converting logits/confidence scores into labels or other structured
            predictions.
        :return datasets.Dataset: Hugging Face dataset.
        """

    @abc.abstractmethod
    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:
        """Distill a model for this task.

        Doc instances must have `.results[task_id]` - otherwise this terminates with an error.

        This method fine-tunes a base model using distillation techniques based on the provided framework. It splits
        the input dataset, trains the model, and saves the resulting model and metadata to the specified output path.

        :param base_model_id: ID of Hugging Face model to use as base for distillation. The chosen model will be
            fine-tuned on the target task's results.
        :param framework: Which distillation framework to use.
        :param data: Docs to extract results from.
        :param output_path: Path to store distilled model and training metadata at.
        :param val_frac: Fraction of data to use for validation set.
        :param seed: RNG seed.
        :param init_kwargs: Kwargs passed on to model/trainer initialization.
        :param train_kwargs: Kwargs passed on to training call.
        :raises KeyError: If expected columns don't exist in `hf_dataset`.
        """

    @staticmethod
    def _split_dataset(
        hf_dataset: datasets.Dataset, train_frac: float, val_frac: float, seed: int | None
    ) -> datasets.DatasetDict:
        """Split dataset.

        :param hf_dataset: Dataset to split.
        :param train_frac: Fractions for training set. `train_frac` + `val_frac` must sum up to 1.
        :param val_frac: Fractions for validation set. `train_frac` + `val_frac` must sum up to 1.
        :param seed: RNG seed.
        :return: Train, val sets; mapping of rows to sets.
        :raises ValueError: If fractions don't sum up to 1.
        """
        if not 0 < val_frac < 1:
            raise ValueError(f"`val_frac` must be greater than 0 and less than 1, but got {val_frac}.")
        if not abs(train_frac + val_frac - 1.0) < 1e-9:
            raise ValueError(f"Split fractions must sum to 1.0, but got {train_frac}, {val_frac}.")

        if val_frac > 0:
            train_val_dataset = hf_dataset.train_test_split(test_size=val_frac, shuffle=True, seed=seed)
            return datasets.DatasetDict({"train": train_val_dataset["train"], "val": train_val_dataset["test"]})

        return datasets.DatasetDict({"train": hf_dataset})
