import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import torch
from gtts import gTTS
import tempfile
import os

models = {
    "English to Hindi": "Helsinki-NLP/opus-mt-en-hi",
    "Hindi to English": "Helsinki-NLP/opus-mt-hi-en",
    "Hindi to Sanskrit": "ai4bharat/indictrans-hin-san",
    "English to Sanskrit": "ai4bharat/indictrans-en-san"
}

loaded_models = {}

def load_model(direction):
    if direction not in loaded_models:
        tokenizer = AutoTokenizer.from_pretrained(models[direction])
        model = AutoModelForSeq2SeqLM.from_pretrained(models[direction])
        loaded_models[direction] = (tokenizer, model)
    return loaded_models[direction]

chat_history = []

def translate(text, direction):
    tokenizer, model = load_model(direction)
    inputs = tokenizer(text, return_tensors="pt", padding=True)
    outputs = model.generate(**inputs)
    translated = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # memory
    chat_history.append((text, translated))
    return translated

def text_to_speech(text):
    tts = gTTS(text)
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
    tts.save(temp_file.name)
    return temp_file.name

with gr.Blocks() as demo:
    gr.Markdown("## 🌐 Translator Chatbot with Voice")

    with gr.Row():
        input_text = gr.Textbox(label="Enter Text", lines=2)
        direction = gr.Dropdown(list(models.keys()), label="Translate Direction")
        output_text = gr.Textbox(label="Translated Text", lines=2)
    
    translate_btn = gr.Button("🔁 Translate")
    tts_btn = gr.Button("🔊 Play Voice")
    audio_output = gr.Audio()

    translate_btn.click(fn=translate, inputs=[