app.py

import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
from IndicTransToolkit.processor import IndicProcessor
import gradio as gr
import requests
from datetime import datetime
import tempfile
from gtts import gTTS
import os

# Supabase configuration
SUPABASE_URL = "https://gptmdbhzblfybdnohqnh.supabase.co"
SUPABASE_API_KEY = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZSIsInJlZiI6ImdwdG1kYmh6YmxmeWJkbm9ocW5oIiwicm9sZSI6ImFub24iLCJpYXQiOjE3NDc0NjY1NDgsImV4cCI6MjA2MzA0MjU0OH0.CfWArts6Kd_x7Wj0a_nAyGJfrFt8F7Wdy_MdYDj9e7U"
SUPABASE_TABLE = "translations"

# Device configuration
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Load translation models
model_en_to_indic = AutoModelForSeq2SeqLM.from_pretrained("ai4bharat/indictrans2-en-indic-1B", trust_remote_code=True).to(DEVICE)
tokenizer_en_to_indic = AutoTokenizer.from_pretrained("ai4bharat/indictrans2-en-indic-1B", trust_remote_code=True)
model_indic_to_en = AutoModelForSeq2SeqLM.from_pretrained("ai4bharat/indictrans2-indic-en-1B", trust_remote_code=True).to(DEVICE)
tokenizer_indic_to_en = AutoTokenizer.from_pretrained("ai4bharat/indictrans2-indic-en-1B", trust_remote_code=True)
ip = IndicProcessor(inference=True)

# Whisper STT pipeline (keep as is)
asr = pipeline("automatic-speech-recognition", model="openai/whisper-small")


# Save to Supabase
def save_to_supabase(input_text, output_text, direction):
    if not input_text.strip() or not output_text.strip():
        return "Nothing to save."

    table_name = "translations" if direction == "en_to_ks" else "ks_to_en_translations"

    payload = {
        "timestamp": datetime.utcnow().isoformat(),
        "input_text": input_text,
        "output_text": output_text
    }

    headers = {
        "apikey": SUPABASE_API_KEY,
        "Authorization": f"Bearer {SUPABASE_API_KEY}",
        "Content-Type": "application/json"
    }

    try:
        response = requests.post(
            f"{SUPABASE_URL}/rest/v1/{table_name}",
            headers=headers,
            json=payload,
            timeout=10
        )
        return "✅ Saved successfully!" if response.status_code == 201 else "❌ Failed to save."
    except Exception as e:
        print("SAVE EXCEPTION:", e)
        return "❌ Save request error."

# Fetch translation history
def get_translation_history(direction="en_to_ks"):
    table_name = "translations" if direction == "en_to_ks" else "ks_to_en_translations"

    headers = {
        "apikey": SUPABASE_API_KEY,
        "Authorization": f"Bearer {SUPABASE_API_KEY}"
    }

    try:
        response = requests.get(
            f"{SUPABASE_URL}/rest/v1/{table_name}?order=timestamp.desc&limit=10",
            headers=headers,
            timeout=10
        )
        if response.status_code == 200:
            records = response.json()
            return "\n\n".join(
                [f"Input: {r['input_text']} → Output: {r['output_text']}" for r in records]
            )
        return "Failed to load history."
    except Exception as e:
        print("HISTORY FETCH ERROR:", e)
        return "Error loading history."

# Translation function
def translate(text, direction):
    if not text.strip():
        return "Please enter some text.", gr.update(), gr.update()

    if direction == "en_to_ks":
        src_lang = "eng_Latn"
        tgt_lang = "kas_Arab"
        model = model_en_to_indic
        tokenizer = tokenizer_en_to_indic
    else:
        src_lang = "kas_Arab"
        tgt_lang = "eng_Latn"
        model = model_indic_to_en
        tokenizer = tokenizer_indic_to_en

    try:
        processed = ip.preprocess_batch([text], src_lang=src_lang, tgt_lang=tgt_lang)
        batch = tokenizer(processed, return_tensors="pt", padding=True).to(DEVICE)

        with torch.no_grad():
            outputs = model.generate(
                **batch,
                max_length=256,
                num_beams=5,
                num_return_sequences=1
            )

        translated = tokenizer.batch_decode(outputs, skip_special_tokens=True)
        result = ip.postprocess_batch(translated, lang=tgt_lang)[0]

        return result, gr.update(), gr.update()
    except Exception as e:
        print("Translation Error:", e)
        return "⚠️ Translation failed.", gr.update(), gr.update()

