from typing import Any, Optional
from smolagents.tools import Tool

class TranslationTool(Tool):
    """
    Example:

    ```py
    translator = TranslationTool()
    translator("This is a super nice API!", src_lang="English", tgt_lang="French")
    ```
    """
    default_checkpoint = "facebook/nllb-200-distilled-600M"
    description = "This is a tool that translates text from a language to another."
    name = "translator"
    inputs = {'text': {'type': 'string', 'description': 'The text to translate'}, 'src_lang': {'type': 'string', 'description': "The language of the text to translate. Written in plain English, such as 'Romanian', or 'Albanian'"}, 'tgt_lang': {'type': 'string', 'description': "The language for the desired output language. Written in plain English, such as 'Romanian', or 'Albanian'"}}
    output_type = "string"

    def __init__(self, lang_to_code=LANGUAGE_CODES, pre_processor_class=AutoTokenizer, model_class=AutoModelForSeq2SeqLM):
        super().__init__()
        self.lang_to_code = lang_to_code
        self.pre_processor_class = pre_processor_class
        self.model_class = model_class
        # self.pre_processor = self.pre_processor_class.from_pretrained(self.default_checkpoint)
        # self.model = self.model_class.from_pretrained(self.default_checkpoint)
        # self.post_processor = self.pre_processor

    def encode(self, text, src_lang, tgt_lang):
        if src_lang not in self.lang_to_code:
            raise ValueError(f"{src_lang} is not a supported language.")
        if tgt_lang not in self.lang_to_code:
            raise ValueError(f"{tgt_lang} is not a supported language.")
        src_lang = self.lang_to_code[src_lang]
        tgt_lang = self.lang_to_code[tgt_lang]
        return self.pre_processor._build_translation_inputs(
            text, return_tensors="pt", src_lang=src_lang, tgt_lang=tgt_lang
        )

    def forward(self, inputs):
        return self.model.generate(**inputs)

    def decode(self, outputs):
        return self.post_processor.decode(outputs[0].tolist(), skip_special_tokens=True)