import gradio as gr
import numpy as np
import torch
import soundfile as sf
import librosa
import noisereduce as nr
from scipy import signal
from transformers import AutoFeatureExtractor, AutoModelForAudioFrameClassification
from recitations_segmenter import segment_recitations, clean_speech_intervals
import tempfile
import os
import zipfile

# --- الإعدادات الافتراضية الثابتة ---
MODEL_ID = "obadx/recitation-segmenter-v2"
device = 'cuda' if torch.cuda.is_available() else 'cpu'
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32

# تحميل النموذج لمرة واحدة عند التشغيل
print(f"Loading model on {device}...")
processor = AutoFeatureExtractor.from_pretrained(MODEL_ID)
model = AutoModelForAudioFrameClassification.from_pretrained(
    MODEL_ID, torch_dtype=dtype, device_map=device
)

# --- وظائف تحسين الصوت (من الملف الثاني) ---
def enhance_audio_logic(audio, sr):
    # 1. إزالة الضوضاء (قوة 0.75 افتراضية)
    noise_clip = audio[:int(sr * 0.5)] if len(audio) > sr else audio
    audio = nr.reduce_noise(y=audio, sr=sr, y_noise=noise_clip, prop_decrease=0.75)
    
    # 2. فلتر الترددات (Bandpass)
    nyquist = sr / 2
    low, high = 80 / nyquist, min(8000 / nyquist, 0.99)
    b, a = signal.butter(6, [low, high], btype='band')
    audio = signal.filtfilt(b, a, audio)
    
    # 3. تحسين الوضوح (Clarity)
    harmonic, _ = librosa.effects.hpss(audio, margin=3.0)
    audio = (0.8 * harmonic + 0.2 * audio).astype(np.float32)
    
    # 4. تطبيع الصوت (-18 dBFS)
    target_rms = 10 ** (-18.0 / 20)
    rms = np.sqrt(np.mean(audio ** 2))
    if rms > 1e-8:
        audio = np.clip(audio * (target_rms / rms), -1.0, 1.0)
        
    return audio

# --- الوظيفة الرئيسية المدمجة ---
def process_pipeline(audio_path):
    if audio_path is None:
        return None, None

    # 1. تحميل الملف
    y, sr = librosa.load(audio_path, sr=16000) # التقطيع يحتاج 16000Hz
    
    # 2. التحسين التلقائي
    enhanced_y = enhance_audio_logic(y, sr)
    
    # حفظ الملف المحسن في مسار مؤقت للعرض
    temp_dir = tempfile.mkdtemp()
    enhanced_file_path = os.path.join(temp_dir, "enhanced_full.wav")
    sf.write(enhanced_file_path, enhanced_y, sr)
    
    # 3. التقطيع التلقائي (باستخدام الصوت المحسن)
    wav_tensor = torch.tensor(enhanced_y).float()
    sampled_outputs = segment_recitations(
        [wav_tensor], model, processor, device=device, dtype=dtype, batch_size=4
    )
    
    clean_out = clean_speech_intervals(
        sampled_outputs[0].speech_intervals,
        sampled_outputs[0].is_complete,
        min_silence_duration_ms=30, # إعدادات افتراضية مثالية
        min_speech_duration_ms=30,
        pad_duration_ms=30,
        return_seconds=True,
    )
    
    intervals = clean_out.clean_speech_intervals
    
    # 4. حفظ المقاطع في ملف ZIP
    zip_path = os.path.join(temp_dir, "segments.zip")
    with zipfile.ZipFile(zip_path, 'w') as zipf:
        for idx, (start_sec, end_sec) in enumerate(intervals):
            start_sample = int(start_sec * sr)
            end_sample = int(end_sec * sr)
            segment = enhanced_y[start_sample:end_sample]
            
            seg_name = f"segment_{idx+1:03d}.wav"
            seg_path = os.path.join(temp_dir, seg_name)
            sf.write(seg_path, segment, sr)
            zipf.write(seg_path, seg_name)
            
    return enhanced_file_path, zip_path

# --- واجهة Gradio البسيطة ---
with gr.Blocks(title="مُعالج التلاوات الشامل") as demo:
    gr.Markdown("# 🕌 مُعالج التلاوات الشامل (تحسين + تقطيع)")
    gr.Markdown("ارفع ملف الصوت وسيتم تحسينه وتقطيعه تلقائياً بأفضل الإعدادات.")
    
    with gr.Row():
        input_audio = gr.Audio(label="📤 ارفع ملف التلاوة", type="filepath")
    
    btn = gr.Button("🚀 ابدأ المعالجة الذكية", variant="primary")
    
    with gr.Row():
        out_enhanced = gr.Audio(label="🎵 الصوت المحسن كاملاً", type="filepath")
        out_zip = gr.File(label="📦 تحميل كل المقاطع المقطعة (ZIP)")

    btn.click(
        fn=process_pipeline,
        inputs=[input_audio],
        outputs=[out_enhanced, out_zip]
    )

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