import os
import gradio as gr
from langdetect import detect, LangDetectException
from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer

try:
    from groq import Groq
except Exception:
    Groq = None

# Config
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
GROQ_MODEL = os.getenv("GROQ_MODEL", "mixtral-8x7b-32768")

groq_client = None
if GROQ_API_KEY and Groq is not None:
    try:
        groq_client = Groq(api_key=GROQ_API_KEY)
    except Exception:
        pass

m2m_model_name = "facebook/m2m100_418M"
m2m_tokenizer = M2M100Tokenizer.from_pretrained(m2m_model_name)
m2m_model = M2M100ForConditionalGeneration.from_pretrained(m2m_model_name)

LANG_UI_TO_CODE = {"English": "en", "Spanish": "es", "French": "fr"}

def call_m2m(user_text, target_code):
    try:
        src_code = detect(user_text)
    except LangDetectException:
        src_code = "en"
    if src_code == target_code:
        return user_text
    m2m_tokenizer.src_lang = src_code
    encoded = m2m_tokenizer(user_text, return_tensors="pt")
    generated = m2m_model.generate(**encoded, forced_bos_token_id=m2m_tokenizer.get_lang_id(target_code))
    return m2m_tokenizer.decode(generated[0], skip_special_tokens=True)

def translate_text(user_text, target_lang_ui):
    if not user_text.strip():
        return "⚠️ Please enter text."
    target_code = LANG_UI_TO_CODE.get(target_lang_ui, "en")
    try:
        return call_m2m(user_text, target_code)
    except:
        return "❌ Translation failed."

with gr.Blocks() as demo:
    gr.Markdown("## 🌐 Universal Translator")
    with gr.Row():
        txt = gr.Textbox(label="Enter your text", lines=6)
        tgt = gr.Dropdown(choices=["English", "Spanish", "French"], value="English", label="Target Language")
    out = gr.Textbox(label="Translated Output", lines=6)
    btn = gr.Button("Translate")
    btn.click(fn=translate_text, inputs=[txt, tgt], outputs=[out])