import os
import subprocess
import tempfile
import gradio as gr
import torch
import shutil
from transformers import (
    AutoTokenizer, 
    AutoModelForCausalLM,
    WhisperForConditionalGeneration,
    WhisperProcessor
)
import librosa

# --- မော်ဒယ်များကို စတင်တင်ဆောင်ခြင်း (Model Loading) ---
print("🔄 မော်ဒယ်များ တင်ဆောင်နေသည်...")

device = "cuda" if torch.cuda.is_available() else "cpu"

# ==============================================
# ၁။ ASR မော်ဒယ် (Hugging Face Transformers ဖြင့်)
# ==============================================
asr_model_id = "Chonlasitk/whisper-burmese"
print(f"⏳ ASR မော်ဒယ် ({asr_model_id}) ကို တင်ဆောင်နေသည်...")

# Whisper Processor က audio ကို input features အဖြစ် ပြောင်းပေးပြီး 
# WhisperForConditionalGeneration က text ထုတ်ပေးပါတယ်။
processor = WhisperProcessor.from_pretrained(asr_model_id)
asr_model = WhisperForConditionalGeneration.from_pretrained(asr_model_id).to(device)
print(f"✅ ASR မော်ဒယ်ကို {device} ပေါ်တွင် တင်ဆောင်ပြီးပါပြီ။")

# ==============================================
# ၂။ ဘာသာပြန် LLM မော်ဒယ် (MIG Burmese LLM)
# ==============================================
translation_model_id = "Ko-Yin-Maung/mig-burmese-llm"
print(f"⏳ ဘာသာပြန် LLM မော်ဒယ် ({translation_model_id}) ကို တင်ဆောင်နေသည်...")

try:
    translation_tokenizer = AutoTokenizer.from_pretrained(translation_model_id)
    if device == "cuda":
        translation_model = AutoModelForCausalLM.from_pretrained(
            translation_model_id,
            torch_dtype=torch.float16,
            device_map="auto"
        )
    else:
        translation_model = AutoModelForCausalLM.from_pretrained(
            translation_model_id,
            device_map="auto"
        )
    print(f"✅ ဘာသာပြန် LLM မော်ဒယ်ကို တင်ဆောင်ပြီးပါပြီ။")
except Exception as e:
    print(f"❌ ဘာသာပြန်မော်ဒယ်ကို တင်ဆောင်ရာတွင် အမှားတွေ့ရှိသည်: {e}")
    exit()

print("🚀 အားလုံး အဆင်သင့်ဖြစ်ပါပြီ။ Gradio App ကို စတင်လိုက်ပါ။")

# --- အကူအညီ ဖန်ရှင်များ (Helper Functions) ---

def extract_audio(video_path):
    """ဗီဒီယိုဖိုင်မှ အသံကို 16kHz Mono WAV format ဖြင့် ထုတ်ယူပါ။"""
    print("🔊 ဗီဒီယိုမှ အသံထုတ်ယူနေသည်...")
    audio_path = "extracted_audio.wav"
    command = [
        "ffmpeg",
        "-i", video_path,
        "-ac", "1",
        "-ar", "16000",
        "-vn",
        "-f", "wav",
        "-y",
        audio_path
    ]
    subprocess.run(command, check=True, capture_output=True)
    print(f"✅ အသံထုတ်ယူပြီးပါပြီ: {audio_path}")
    return audio_path

