# app.py
import os
import json
import requests
import gradio as gr
from langdetect import detect, LangDetectException

# Hugging Face Transformers
from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer

# Groq SDK
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")
HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")

# Init Groq
groq_client = None
if GROQ_API_KEY and Groq is not None:
    try:
        groq_client = Groq(api_key=GROQ_API_KEY)
    except Exception as e:
        print("Groq client init failed:", repr(e))

# Universal translation model
m2m_model_name = "facebook/m2m100_418M"
m2m_tokenizer = M2M100Tokenizer.from_pretrained(m2m_model_name)
m2m_model = M2M100ForConditionalGeneration.from_pretrained(m2m_model_name)

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

SYSTEM_PROMPT = """
You are a multilingual translation assistant.
Task:
1. Detect the input language automatically.
2. Translate into the requested target language.
3. Preserve meaning, tone, and formatting.
4. Keep numbers, symbols, names, and special characters unchanged.
5. If the input is already in the target language, return it unchanged.
"""

def call_groq(user_text, target_lang_ui):
    if not groq_client:
        raise RuntimeError("Groq client not configured")
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"Target language: {target_lang_ui}\n\n{user_text}"},
    ]
    chat = groq_client.chat.completions.create(
        model=GROQ_MODEL,
        messages=messages,
        temperature=0,
        max_tokens=2048,
    )
    try:
        return chat.choices[0].message.content.strip()
    except Exception:
        try:
            return chat["choices"][0]["message"]["content"].strip()
        except Exception as e:
            print("Unexpected Groq response:", repr(e))
            raise

def call_m2m(user_text, target_code):
    try:
        src_code = detect(user_text)
    except LangDetectException:
        src_code = "en"  # fallback

    # if already target language → return as-is
    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):
    user_text = (user_text or "").strip()
    if not user_text:
        return "⚠️ Please enter some text to translate."
    target_code = LANG_UI_TO_CODE.get(target_lang_ui, "en")

    # Try Groq first
    try:
        if groq_client:
            out = call_groq(user_text, target_lang_ui)
            if out:
                return out
    except Exception as e:
        print("Groq call failed:", repr(e))

    # Fallback → M2M100 universal translator
    try:
        return call_m2m(user_text, target_code)
    except Exception as e:
        print("M2M100 translation failed:", repr(e))
        return "❌ Translation failed. Check logs."

# ----------------- Gradio UI -----------------
with gr.Blocks() as demo:
    gr.Markdown("## 🌐 Hackathon Translator (Universal)")

    with gr.Row():
        txt = gr.Textbox(label="Enter your text", lines=6, placeholder="Type or paste text here...")
        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])

if __name__ == "__main__":
    demo.launch()