from transformers import pipeline, MarianMTModel, MarianTokenizer
from typing import Optional

# Load English->French pipeline once (using HuggingFace translation pipeline)
_en_fr_translator = pipeline("translation_en_to_fr", model="Helsinki-NLP/opus-mt-en-fr")

# Load English->Spanish MarianMT model and tokenizer once
_es_model_name = "Helsinki-NLP/opus-mt-en-es"
_es_tokenizer = MarianTokenizer.from_pretrained(_es_model_name)
_es_model = MarianMTModel.from_pretrained(_es_model_name)

def translate_to_french(text: str) -> str:
    """Translate English text to French using HuggingFace pipeline."""
    try:
        result = _en_fr_translator(text)
        return result[0]['translation_text']
    except Exception as e:
        return f"Translation error (EN->FR): {str(e)}"

def translate_to_spanish(text: str) -> str:
    """Translate English text to Spanish using MarianMT."""
    try:
        inputs = _es_tokenizer(text, return_tensors="pt", padding=True)
        translated_tokens = _es_model.generate(**inputs)
        return _es_tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
    except Exception as e:
        return f"Translation error (EN->ES): {str(e)}"

def translate(text: str, to_lang: Optional[str] = "fr") -> str:
    """
    General translate function.
    Supported: 'fr' (French), 'es' (Spanish).
    Defaults to French.
    """
    if to_lang == "fr":
        return translate_to_french(text)
    elif to_lang == "es":
        return translate_to_spanish(text)
    else:
        return f"Unsupported target language '{to_lang}'. Supported: 'fr', 'es'."