# Transcribe English audio
def transcribe_audio(audio_path):
    try:
        result = asr(audio_path)
        return result["text"]
    except Exception as e:
        print("STT Error:", e)
        return "⚠️ Transcription failed."

# Synthesize English TTS using gTTS for ks_to_en direction
def synthesize_tts(text, direction):
    if direction == "ks_to_en" and text.strip():
        try:
            tts = gTTS(text=text, lang="en")
            tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
            tts.save(tmp_file.name)
            tmp_file.close()
            return tmp_file.name
        except Exception as e:
            print("TTS Error:", e)
    return None

# Direction switch
def switch_direction(direction, input_text_val, output_text_val):
    new_direction = "ks_to_en" if direction == "en_to_ks" else "en_to_ks"
    input_label = "Kashmiri Text" if new_direction == "ks_to_en" else "English Text"
    output_label = "English Translation" if new_direction == "ks_to_en" else "Kashmiri Translation"
    return (
        new_direction,
        gr.update(value=output_text_val, label=input_label),
        gr.update(value=input_text_val, label=output_label)
    )

# Gradio interface
with gr.Blocks() as interface:
    gr.HTML("""
    <div style="display: flex; justify-content: space-between; align-items: center; padding: 10px;">
        <img src="https://raw.githubusercontent.com/BurhaanRasheedZargar/Images/211321a234613a9c3dd944fe9367cf13d1386239/assets/left_logo.png" style="height:150px; width:auto;">
        <h2 style="margin: 0; text-align: center;">English ↔ Kashmiri Translator</h2>
        <img src="https://raw.githubusercontent.com/BurhaanRasheedZargar/Images/77797f7f7cbee328fa0f9d31cf3e290441e04cd3/assets/right_logo.png">
    </div>
    """)

    translation_direction = gr.State(value="en_to_ks")

    with gr.Row():
        input_text = gr.Textbox(lines=2, label="English Text", placeholder="Enter text....")
        output_text = gr.Textbox(lines=2, label="Kashmiri Translation", placeholder="Translated text....")

    with gr.Row():
        translate_button = gr.Button("Translate")
        save_button = gr.Button("Save Translation")
        switch_button = gr.Button("Switch")

    save_status = gr.Textbox(label="Save Status", interactive=False)
    history_box = gr.Textbox(lines=10, label="Translation History", interactive=False)

    with gr.Row():
        audio_input = gr.Audio(source="microphone", type="filepath", label="🎙️ Speak in English")
        audio_output = gr.Audio(label="🔊 English Output Audio")
    stt_translate_button = gr.Button("🎤 Transcribe & Translate")

    # Click events
    translate_button.click(
        fn=translate,
        inputs=[input_text, translation_direction],
        outputs=[output_text, input_text, output_text]
    )

    save_button.click(
        fn=save_to_supabase,
        inputs=[input_text, output_text, translation_direction],
        outputs=save_status
    ).then(
        fn=get_translation_history,
        inputs=translation_direction,
        outputs=history_box
    )

    switch_button.click(
        fn=switch_direction,
        inputs=[translation_direction, input_text, output_text],
        outputs=[translation_direction, input_text, output_text]
    )

    stt_translate_button.click(
        fn=transcribe_audio,
        inputs=audio_input,
        outputs=input_text
    ).then(
        fn=translate,
        inputs=[input_text, translation_direction],
        outputs=[output_text, input_text, output_text]
    ).then(
        fn=synthesize_tts,
        inputs=[output_text, translation_direction],
        outputs=audio_output
    )

if __name__ == "__main__":
    interface.queue().launch()