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
import requests  

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

# --- LOAD 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)

print("Loading Voice Model...")
asr_pipeline = pipeline("automatic-speech-recognition", model="openai/whisper-tiny")

# --- HELPER FUNCTIONS ---

def transliterate_to_hindi(text):
    
    try:
        url = "https://inputtools.google.com/request?text={}&itc=hi-t-i0-und&num=1"
        response = requests.get(url.format(text))
        result = response.json()
        # The API returns a nested list; we grab the first suggestion
        return result[1][0][1][0]
    except:
        return text  # If it fails (no internet), return original text

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)}"

def speech_to_text(audio_path):
    if audio_path is None: return ""
    return asr_pipeline(audio_path)["text"]

def text_to_speech(text, language="hi"):
    if not text: return None
    try:
        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

# --- MAIN TRANSLATION LOGIC ---
def process_translation(text, audio_input, direction, is_hinglish):
    
    # 1. Get Text from Voice or Typing
    original_text = speech_to_text(audio_input) if audio_input else text
    if not original_text: return "", "", None

    # 2. Handle Hinglish (NEW FEATURE)
    # Only applies if translating TO Kurukh (User typing Hindi in English letters)
    if direction == "Hindi -> Kurukh" and is_hinglish:
        original_text = transliterate_to_hindi(original_text)

    # 3. 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

    # 4. Audio Output (For Hindi)
    audio_output = None
    if direction == "Kurukh -> Hindi":
        audio_output = text_to_speech(translated_text, "hi")
    
    # Return: (Updated Input Box), (Translation), (Audio)
    return original_text, translated_text, audio_output

# --- THE UI ---
with gr.Blocks() as demo:
    gr.Markdown("# 🇮🇳 AI Kurukh (Kurux) Translator")
    gr.Markdown("### Voice & Hinglish Supported")
    
    with gr.Tabs():
        with gr.TabItem("🗣️ Translator"):
            
            with gr.Row():
                direction = gr.Radio(["Kurukh -> Hindi", "Hindi -> Kurukh"], label="Mode", value="Kurukh -> Hindi")
                # NEW CHECKBOX
                is_hinglish = gr.Checkbox(label="I am typing Hindi in English (e.g. 'Tumhara')", value=False)
            
            with gr.Row():
                with gr.Column():
                    input_text = gr.Textbox(label="Input", placeholder="Type here...", lines=3)
                    input_audio = gr.Audio(sources=["microphone"], type="filepath", label="Voice Input (Hindi)")
                    translate_btn = gr.Button("Translate 🚀", variant="primary")
                
                with gr.Column():
                    output_text = gr.Textbox(label="Translation", lines=3, interactive=False)
                    output_audio = gr.Audio(label="Listen (Hindi Only)", interactive=False)
            
            translate_btn.click(
                fn=process_translation, 
                inputs=[input_text, input_audio, direction, is_hinglish], 
                outputs=[input_text, output_text, output_audio]
            )

        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()