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

from __future__ import annotations

import abc
import functools
import itertools
import json
import logging
import sys
import warnings
from collections.abc import Iterable, Sequence
from pathlib import Path
from typing import Any, Generic, Self

import datasets
import dspy
import pydantic

from sieves.data import Doc
from sieves.engines import Engine, 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 import optimization
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 EvaluationSignature(dspy.Signature):
    """Evaluate similarity between ground truth and predicted outputs."""

    target_fields: str = dspy.InputField(desc="Names of output fields being compared.")
    ground_truth: str = dspy.InputField(desc="Ground truth output values.")
    prediction: str = dspy.InputField(desc="Predicted output values.")

    reasoning: str = dspy.OutputField(desc="Step-by-step reasoning for the similarity assessment")
    similarity_score: float = dspy.OutputField(
        desc="Similarity score between 0.0 and 1.0, where 1.0 means identical and 0.0 means completely different."
    )


class FewshotExample(pydantic.BaseModel):
    """Few-shot example.

    :params text: Input text.
    """

    text: str

    @property
    def input_fields(self) -> Sequence[str]:
        """Defines which fields are inputs.

        :return: Sequence of field names.
        """
        return ("text",)

    @property
    @abc.abstractmethod
    def target_fields(self) -> Sequence[str]:
        """Define which fields are targets, i.e. the end results the task aims to produce.

        :return: Sequence of field names.
        """

    def to_dspy(self) -> dspy.Example:
        """Convert to `dspy.Example`.

        :returns: Example as `dspy.Example`.
        """
        return dspy.Example(**Engine.convert_fewshot_examples([self])[0]).with_inputs(self.input_fields)

    @classmethod
    def from_dspy(cls, example: dspy.Example) -> Self:
        """Convert from `dspy.Example`.

        :param example: Example as `dspy.Example`.
        :returns: Example as `FewshotExample`.
        """
        return cls(**example)


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_instructions: str | None,
        fewshot_examples: Sequence[FewshotExample],
        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_instructions: Custom prompt instructions. If None, default instructions are 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_instructions = prompt_instructions
        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:
        """Return prompt template.

        :return str | None: Prompt template.
        """
        prompt_template = self._bridge.prompt_template
        assert 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 fewshot_examples(self) -> Sequence[FewshotExample]:
        """Return few-shot examples.

        :return: Few-shot examples.
        """
        return self._fewshot_examples

    @property
    def _state(self) -> dict[str, Any]:
        return {
            **super()._state,
            "model": self._engine.model,
            "generation_settings": self._engine.generation_settings.model_dump(),
            "prompt_instructions": self._custom_prompt_instructions,
            "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})

    def _get_task_signature(self) -> type[dspy.Signature] | type[dspy.Module]:
        """Get DSPy signature for this task.

        By default this uses the task signature of the DSPy bridge for this task. If none is found, this fails. Can be
        overwritten with a custom task signature, if no DSPy bridge will be implemented for this task.

        :return: DSPy signature for this task.
        :raises KeyError: If no DSPy bridge defined for this task.
        """
        try:
            dspy_bridge = self._bridge if self._engine == EngineType.dspy else self._init_bridge(EngineType.dspy)
            return dspy_bridge.prompt_signature

        except KeyError as err:
            raise KeyError(f"DSPy bridge not available for task {self.__class__.__name__}.") from err

    def _evaluate_optimization_example(
        self, truth: dspy.Example, pred: dspy.Prediction, model: dspy.LM, trace: Any | None = None
    ) -> float:
        """Evaluate DSPy example for optimization.

        By default this implements an LLM-based evaluation that uses the model that drives optimization to compare
        results. Where possible this should be replaced by a targeted, deterministic evaluation.

        :param truth: Ground truth.
        :param pred: Predicted value.
        :param model: Model used by optimizer.
        :param trace: Optional trace information.
        :return: Metric value between 0.0 and 1.0.
        :raises KeyError: If target fields are missing from truth or prediction.
        :raises ValueError: If similarity score cannot be parsed from LLM response.
        """
        target_fields = list(self._fewshot_examples[0].target_fields)

        # Filter truth and pred to only include target fields.
        truth_filtered: dict[str, Any] = {}
        for field in target_fields:
            if field not in truth:
                raise KeyError(f"Target field '{field}' missing from ground truth example: {dict(truth)}.")
            truth_filtered[field] = truth[field]

        pred_filtered: dict[str, Any] = {}
        for field in target_fields:
            if field not in pred:
                raise KeyError(f"Target field '{field}' missing from prediction: {dict(pred)}.")
            pred_filtered[field] = pred[field]

        # Serialize to readable format
        target_fields_str = ", ".join(target_fields)
        ground_truth_str = json.dumps(truth_filtered, indent=2, default=str)
        prediction_str = json.dumps(pred_filtered, indent=2, default=str)

        # Create DSPy evaluator module
        evaluator = dspy.Predict(EvaluationSignature)

        # Call evaluator with the model
        with dspy.context(lm=model):
            result = evaluator(
                target_fields=target_fields_str,
                ground_truth=ground_truth_str,
                prediction=prediction_str,
            )

        # Parse and validate score
        try:
            score = float(result.similarity_score)
        except (ValueError, AttributeError, TypeError) as err:
            raise ValueError(
                f"Failed to parse similarity score from LLM response. Got: {getattr(result, 'similarity_score', None)}"
            ) from err

        # Clamp to [0, 1].
        return max(0.0, min(1.0, score))

    def optimize(self, optimizer: optimization.Optimizer, verbose: bool = True) -> tuple[str, Sequence[FewshotExample]]:
        """Optimize task prompt and few-shot examples with the available optimization config.

        Updates task to use best prompt and few-shot examples found by the optimizer.

        :param optimizer: Optimizer to run.
        :param verbose: Whether to suppress output. DSPy produces a good amount of logs, so this can be useful to
            not pollute your terminal. Only warnings and errors will be printed.

        :return tuple[str, Sequence[FewshotExample]]: Best found prompt and few-shot examples.
        """
        assert len(self._fewshot_examples) > 1, "At least two few-shot examples need to be provided to optimize."

        # Run optimizer to get best prompt and few-shot examples.
        signature = self._get_task_signature()
        dspy_examples = [ex.to_dspy() for ex in self._fewshot_examples]
        pred_eval = functools.partial(self._evaluate_optimization_example, model=optimizer.model)

        if verbose:
            best_prompt, best_examples = optimizer(signature, dspy_examples, pred_eval, verbose=verbose)
        else:
            # Temporarily suppress DSPy logs.
            dspy_logger = logging.getLogger("dspy")
            optuna_logger = logging.getLogger("optuna")
            original_dspy_level = dspy_logger.level
            original_optuna_level = optuna_logger.level

            try:
                dspy_logger.setLevel(logging.ERROR)
                optuna_logger.setLevel(logging.ERROR)
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")
                    best_prompt, best_examples = optimizer(signature, dspy_examples, pred_eval, verbose=verbose)
            finally:
                dspy_logger.setLevel(original_dspy_level)
                optuna_logger.setLevel(original_optuna_level)

        # Update few-shot examples and prompt instructions.
        fewshot_example_cls = self._fewshot_examples[0].__class__
        self._fewshot_examples = [fewshot_example_cls.from_dspy(ex) for ex in best_examples]
        self._validate_fewshot_examples()
        self._custom_prompt_instructions = best_prompt

        # Reinitialize bridge to use new prompt and few-shot examples.
        self._bridge = self._init_bridge(EngineType.get_engine_type(self._engine))

        return best_prompt, self._fewshot_examples
