import asyncio
from smolagents import Tool
from googletrans import Translator
from typing import Optional


class TranslateTool(Tool):
    name = "google_translate"
    description = """
    Translates text from one language to another using the Google Translate API.
    Supports a wide range of languages using standard language codes (e.g., 'en' for English, 'fr' for French).
    If source_language is not specified, it will be auto-detected. If target_language is not specified, it defaults to English ('en').
    """
    inputs = {
        "text": {
            "type": "string",
            "description": "The text to translate."
        },
        "source_language": {
            "type": "string",
            "description": "The source language code (e.g., 'en' for English, 'fr' for French).",
        },
        "target_language": {
            "type": "string",
            "description": "The target language code (e.g., 'es' for Spanish, 'de' for German).",
        }
    }
    output_type = "string"

    async def translate_text(text: str, src: str, dest: str):
        async with Translator() as translator:
            res = await translator.translate(text, src, dest)
            return res

    def forward(self, text: str,  source_language: str, target_language: str):
        try:
            translator = Translator()
            translation = asyncio.run(translator.translate(text, src=source_language, dest=target_language))
            return translation.text
        except Exception as e:
            return f"Error: Unable to translate due to {e}"