"""Translation predictive task."""

from __future__ import annotations

from collections.abc import 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.postprocessing.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(),
    ) -> 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.
        """
        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,
        )

    @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,
            )
        except KeyError as err:
            raise KeyError(f"Engine type {engine_type} is not supported by {self.__class__.__name__}.") from err

        return bridge

    @override
    @property
    def supports(self) -> 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