def transcribe_audio(audio_path):
    """
    Hugging Face Whisper မော်ဒယ်ကို သုံး၍ အသံကို စာသားပြောင်းပါ။
    အချိန်ကုဒ် segment များပါ ထုတ်ပေးပါသည်။
    """
    print("🎤 အသံကို စာသားပြောင်းနေသည်...")
    
    # Audio ဖိုင်ကို 16kHz sampling rate နဲ့ ဖတ်ပါ
    audio_array, sampling_rate = librosa.load(audio_path, sr=16000)
    
    # Whisper Processor က input features ပြင်ဆင်ပေးပါတယ်
    inputs = processor(audio_array, sampling_rate=16000, return_tensors="pt").input_features
    inputs = inputs.to(device)
    
    # Model ကို generate လုပ်ပြီး token ids ရယူပါ
    with torch.no_grad():
        generated_ids = asr_model.generate(
            inputs,
            task="transcribe",
            return_timestamps=True,   # အချိန်ကုဒ်ပါ ထုတ်ပေးရန် အရေးကြီးပါသည်
            max_new_tokens=448,
            language="burmese",       # မြန်မာဘာသာစကားကို တိုက်ရိုက်သတ်မှတ်
        )
    
    # Token ids မှ text နှင့် timestamps ပါတဲ့ result ကို decode လုပ်ပါ
    transcription = processor.decode(generated_ids[0], skip_special_tokens=True, output_offsets=True)
    
    # Whisper ရဲ့ decode က dictionary ပုံစံ {text: "...", offsets: [...]} ပြန်ပေးပါတယ်
    # offsets ထဲမှာ တစ်ခုချင်းစီရဲ့ timestamp ပါပါတယ်။
    segments = []
    if "offsets" in transcription:
        for offset in transcription["offsets"]:
            if "text" in offset and "timestamp" in offset:
                start, end = offset["timestamp"]
                segments.append({
                    "start": start,
                    "end": end,
                    "text": offset["text"].strip()
                })
    
    # အကယ်၍ segments မရခဲ့ရင် fallback အနေနဲ့ text တစ်ခုလုံးကို ယူပါမယ်
    if not segments:
        full_text = transcription.get("text", "").strip()
        if full_text:
            segments.append({"start": 0.0, "end": len(audio_array)/16000, "text": full_text})
    
    print(f"✅ စာသားပြောင်းပြီးပါပြီ။ Segment {len(segments)} ခုရရှိသည်။")
    return segments

def translate_to_burmese(text):
    """အင်္ဂလိပ်စာသားကို MIG Burmese LLM သုံး၍ မြန်မာဘာသာသို့ ပြန်ဆိုပါ။"""
    if not text or text.isspace():
        return ""
    
    # MIG Burmese LLM အတွက် prompt format
    prompt = f"<start_of_turn>user\nTranslate to Myanmar: {text}\n<end_of_turn><start_of_turn>model\n"
    inputs = translation_tokenizer(prompt, return_tensors="pt").to(translation_model.device)
    
    with torch.no_grad():
        outputs = translation_model.generate(
            **inputs,
            max_new_tokens=512,
            do_sample=False,
            temperature=0.7,
        )
    translated = translation_tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # "model\n" နောက်မှ ရလဒ်ကို ထုတ်ယူပါ
    try:
        result = translated.split("model\n")[-1].strip()
        return result
    except:
        return translated

def format_timestamp(seconds):
    """စက္ကန့်ကို SRT အချိန်ကုဒ် format (HH:MM:SS,ms) သို့ပြောင်းပါ။"""
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    secs = int(seconds % 60)
    millis = int((seconds - int(seconds)) * 1000)
    return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"

