import torch
import spaces
import gradio as gr
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline

# आपका पसंदीदा और सटीक मॉडल आईडी (या आप यहाँ "openai/whisper-large-v3" भी रख सकते हैं)
MODEL_NAME = "collabora/whisper-large-v2-hindi"
BATCH_SIZE = 8

print("मॉडल और प्रोसेसर को आरंभिक रैम (RAM) में लोड किया जा रहा है...")

# १. मॉडल को शुरुआत में CPU पर लोड करें (ZeroGPU को ब्लॉक होने से बचाने के लिए अनिवार्य)
model = AutoModelForSpeechSeq2Seq.from_pretrained(
    MODEL_NAME, 
    torch_dtype=torch.float16, 
    low_cpu_mem_usage=True, 
    use_safetensors=True
)
processor = AutoProcessor.from_pretrained(MODEL_NAME)

print("मॉडल सफलतापूर्वक लोड हो गया। अब ऐप यूजर के अनुरोध के लिए तैयार है।")

# २. मुख्य ट्रांसक्रिप्शन प्रक्रिया - जिसे ZeroGPU नियंत्रित करेगा
@spaces.GPU(duration=120)
def transcribe(inputs, task):
    if inputs is None:
        raise gr.Error("No audio file submitted! Please upload or record an audio file.")
    
    try:
        # मॉडल को केवल गणना के क्षणों में CUDA पर स्थानांतरित करें
        model.to("cuda")
        
        # फ़ंक्शन के भीतर ही तात्कालिक पाइपलाइन का निर्माण
        pipe = pipeline(
            task="automatic-speech-recognition",
            model=model,
            tokenizer=processor.tokenizer,
            feature_extractor=processor.feature_extractor,
            chunk_length_s=30,
            device="cuda",
            torch_dtype=torch.float16,
        )
        
        # अनुवाद जनरेट करें
        result = pipe(inputs, batch_size=BATCH_SIZE, generate_kwargs={"task": task, "language": "hindi"})
        text_output = result["text"].strip()
        
        # काम पूरा होते ही मॉडल को वापस CPU पर भेजें
        model.to("cpu")
        
        # डाउनलोड के लिए टेक्स्ट फ़ाइल का निर्माण
        file_path = "transcription.txt"
        with open(file_path, "w", encoding="utf-8") as f:
            f.write(text_output)
            
        return text_output, file_path

    except Exception as e:
        # किसी भी व्यवधान की स्थिति में मॉडल को सुरक्षित स्थिति में लाएँ
        try:
            model.to("cpu")
        except:
            pass
        raise gr.Error(f"प्रक्रिया में व्यवधान: {str(e)}")

# ३. आधुनिक Gradio Blocks इंटरफ़ेस (डाउनलोड लिंक के साथ)
custom_css = """
footer {visibility: hidden}
.gradio-container {background-color: #fcfcfc}
#header {text-align: center; margin-bottom: 20px}
"""

with gr.Blocks(title="IndicWhisper Optimized") as demo:
    gr.HTML("<div id='header'><h1>🎙️ Hindi Whisper (Optimized ZeroGPU Pipeline)</h1></div>")
    
    with gr.Row():
        with gr.Column():
            audio_input = gr.Audio(sources=["microphone", "upload"], type="filepath", label="ऑडियो इनपुट")
            task_radio = gr.Radio(["transcribe", "translate"], label="Task", value="transcribe")
            submit_btn = gr.Button("अनुवाद करें (Transcribe)", variant="primary")
            
        with gr.Column():
            output_text = gr.Textbox(label="Transcription Output", lines=10, placeholder="टेक्स्ट यहाँ दिखेगा...")
            download_file = gr.File(label="टैक्स्ट फ़ाइल डाउनलोड करें")
            
    gr.Markdown("""
    ---
    **सुझाव:** यहाँ से प्राप्त आउटपुट को कॉपी करें या फ़ाइल डाउनलोड करके अपने **Gemini Gem** में पेस्ट करें ताकि **पञ्चमाक्षर नियमों** (ङ्, ञ्, ण्, न्, म्) के अनुसार शुद्धिकरण किया जा सके।
    """)

    submit_btn.click(
        fn=transcribe, 
        inputs=[audio_input, task_radio], 
        outputs=[output_text, download_file]
    )

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