"""Question Answering 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.question_answering.bridges import DSPyQuestionAnswering, PydanticQA
from sieves.tasks.predictive.schemas.question_answering import (
    FewshotExample,
    TaskModel,
    TaskPromptSignature,
    TaskResult,
)

_TaskBridge = DSPyQuestionAnswering | PydanticQA


class QuestionAnswering(PredictiveTask[TaskPromptSignature, TaskResult, _TaskBridge]):
    """Answer questions about a text using structured model wrappers."""

    def __init__(
        self,
        questions: list[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] = (),
        model_settings: ModelSettings = ModelSettings(),
        condition: Callable[[Doc], bool] | None = None,
    ) -> None:
        """
        Initialize QuestionAnswering task.

        :param questions: Questions to answer.
        :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 model_settings: Settings for structured generation.
        :param condition: Optional callable that determines whether to process each document.
        """
        self._questions = questions
        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.
        """
        return FewshotExample

    @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)

    def _validate_fewshot_examples(self) -> None:
        """Validate that questions, answers and scores have the same length.

        :raises ValueError: If lengths don't match.
        """
        for i, example in enumerate(self._fewshot_examples):
            if len(example.questions) != len(example.answers):
                raise ValueError(
                    f"Length mismatch in few-shot example {i}: {len(example.questions)} questions "
                    f"vs {len(example.answers)} answers."
                )
            if example.scores is not None and len(example.scores) != len(example.answers):
                raise ValueError(
                    f"Length mismatch in few-shot example {i}: {len(example.answers)} answers "
                    f"vs {len(example.scores)} scores."
                )

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

        try:
            return bridge_types[model_type](
                task_id=self._task_id,
                prompt_instructions=self._custom_prompt_instructions,
                questions=self._questions,
                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 prompt_signature(self) -> type[pydantic.BaseModel]:
        return TaskResult

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

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

        # Fetch data used for generating dataset.
        try:
            data: list[tuple[str, list[str], list[float]]] = []
            for doc in docs:
                result = doc.results[self._task_id]
                answers = [qa.answer for qa in result.qa_pairs]
                scores = [qa.score if qa.score is not None else 1.0 for qa in result.qa_pairs]
                data.append((doc.text, answers, scores))
        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, answers, scores in data:
                yield {"text": text, "answers": answers, "scores": scores}

        # 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
