app.py

import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
import os
import json
import gspread
from oauth2client.service_account import ServiceAccountCredentials
from datetime import datetime
from gtts import gTTS
import tempfile

# --- CONFIGURATION ---
MODEL_K2H_REPO = "ankitklakra/kurukh-to-hindi"
MODEL_H2K_REPO = "ankitklakra/hindi-to-kurukh"
SHEET_NAME = "Kurukh_Feedback_Log"

# --- LOAD TEXT MODELS ---
print("Loading Translation Models...")
tokenizer = AutoTokenizer.from_pretrained("google/mt5-small")
model_k2h = AutoModelForSeq2SeqLM.from_pretrained(MODEL_K2H_REPO)
model_h2k = AutoModelForSeq2SeqLM.from_pretrained(MODEL_H2K_REPO)

pipe_k2h = pipeline("text2text-generation", model=model_k2h, tokenizer=tokenizer)
pipe_h2k = pipeline("text2text-generation", model=model_h2k, tokenizer=tokenizer)

# --- LOAD AUDIO MODEL (WHISPER) ---
print("Loading Voice Model...")

asr_pipeline = pipeline("automatic-speech-recognition", model="openai/whisper-tiny")

# --- GOOGLE SHEETS FUNCTION  ---
def save_to_sheet(original, translation, correction, direction):
    try:
        json_creds = os.getenv("GOOGLE_CREDENTIALS")
        if not json_creds: return "⚠️ Error: Credentials missing."
        creds_dict = json.loads(json_creds)
        scope = ["https://spreadsheets.google.com/feeds", "https://www.googleapis.com/auth/drive"]
        creds = ServiceAccountCredentials.from_json_keyfile_dict(creds_dict, scope)
        client = gspread.authorize(creds)
        sheet = client.open(SHEET_NAME).sheet1
        if not sheet.get_all_values():
            sheet.append_row(["Timestamp", "Direction", "Original Text", "AI Translation", "User Correction"])
        sheet.append_row([str(datetime.now()), direction, original, translation, correction])
        return "✅ Saved to Google Sheets."
    except Exception as e:
        return f"❌ Error: {str(e)}"

# --- CORE FUNCTIONS ---

def speech_to_text(audio_path):
    """Converts Microphone Audio -> Hindi/English Text"""
    if audio_path is None:
        return ""
    print("Transcribing audio...")
    text = asr_pipeline(audio_path)["text"]
    return text

def text_to_speech(text, language="hi"):
    """Converts Text -> Audio File"""
    if not text:
        return None
    try:
        # Save audio to a temporary file
        tts = gTTS(text=text, lang=language)
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
        tts.save(temp_file.name)
        return temp_file.name
    except:
        return None

def process_translation(text, audio_input, direction):
    # 1. Handle Input Source (Audio vs Text)
    if audio_input is not None:
        # If user spoke, transcribe it first
        original_text = speech_to_text(audio_input)
    else:
        original_text = text

    if not original_text:
        return "", "", None

    # 2. Translate
    target_pipeline = pipe_k2h if direction == "Kurukh -> Hindi" else pipe_h2k
    try:
        results = target_pipeline(original_text, max_length=128)
        translated_text = results[0]['generated_text']
    except Exception as e:
        return str(e), "", None

    # 3. Handle Audio Output (Read aloud if Hindi)
    audio_output = None
    
    if direction == "Kurukh -> Hindi":
        audio_output = text_to_speech(translated_text, "hi")
    
    # Return: (Input Text shown to user), (Translation), (Audio File)
    return original_text, translated_text, audio_output

# --- THE UI ---
with gr.Blocks() as demo:
    gr.Markdown("# 🇮🇳 AI Kurukh (Kurux) Translator")
    gr.Markdown("### Voice-Enabled Tribal Language Preserver")
    
    with gr.Tabs():
        with gr.TabItem("🗣️ Voice & Text Translator"):
            
            with gr.Row():
                direction = gr.Radio(["Kurukh -> Hindi", "Hindi -> Kurukh"], label="Mode", value="Kurukh -> Hindi")
            
            # INPUT SECTION
            with gr.Row():
                with gr.Column():
                    input_text = gr.Textbox(label="Type Here", placeholder="Or use the microphone below...", lines=3)
                    input_audio = gr.Audio(sources=["microphone"], type="filepath", label="Or Speak (Hindi Only)")
                    translate_btn = gr.Button("Translate 🚀", variant="primary")
                
                # OUTPUT SECTION
                with gr.Column():
                    output_text = gr.Textbox(label="Translation", lines=3, interactive=False)
                    output_audio = gr.Audio(label="Listen to Translation (Hindi)", interactive=False)
            
            # LOGIC
            translate_btn.click(
                fn=process_translation, 
                inputs=[input_text, input_audio, direction], 
                outputs=[input_text, output_text, output_audio] # Updates input box with transcribed text too!
            )

        with gr.TabItem("📝 Improve the AI"):
            gr.Markdown("### Help us improve!")
            with gr.Row():
                fb_direction = gr.Radio(["Kurukh -> Hindi", "Hindi -> Kurukh"], label="Direction", value="Kurukh -> Hindi")
            with gr.Row():
                fb_original = gr.Textbox(label="Original")
                fb_ai_output = gr.Textbox(label="AI Said")
                fb_user_correct = gr.Textbox(label="Correction", lines=2)
            submit_btn = gr.Button("Submit to Google Sheets ✅")
            status_lbl = gr.Label(label="Status")
            submit_btn.click(fn=save_to_sheet, inputs=[fb_original, fb_ai_output, fb_user_correct, fb_direction], outputs=status_lbl)

demo.launch()