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

# انتخاب نسخه مرجع جهانی چندزبانه OpenAI Whisper V3
model_id = "openai/whisper-large-v3"  

# ۱. انتخاب بهینه‌ترین فرمت عددی برای سخت‌افزار A100/H200
if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
    torch_dtype = torch.bfloat16
    print("Using optimized BFloat16 for A100/H200")
else:
    torch_dtype = torch.float16
    print("Using Float16")

# ۲. فعال‌سازی مکانیسم شتاب‌دهنده توجه (Attention)
attn_implementation = "sdpa"
try:
    import flash_attn
    attn_implementation = "flash_attention_2"
    print("Using Flash Attention 2 for maximum hardware utilization")
except ImportError:
    print("Flash Attention 2 not found. Falling back to high-performance PyTorch SDPA")

# ۳. بارگذاری بهینه مدل چندزبانه با تکیه بر Safetensors و Memory optimization
model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_id,
    torch_dtype=torch_dtype,
    low_cpu_mem_usage=True,
    use_safetensors=True,
    attn_implementation=attn_implementation
)

# انتقال مدل به حافظه گرافیکی در صورت در دسترس بودن
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

processor = AutoProcessor.from_pretrained(model_id)

# تعریف متن راهنمای پرامپت فارسی به صورت سراسری
prompt_text = "کلمات دقیقاً همان‌طور که تلفظ می‌شوند، با رعایت حروف اضافه و ساختار عامیانه، محاوره‌ای و شکسته نوشته شوند."

# ۴. ساخت پایپ‌لاین تشخیص گفتار چندزبانه با استفاده از پارامتر اصلاح شده‌ی dtype
pipe = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor,
    dtype=torch_dtype,
    device=device,
)

# نقشه کدهای زبانی برای منوی کشویی انتخاب زبان
LANGUAGE_MAPPING = {
    "Auto-Detect (تشخیص خودکار)": None,
    "Persian / فارسی": "fa",
    "English / انگلیسی": "en",
    "Spanish / اسپانیایی": "es",
    "French / فرانسوی": "fr",
    "German / آلمانی": "de",
    "Arabic / عربی": "ar",
    "Turkish / ترکی": "tr",
    "Italian / ایتالیایی": "it",
    "Russian / روسی": "ru",
    "Chinese / چینی": "zh",
    "Japanese / ژاپنی": "ja",
    "Korean / کره‌ای": "ko"
}