def process_video(video_file, progress=gr.Progress()):
    """
    ဗီဒီယိုကို အဆင့်လိုက် လုပ်ဆောင်ပေးမယ့် အဓိက ဖန်ရှင်။
    """
    if video_file is None:
        raise gr.Error("ကျေးဇူးပြု၍ ဗီဒီယိုဖိုင်တစ်ခု အပ်လုဒ်လုပ်ပါ။")
    
    progress(0, desc="စတင်နေပါပြီ...")
    audio_path = None
    srt_file_path = None
    
    try:
        # --- အဆင့် ၁: အသံထုတ်ယူခြင်း ---
        progress(0.1, desc="ဗီဒီယိုမှ အသံထုတ်ယူနေသည်...")
        audio_path = extract_audio(video_file)

        # --- အဆင့် ၂: ASR ပြုလုပ်ခြင်း (Whisper) ---
        progress(0.3, desc="အသံကို စာသားပြောင်းနေသည်... (ASR)")
        segments = transcribe_audio(audio_path)
        
        if not segments:
            raise gr.Error("အသံဖိုင်ထဲတွင် စာသားပြောင်းရန် အကြောင်းအရာ မရှိပါ။")
        
        # --- အဆင့် ၃: ဘာသာပြန်ခြင်း ---
        progress(0.6, desc="မြန်မာဘာသာသို့ ပြန်ဆိုနေသည်...")
        total_segments = len(segments)
        burmese_segments = []
        
        for i, segment in enumerate(segments):
            progress(0.6 + (0.3 * (i / total_segments)), 
                     desc=f"မြန်မာပြန်ဆိုနေသည်... ({i+1}/{total_segments})")
            
            english_text = segment["text"].strip()
            if english_text:
                burmese_text = translate_to_burmese(english_text)
                burmese_segments.append(burmese_text)
            else:
                burmese_segments.append("")
        
        # --- အဆင့် ၄: SRT ဖိုင်ဖန်တီးခြင်း ---
        progress(0.95, desc="SRT ဖိုင်ကို ဖန်တီးနေသည်...")
        srt_filename = "subtitles.srt"
        with open(srt_filename, "w", encoding="utf-8") as f:
            for i, segment in enumerate(segments):
                if i < len(burmese_segments) and burmese_segments[i]:
                    start_time = format_timestamp(segment["start"])
                    end_time = format_timestamp(segment["end"])
                    f.write(f"{i+1}\n")
                    f.write(f"{start_time} --> {end_time}\n")
                    f.write(f"{burmese_segments[i]}\n\n")
        
        srt_file_path = srt_filename
        progress(1.0, desc="ပြီးစီးပါပြီ။")
        
        return "✅ လုပ်ငန်းစဉ်အောင်မြင်စွာ ပြီးစီးပါပြီ။ SRT ဖိုင်ကို အောက်ပါလင့်ခ်မှ ဒေါင်းလုဒ်ရယူနိုင်ပါသည်။", srt_file_path

    except subprocess.CalledProcessError as e:
        raise gr.Error(f"FFmpeg အမှား: {e.stderr.decode('utf-8')}")
    except Exception as e:
        raise gr.Error(f"မမျှော်လင့်ထားသော အမှားတစ်ခု ဖြစ်ပွားခဲ့သည်: {e}")
    finally:
        if audio_path and os.path.exists(audio_path):
            os.remove(audio_path)
            print(f"🧹 ယာယီအသံဖိုင် {audio_path} ကို ရှင်းလင်းပြီးပါပြီ။")

# --- Gradio Interface ဖန်တီးခြင်း ---
with gr.Blocks(theme=gr.themes.Soft(), title="မြန်မာစာတန်းထိုး ဖန်တီးရေး") as demo:
    gr.Markdown("""
    # 🇲🇲 မြန်မာစာတန်းထိုး အလိုအလျောက် ဖန်တီးရေး (Auto Burmese Subtitle Generator)
    ဗီဒီယိုဖိုင်တစ်ခု အပ်လုဒ်လုပ်လိုက်ရုံနဲ့ အသံကို နားထောင်ပြီး မြန်မာလို အချိန်နဲ့တစ်ပြေးညီ စာတန်းထိုး (SRT) ဖိုင်ကို အလိုအလျောက် ထုတ်ပေးပါမယ်။
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            video_input = gr.Video(label="ဗီဒီယိုဖိုင် ထည့်ရန်", sources=["upload"])
            submit_btn = gr.Button("စာတန်းထိုးဖန်တီးမည်", variant="primary")
        
        with gr.Column(scale=1):
            status_output = gr.Textbox(label="အခြေအနေ", interactive=False)
            srt_output = gr.File(label="မြန်မာစာတန်းထိုး (SRT) ဖိုင်", interactive=False)
    
    submit_btn.click(
        fn=process_video,
        inputs=[video_input],
        outputs=[status_output, srt_output]
    )
    
    gr.Markdown("""
    ---
    ### ℹ️ အသုံးပြုနည်း
    ၁။ `ဗီဒီယိုဖိုင် ထည့်ရန်` အကွက်ကို နှိပ်၍ သင့်စက်ထဲမှ ဗီဒီယိုတစ်ခုကို ရွေးချယ်ပါ။
    ၂။ `စာတန်းထိုးဖန်တီးမည်` ခလုတ်ကို နှိပ်ပါ။
    ၃။ လုပ်ငန်းစဉ် ပြီးစီးသည်အထိ စောင့်ပါ။
    ၄။ ပြီးစီးသွားပါက `မြန်မာစာတန်းထိုး (SRT) ဖိုင်` အကွက်မှ ဒေါင်းလုဒ်ရယူနိုင်ပါပြီ။
    """)

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