"""Translation 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 pydantic

from sieves.data import Doc
from sieves.engines import EngineType, dspy_, langchain_, outlines_
from sieves.engines.types import GenerationSettings
from sieves.serialization import Config
from sieves.tasks.distillation.types import DistillationFramework
from sieves.tasks.predictive.core import FewshotExample as BaseFewshotExample
from sieves.tasks.predictive.core import PredictiveTask
from sieves.tasks.predictive.translation.bridges import (
    DSPyTranslation,
    LangChainTranslation,
    OutlinesTranslation,
)

_TaskModel = dspy_.Model | langchain_.Model | outlines_.Model
_TaskPromptSignature = pydantic.BaseModel | dspy_.PromptSignature
_TaskResult = outlines_.Result | dspy_.Result
_TaskBridge = DSPyTranslation | LangChainTranslation | OutlinesTranslation


class FewshotExample(BaseFewshotExample):
    """Few-shot example with a target translation."""

    to: str
    translation: str

    @override
    @property
    def input_fields(self) -> Sequence[str]:
        return "text", "to"

    @override
    @property
    def target_fields(self) -> Sequence[str]:
        return ("translation",)


class Translation(PredictiveTask[_TaskPromptSignature, _TaskResult, _TaskBridge]):
    """Translate documents into a target language using structured engines."""

    def __init__(
        self,
        to: str,
        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] = (),
        generation_settings: GenerationSettings = GenerationSettings(),
        condition: Callable[[Doc], bool] | None = None,
    ) -> None:
        """
        Initialize Translation task.

        :param to: Language to translate to.
        :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 generation_settings: Settings for structured generation.
        :param condition: Optional callable that determines whether to process each document.
        """
        self._to = to

        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,
            generation_settings=generation_settings,
            condition=condition,
        )

    @override
    def _init_bridge(self, engine_type: EngineType) -> _TaskBridge:
        bridge_types: dict[EngineType, type[_TaskBridge]] = {
            EngineType.dspy: DSPyTranslation,
            EngineType.langchain: LangChainTranslation,
            EngineType.outlines: OutlinesTranslation,
        }

        try:
            bridge = bridge_types[engine_type](
                task_id=self._task_id,
                prompt_instructions=self._custom_prompt_instructions,
                overwrite=self._overwrite,
                language=self._to,
                generation_settings=self._generation_settings,
            )
        except KeyError as err:
            raise KeyError(f"Engine type {engine_type} is not supported by {self.__class__.__name__}.") from err

        return bridge

    @staticmethod
    @override
    def supports() -> set[EngineType]:
        return {EngineType.dspy, EngineType.langchain, EngineType.outlines}

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

    @override
    def to_hf_dataset(self, docs: Iterable[Doc], threshold: float = 0.5) -> datasets.Dataset:
        # Define metadata.
        features = datasets.Features({"text": datasets.Value("string"), "translation": datasets.Value("string")})
        info = datasets.DatasetInfo(
            description=f"Translation dataset with target language {self._to}."
            f"Generated with sieves v{Config.get_version()}.",
            features=features,
        )

        # Fetch data used for generating dataset.
        try:
            data = [(doc.text, doc.results[self._task_id]) 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, translation in data:
                yield {"text": text, "translation": translation}

        # 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