# تابع کمکی برای فرمت‌بندی زمان در ساختار استاندارد زیرنویس (HH:MM:SS,mmm)
def format_timestamp(seconds):
    hours = int(seconds // 3600)
    minutes = int((seconds % 3600) // 60)
    secs = int(seconds % 60)
    millis = int(round((seconds % 1) * 1000))
    return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"

# مبدل داده‌های کلمه به کلمه ویسپر به ساختار استاندارد خوانای SRT
def create_srt_from_chunks(chunks, max_words_per_segment=6, max_duration_per_segment=3.0):
    srt_lines = []
    segment_idx = 1
    current_segment_words = []
    current_start = None
    
    for chunk in chunks:
        text = chunk.get("text", "").strip()
        timestamp = chunk.get("timestamp")
        
        if not text or timestamp is None:
            continue
            
        start, end = timestamp
        if start is None:
            continue
        if end is None:
            end = start + 0.3  # مقدار جایگزین در صورت عدم ثبت زمان پایان توکن
            
        if current_start is None:
            current_start = start
            
        current_segment_words.append(text)
        duration = end - current_start
        
        # دسته‌بندی کلمات در قالب خطوط زیرنویس خوانا (مثلا حداکثر ۶ کلمه یا ۳ ثانیه برای هر خط)
        if len(current_segment_words) >= max_words_per_segment or duration >= max_duration_per_segment:
            line_text = " ".join(current_segment_words)
            srt_lines.append(f"{segment_idx}")
            srt_lines.append(f"{format_timestamp(current_start)} --> {format_timestamp(end)}")
            srt_lines.append(line_text)
            srt_lines.append("")
            segment_idx += 1
            current_segment_words = []
            current_start = None
            
    if current_segment_words:
        last_end = chunks[-1].get("timestamp")[1] if chunks and chunks[-1].get("timestamp") else current_start + 1.0
        line_text = " ".join(current_segment_words)
        srt_lines.append(f"{segment_idx}")
        srt_lines.append(f"{format_timestamp(current_start)} --> {format_timestamp(last_end)}")
        srt_lines.append(line_text)
        srt_lines.append("")
        
    return "\n".join(srt_lines)

# تابع تولید دیتای خام زمان‌بندی کلمات (JSON) برای دسترسی توسعه‌دهندگان
def get_raw_word_timestamps(chunks):
    word_timestamps = []
    for chunk in chunks:
        text = chunk.get("text", "").strip()
        timestamp = chunk.get("timestamp")
        if text and timestamp is not None:
            word_timestamps.append({
                "word": text,
                "start": round(timestamp[0], 3),
                "end": round(timestamp[1], 3)
            })
    return json.dumps(word_timestamps, ensure_ascii=False, indent=2)

@spaces.GPU
def transcribe_audio(audio_path, language_selection, optimize_persian):
    if audio_path is None:
        return "Please upload or record an audio file first. / لطفاً ابتدا فایل صوتی را آپلود یا ضبط کنید.", "", ""
    
    try:
        # ۱. بارگذاری ایمن فایل صوتی به فرکانس ۱۶۰۰۰ هرتز (سازگار با همه فرمت‌ها مانند میکروفون وب‌ام یا m4a)
        audio_data, sampling_rate = librosa.load(audio_path, sr=16000)
        
        # ۲. پیکربندی پویا بر اساس زبان انتخاب شده توسط کاربر
        lang_code = LANGUAGE_MAPPING.get(language_selection, None)
        
        # تنظیمات پایه تولید متن با استفاده از روش پیشرفته Beam Search برای افزایش حداکثری دقت تفکیک کلمات
        generate_kwargs = {
            "task": "transcribe",
            "do_sample": False,
            "num_beams": 5  # ارزیابی موازی ۵ مسیر واژگانی برای کاهش خطاهای شنیداری صوتی
        }
        
        # اگر کاربر زبانی به جز تشخیص خودکار انتخاب کرده باشد
        if lang_code is not None:
            generate_kwargs["language"] = lang_code
            
        # ۳. اعمال پرامپت اختصاصی فارسی فقط در صورت تمایل کاربر (به عنوان تانسور ۱ بعدی پایتورچِ روی GPU)
        if optimize_persian:
            prompt_ids = processor.get_prompt_ids(prompt_text, return_tensors="pt")
            prompt_ids = prompt_ids.to(device)
            generate_kwargs["prompt_ids"] = prompt_ids
            generate_kwargs["language"] = "fa"  # اجبار به فارسی در صورت فعال بودن این گزینه
        
        # ۴. پردازش نهایی با حداکثر ظرفیت سخت‌افزاری روی A100/H200 و درخواست ثبت تایم‌استمپ کلمات
        result = pipe(
            {"raw": audio_data, "sampling_rate": sampling_rate},
            chunk_length_s=30,
            stride_length_s=(6, 6),
            batch_size=24,
            return_timestamps="word",  # فعال‌سازی ثبت دقیق زمان‌بندی برای هر کلمه
            generate_kwargs=generate_kwargs
        )
        
        full_text = result["text"].strip()
        chunks = result.get("chunks", [])
        
        # ۵. تولید ساختارهای سه‌گانه دیتای خروجی
        srt_content = create_srt_from_chunks(chunks)
        word_timestamps_json = get_raw_word_timestamps(chunks)
        
        # بازگرداندن هر سه خروجی در قالب تاپل
        return full_text, srt_content, word_timestamps_json
        
    except Exception as e:
        return f"An error occurred during processing / خطایی در حین پردازش رخ داد: {str(e)}", "", ""

# رابط کاربری چندزبانه و مدرن بر پایه Gradio 6.0
with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 🌐 Global Speech-to-Text & Subtitle Studio (Whisper Large V3)
        ### سیستم هوشمند و بین‌المللی تبدیل گفتار به متن و زیرنویس خودکار (بهینه‌سازی شده روی A100/H200)
        
        *Supports over 99 languages with automatic detection and Word-Level Timestamps.*  
        *پشتیبانی از بیش از ۹۹ زبان زنده دنیا به همراه تشخیص خودکار زبان و زمان‌بندی دقیق کلمه به کلمه.*
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            audio_input = gr.Audio(
                sources=["upload", "microphone"], 
                type="filepath", 
                label="Audio Input / ورودی صدا"
            )
            
            language_dd = gr.Dropdown(
                choices=list(LANGUAGE_MAPPING.keys()),
                value="Auto-Detect (تشخیص خودکار)",
                label="Select Language / انتخاب زبان صوتی"
            )
            
            optimize_fa_chk = gr.Checkbox(
                label="Optimize for Spoken Persian / بهینه‌سازی برای محاوره عامیانه فارسی (فقط برای ویس‌های فارسی)",
                value=False
            )
            
            submit_btn = gr.Button("Transcribe / شروع فرآیند تبدیل", variant="primary")
            
        with gr.Column(scale=1):
            with gr.Tabs():
                with gr.Tab("Full Text / متن کامل"):
                    text_output = gr.Textbox(
                        label="Transcription Output / متن پیوسته", 
                        lines=14, 
                        buttons=["copy"],
                        interactive=False
                    )
                with gr.Tab("SRT Subtitles / زیرنویس استاندارد"):
                    srt_output = gr.Textbox(
                        label="Subtitles (SRT Format) / فایل زیرنویس", 
                        lines=14, 
                        buttons=["copy"],
                        interactive=False
                    )
                with gr.Tab("Word Timestamps / زمان‌بندی کلمات (JSON)"):
                    json_output = gr.Textbox(
                        label="Word-Level Timestamps (JSON) / زمان خام کلمات", 
                        lines=14, 
                        buttons=["copy"],
                        interactive=False
                    )
            
    submit_btn.click(
        fn=transcribe_audio,
        inputs=[audio_input, language_dd, optimize_fa_chk],
        outputs=[text_output, srt_output, json_output],
        api_name="transcribe"  # نام اختصاصی متد برای فراخوانی API کلاینت
    )

if __name__ == "__main__":
    # فعال‌سازی تم آبی‌رنگ در زمان راه‌اندازی بر اساس استانداردهای Gradio 6.0
    demo.launch(theme=gr.themes.Default(primary_hue="blue"))