from abc import ABC, abstractmethod
from typing import ClassVar

from pydantic import BaseModel, Field

from yaduha.tool import Tool


class BackTranslation(BaseModel):
    source: str = Field(..., description="The back translated source-language text.")
    target: str = Field(..., description="The original target-language text.")
    translation_time: float = Field(..., description="The time taken for back translation.")
    prompt_tokens: int = Field(
        0, description="The number of prompt tokens used for back translation."
    )
    completion_tokens: int = Field(
        0, description="The number of completion tokens used for back translation."
    )


class Translation(BaseModel):
    source: str = Field(..., description="The source-language text.")
    target: str = Field(..., description="The target-language text.")
    translation_time: float = Field(..., description="The time taken for translation.")
    prompt_tokens: int = Field(
        0, description="The number of prompt tokens used for the entire translation."
    )
    completion_tokens: int = Field(
        0, description="The number of completion tokens used for the entire translation."
    )
    back_translation: BackTranslation | None = Field(
        None, description="The back translation details, if available."
    )
    evaluations: dict[str, float] = Field(
        default_factory=dict, description="Evaluation scores keyed by evaluator name."
    )


class Translator(Tool[Translation], ABC):
    """Base class for translators that translate text to a target language and back to the source language."""

    name: ClassVar[str] = "translator"
    description: ClassVar[str] = (
        "Translate text to the target language and back to the source language."
    )

    def _run(self, text: str) -> Translation:
        """Translate the text to the target language and back to the source language.

        Args:
            text (str): The text to translate.
        Returns:
            Translation: The translation
        """
        return self.translate(text)

    @abstractmethod
    def translate(self, text: str) -> Translation:
        pass
