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

# 1. BRAIN CONFIGURATION: Bi-directional Model Mapping
# 'to_indic' for English -> Language
# 'to_english' for Language -> English
CONFIG = {
    "Hindi": {
        "to_indic": "Helsinki-NLP/opus-mt-en-hi", 
        "to_english": "Helsinki-NLP/opus-mt-hi-en", 
        "v_code": "hi", "prefix": ""
    },
    "Tamil": {
        "to_indic": "Helsinki-NLP/opus-mt-en-dra", 
        "to_english": "Helsinki-NLP/opus-mt-dra-en", 
        "v_code": "ta", "prefix": ">>tam<< "
    },
    "Telugu": {
        "to_indic": "Helsinki-NLP/opus-mt-en-dra", 
        "to_english": "Helsinki-NLP/opus-mt-dra-en", 
        "v_code": "te", "prefix": ">>tel<< "
    },
    "Gujarati": {
        "to_indic": "Helsinki-NLP/opus-mt-en-inc", 
        "to_english": "Helsinki-NLP/opus-mt-inc-en", 
        "v_code": "gu", "prefix": ">>guj<< "
    },
    "Marathi": {
        "to_indic": "Helsinki-NLP/opus-mt-en-inc", 
        "to_english": "Helsinki-NLP/opus-mt-inc-en", 
        "v_code": "mr", "prefix": ">>mar<< "
    },
    "Assamese": {
        "to_indic": "Helsinki-NLP/opus-mt-en-inc", 
        "to_english": "Helsinki-NLP/opus-mt-inc-en", 
        "v_code": "as", "prefix": ">>asm<< "
    }
}

# Cache for models
loaded_models = {}
# Load Whisper for Speech-to-Text (The "Ears")
stt_model = whisper.load_model("base") 

def get_model(model_name):
    if model_name not in loaded_models:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
        loaded_models[model_name] = (tokenizer, model, device)
    return loaded_models[model_name]

# --- MODE 1: YOU TO THEM (Text -> Translation -> Voice) ---
def you_to_them(text, language):
    if not text.strip(): return "Please enter text...", None
    cfg = CONFIG[language]
    tokenizer, model, device = get_model(cfg["to_indic"])
    
    input_text = cfg["prefix"] + text
    inputs = tokenizer(input_text, return_tensors="pt", padding=True).to(device)
    with torch.no_grad():
        tokens = model.generate(**inputs, max_length=200)
    
    translation = tokenizer.decode(tokens[0], skip_special_tokens=True)
    tts = gTTS(translation, lang=cfg["v_code"])
    tts.save("outgoing.mp3")
    return translation, "outgoing.mp3"

# --- MODE 2: THEM TO YOU (Microphone -> Transcription -> English) ---
def them_to_you(audio_path, language):
    if audio_path is None: return "Please record audio first."
    
    # 1. Speech to Text (Whisper)
    result = stt_model.transcribe(audio_path)
    spoken_text = result['text']
    
    # 2. Text to English Translation
    cfg = CONFIG[language]
    tokenizer, model, device = get_model(cfg["to_english"])
    
    inputs = tokenizer(spoken_text, return_tensors="pt", padding=True).to(device)
    with torch.no_grad():
        tokens = model.generate(**inputs, max_length=200)
    
    english_text = tokenizer.decode(tokens[0], skip_special_tokens=True)
    return f"They Spoke ({language}): {spoken_text}\n\nEnglish Meaning: {english_text}"

# --- ARTISTIC INTERFACE ---
artistic_theme = gr.themes.Soft(primary_hue="purple", neutral_hue="slate").set(
    body_background_fill='*neutral_950',
    block_background_fill='*neutral_900',
    block_radius='20px',
    button_primary_background_fill='linear-gradient(90deg, #9333ea 0%, #3b82f6 100%)',
    button_primary_text_color='white',
    block_shadow='0 0 25px rgba(147, 51, 234, 0.2)'
)

with gr.Blocks(theme=artistic_theme, title="Polyglot Prophet v3") as demo:
    gr.Markdown("# <center>🌌 The Polyglot Prophet: Conversation Suite</center>")
    gr.Markdown("<center>A professional bidirectional translation engine for the Indian subcontinent.</center>")
    
    with gr.Tabs():
        # TAB 1: YOU TO THEM
        with gr.TabItem("1. You ➔ Them (Voice Out)"):
            with gr.Row():
                with gr.Column():
                    in_text = gr.Textbox(label="Type your English message", lines=4)
                    target_lang = gr.Dropdown(choices=list(CONFIG.keys()), value="Hindi", label="Their Language")
                    btn_send = gr.Button("Awaken Translation", variant="primary")
                with gr.Column():
                    out_text = gr.Textbox(label="Translated Text", interactive=False)
                    out_audio = gr.Audio(label="Spoken Revelation", autoplay=True)
            btn_send.click(you_to_them, inputs=[in_text, target_lang], outputs=[out_text, out_audio])

        # TAB 2: THEM TO YOU
        with gr.TabItem("2. Them ➔ You (Listen & Decode)"):
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### Ask the other person to speak.")
                    mic_input = gr.Audio(sources=["microphone"], type="filepath", label="Record their voice")
                    source_lang = gr.Dropdown(choices=list(CONFIG.keys()), value="Hindi", label="Language they are speaking")
                    btn_listen = gr.Button("Decode to English", variant="primary")
                with gr.Column():
                    final_transcript = gr.Textbox(label="Prophet's Interpretation", lines=6, interactive=False)
            btn_listen.click(them_to_you, inputs=[mic_input, source_lang], outputs=final_transcript)

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