"""Base class for predictive tasks composed with model wrappers and bridges."""

from __future__ import annotations

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

import datasets
import dspy
import pydantic

from sieves.data import Doc
from sieves.model_wrappers import ModelType, ModelWrapperInferenceMode  # noqa: F401
from sieves.model_wrappers.types import ModelSettings, TokenUsage
from sieves.model_wrappers.utils import init_model_wrapper
from sieves.serialization import Config
from sieves.tasks import optimization
from sieves.tasks.core import Task
from sieves.tasks.distillation.types import DistillationFramework
from sieves.tasks.predictive.evaluation import TaskEvaluationReport
from sieves.tasks.predictive.schemas.core import (
    EvaluationSignature,
    FewshotExample,
    TaskBridge,
    TaskPromptSignature,
    TaskResult,
)
from sieves.tasks.types import Model


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

    # Threshold used for evaluation and any other situatoin in which we need to convert scores into a binary
    # representation.
    THRESHOLD = 0.5

    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],
        model_settings: ModelSettings,
        condition: Callable[[Doc], bool] | None = None,
    ):
        """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 model_settings: Settings for structured generation. Use the `inference_mode` field to specify the
            inference mode for the model wrapper. If not provided, the model wrapper will use its default mode.
        :param condition: Optional callable that determines whether to process each document.
        """
        super().__init__(task_id=task_id, include_meta=include_meta, batch_size=batch_size, condition=condition)

        self._model_wrapper = init_model_wrapper(model, model_settings)
        self._model_settings = model_settings
        self._overwrite = overwrite
        self._custom_prompt_instructions = prompt_instructions
        self._fewshot_examples = fewshot_examples
        self._bridge = self._init_bridge(ModelType.get_model_type(self._model_wrapper))

        self._validate_fewshot_examples()

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

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

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

        :return: Few-shot example type.
        """

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

        :param model_type: Type of model to initialize bridge for.
        :return _TaskBridge: ModelWrapper task bridge.
        """

    @staticmethod
    @abc.abstractmethod
    def supports() -> set[ModelType]:
        """Return supported model types.

        :return set[ModelType]: Supported model 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.
        """
        assert hasattr(self._bridge, "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: Processed documents as an iterable.
        """
        # 1. Compile expected prompt signatures.
        signature = self._bridge.prompt_signature

        # 2. Build executable.
        executable = self._model_wrapper.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.
        docs = iter(docs)
        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 = 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_raw = executable(docs_chunks)
            assert len(results_raw) == len(docs_chunks)

            # Extract metadata.
            results, raw_outputs, usages = (list(t) for t in zip(*results_raw))

            # 6. Consolidate chunk results.
            results = 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)

            # 8. Integrate raw outputs and token usage into docs.
            self._integrate_usage(docs_batch, usages, docs_chunks_offsets, raw_outputs)

            yield from docs_batch

    def _integrate_usage(
        self,
        docs_batch: list[Doc],
        usages: Sequence[TokenUsage],
        docs_chunks_offsets: list[tuple[int, int]],
        raw_outputs: Sequence[Any],
    ) -> None:
        """Integrate raw outputs and token usage into document instances.

        :param docs_batch: Current batch of documents.
        :param usages: Token usage per chunk.
        :param docs_chunks_offsets: Mapping of chunks to documents.
        :param raw_outputs: Raw model outputs per chunk.
        """
        for i, (start, end) in enumerate(docs_chunks_offsets):
            doc = docs_batch[i]
            doc_usages = usages[start:end]

            # Aggregate usage for this task.
            input_tokens_list = [u.input_tokens for u in doc_usages if u.input_tokens is not None]
            output_tokens_list = [u.output_tokens for u in doc_usages if u.output_tokens is not None]

            task_input_tokens = sum(input_tokens_list) if input_tokens_list else None
            task_output_tokens = sum(output_tokens_list) if output_tokens_list else None

            if self._include_meta:
                doc.meta[self._task_id] = {
                    "raw": raw_outputs[start:end],
                    "usage": {
                        "input_tokens": task_input_tokens,
                        "output_tokens": task_output_tokens,
                        # Include per-chunk usage in raw meta.
                        "chunks": [u.model_dump() for u in doc_usages],
                    },
                }

            # Update top-level usage.
            usage = doc.meta.get("usage", {"input_tokens": None, "output_tokens": None})
            if task_input_tokens is not None:
                current_input = usage.get("input_tokens")
                usage["input_tokens"] = (current_input or 0) + task_input_tokens
            if task_output_tokens is not None:
                current_output = usage.get("output_tokens")
                usage["output_tokens"] = (current_output or 0) + task_output_tokens

            doc.meta["usage"] = usage

    @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._model_wrapper.model,
            "model_settings": self._model_wrapper.model_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["model_settings"] = ModelSettings.model_validate(init_dict["model_settings"])

        return cls(**init_dict)

    @property
    @abc.abstractmethod
    def prompt_signature(self) -> type[pydantic.BaseModel]:
        """Return the Pydantic prompt signature for this task.

        :return: Pydantic model class.
        """

    @property
    def metric(self) -> str:
        """Return metric name.

        :return str: Metric name.
        """
        return "Similarity (LLM-judged)"

    def evaluate(
        self, docs: Iterable[Doc], judge: dspy.LM | None = None, failure_threshold: float = 0.5
    ) -> TaskEvaluationReport:
        """Evaluate task performance using DSPy-based evaluation.

        :param docs: Documents to evaluate.
        :param judge: Optional DSPy LM instance to use as judge for generative tasks.
        :param failure_threshold: Decision threshold for whether to mark predicitions as failures.
        :return: Evaluation report.
        """
        truths: list[Any] = []
        preds: list[Any] = []
        failures: list[Doc] = []

        # Evaluate each doc individually to identify failed predictions.
        for doc in docs:
            if self.id not in doc.results:
                continue

            pred = doc.results[self.id]
            gold = doc.gold.get(self.id, None)

            # Accumulate for corpus-level metrics.
            truths.append(gold)
            preds.append(pred)

            # If gold or prediction is None: we cannot do proper evalution, so we just check whether they're both None
            # to compute score for failure analysis.
            if gold is None or pred is None:
                if gold is not None or pred is not None:
                    failures.append(doc)
            else:
                # Convert result and gold to DSPy representation.
                truth = dspy.Example(**self._task_result_to_dspy_dict(gold))
                pred_dspy = dspy.Prediction(**self._task_result_to_dspy_dict(pred))

                # Call internal evaluation logic for per-doc failure analysis.
                score = self._evaluate_dspy_example(truth, pred_dspy, trace=None, model=judge)

                if score < failure_threshold:
                    failures.append(doc)

        # Evaluate on corpus level to obtain representative metrics.
        metrics = self._compute_metrics(truths, preds, judge=judge)

        return TaskEvaluationReport(
            metrics=metrics,
            task_id=self.id,
            failures=failures,
        )

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

        By default, this computes the average of `_evaluate_dspy_example` over all documents.

        :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.
        """
        scores: list[float] = []
        for gold, pred in zip(truths, preds):
            if gold is None or pred is None:
                scores.append(1.0 if gold is None and pred is None else 0.0)
            else:
                truth = dspy.Example(**self._task_result_to_dspy_dict(gold))
                pred_dspy = dspy.Prediction(**self._task_result_to_dspy_dict(pred))
                scores.append(self._evaluate_dspy_example(truth, pred_dspy, trace=None, model=judge))

        avg_score = sum(scores) / len(scores) if scores else 0.0
        return {self.metric: avg_score}

    def _task_result_to_dspy_dict(self, result: Any) -> dict[str, Any]:
        """Convert a task result or gold data to a dictionary suitable for DSPy evaluation.

        :param result: Task result or gold data.
        :return: Dictionary representation.
        """
        if hasattr(result, "model_dump"):
            return result.model_dump(serialize_as_any=True)
        if isinstance(result, dict):
            return result
        return {"output": result}

    @abc.abstractmethod
    def to_hf_dataset(self, docs: Iterable[Doc], threshold: float | None = None) -> 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. Defaults to `THRESHOLD`.
        :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._model_wrapper == ModelType.dspy else self._init_bridge(ModelType.dspy)
            assert hasattr(dspy_bridge, "prompt_signature")
            assert issubclass(dspy_bridge.prompt_signature, dspy.Signature | dspy.Module)
            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_dspy_example(self, truth: dspy.Example, pred: dspy.Prediction, trace: Any, model: dspy.LM) -> float:
        """Evaluate DSPy example for optimization or evaluation.

        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 or no judge model provided.
        """
        if model is None:
            raise ValueError(
                f"No judge model provided for task {self.__class__.__name__}. "
                "Generative tasks require a DSPy LM judge for evaluation."
            )

        if self._fewshot_examples:
            target_fields = list(self._fewshot_examples[0].target_fields)
        else:
            # Fallback to all fields in truth that are not input fields.
            target_fields = [k for k in truth.keys() if k not in ("text",)]

        # 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]

        def _pred_eval(truth: dspy.Example, pred: dspy.Prediction, trace: Any | None = None) -> float:
            """Wrap optimization evaluation, inject model.

            :param truth: Ground truth.
            :param pred: Predicted value.
            :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.
            """
            return self._evaluate_dspy_example(truth, pred, trace, 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(ModelType.get_model_type(self._model_wrapper))

        return best_prompt, self._fewshot_examples
