import gradio as gr
import tensorflow as tf
import speechbrain as sb
from speechbrain.inference.classifiers import EncoderClassifier
import torch
import numpy as np
import librosa
import os
import torch.nn.functional as F
from speechbrain.dataio.dataio import read_audio
import math
import tempfile
from huggingface_hub import snapshot_download

print("=== MEMULAI APLIKASI WEB ===")

# -----------------------------------------------------------
# 1. KONFIGURASI & PARAMETER
# -----------------------------------------------------------
print("Mengatur konfigurasi...")

# Path ini akan bekerja di dalam environment Hugging Face Spaces
# Kita akan buat folder 'voiceprint_db' di dalam Space kita
VOICEPRINT_DB_PATH = 'voiceprint_db'
os.makedirs(VOICEPRINT_DB_PATH, exist_ok=True)

# Path ke Model Pilar 1 (unggah file ini ke Space Anda)
PILAR_1_MODEL_PATH = 'model_pilar_1_indonesia.keras' 

# Parameter Audio (HARUS SAMA DENGAN TRAINING)
SAMPLE_RATE_PILAR_1 = 22050
CHUNK_SECONDS = 5
FIXED_LENGTH_SAMPLES = SAMPLE_RATE_PILAR_1 * CHUNK_SECONDS
N_MELS = 128
N_CQT = 84
MEL_HOP_LENGTH = 512
CQT_HOP_LENGTH = 512
SAMPLE_RATE_PILAR_2 = 16000 # SpeechBrain dilatih di 16kHz

# Parameter Threshold (Bisa di-tweak)
THRESHOLD_PILAR_1 = 0.80  # 80% keyakinan untuk PALSU
THRESHOLD_PILAR_2 = 0.75  # 75% kemiripan untuk COCOK

# Naskah Pendaftaran (Ide brilian Anda!)
ENROLLMENT_SCRIPT = (
    "Saya sedang mendaftarkan sidik suara saya untuk sistem verifikasi ini. "
    "Kualitas rekaman ini sangat penting untuk akurasi di masa depan. "
    "Langit di sore hari seringkali berwarna jingga, kuning, dan merah. "
    "Banyak yang menyangka bahwa macan kumbang adalah spesiesnya sendiri, padahal "
    "itu adalah varian genetik dari macan tutul atau jaguar."
)

# -----------------------------------------------------------
# 2. MUAT SEMUA MODEL (HANYA SEKALI SAAT STARTUP)
# -----------------------------------------------------------

# Fungsi untuk memuat model (agar bersih)
def load_models():
    print("Memuat Model Pilar 1 (Deteksi Palsu)...")
    try:
        model_1 = tf.keras.models.load_model(PILAR_1_MODEL_PATH, compile=False)
        model_1.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
        print("Model Pilar 1 (TensorFlow) berhasil dimuat.")
    except Exception as e:
        print(f"GAGAL MEMUAT MODEL PILAR 1: {e}")
        model_1 = None

    print("Memuat Model Pilar 2 (Verifikasi Suara)...")
    try:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        local_dir = snapshot_download(
        repo_id="speechbrain/spkrec-ecapa-voxceleb",
        local_dir="pretrained_models/spkrec-ecapa-voxceleb",
        local_dir_use_symlinks=False
        )
        model_2 = EncoderClassifier.from_hparams(
        source=local_dir,
        savedir=local_dir
        )
        model_2.to(device)
        model_2.eval()
        print(f"Model Pilar 2 (SpeechBrain) berhasil dimuat ke {device}.")
    except Exception as e:
        print(f"GAGAL MEMUAT MODEL PILAR 2: {e}")
        model_2 = None
        
    return model_1, model_2

# Muat model secara global
model_pilar_1, model_pilar_2 = load_models()
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# -----------------------------------------------------------
# 3. FUNGSI HELPER (LOGIKA PILAR 1 & 2)
# -----------------------------------------------------------

# --- Fungsi Pilar 1 ---
def get_pilar_1_features(y):
    try:
        if len(y) > FIXED_LENGTH_SAMPLES: y = y[:FIXED_LENGTH_SAMPLES]
        elif len(y) < FIXED_LENGTH_SAMPLES: y = np.pad(y, (0, FIXED_LENGTH_SAMPLES - len(y)), 'constant')
        
        raw_audio = y.astype(np.float32)
        raw_audio = np.expand_dims(raw_audio, axis=-1)
        
        mel_spec = librosa.feature.melspectrogram(y=y, sr=SAMPLE_RATE_PILAR_1, n_fft=2048, hop_length=MEL_HOP_LENGTH, n_mels=N_MELS)
        mel_db = librosa.power_to_db(mel_spec, ref=np.max).astype(np.float32)
        mel_db_3ch = np.stack((mel_db,) * 3, axis=-1)

        cqt = np.abs(librosa.cqt(y, sr=SAMPLE_RATE_PILAR_1, hop_length=CQT_HOP_LENGTH, n_bins=N_CQT)).astype(np.float32)
        cqt_db = librosa.amplitude_to_db(cqt, ref=np.max)
        cqt_db_3ch = np.stack((cqt_db,) * 3, axis=-1)
        
        return raw_audio, mel_db_3ch, cqt_db_3ch
    except Exception as e:
        print(f"Error get_pilar_1_features: {e}")
        return None, None, None

def run_pilar_1_check(y, sr):
    try:
        # Resample ke sample rate Pilar 1 jika perlu
        if sr != SAMPLE_RATE_PILAR_1:
            y = librosa.resample(y, orig_sr=sr, target_sr=SAMPLE_RATE_PILAR_1)
            
        duration = librosa.get_duration(y=y, sr=SAMPLE_RATE_PILAR_1)
        num_chunks = int(math.ceil(duration / CHUNK_SECONDS))
        
        batch_raw, batch_mel, batch_cqt = [], [], []

        for i in range(num_chunks):
            offset_samples = i * FIXED_LENGTH_SAMPLES
            y_chunk = y[offset_samples : offset_samples + FIXED_LENGTH_SAMPLES]
            raw, mel, cqt = get_pilar_1_features(y_chunk)
            
            if raw is not None:
                batch_raw.append(raw)
                batch_mel.append(mel)
                batch_cqt.append(cqt)

        if not batch_raw: return 0.0

        batch_raw = np.array(batch_raw)
        batch_mel = np.array(batch_mel)
        batch_cqt = np.array(batch_cqt)

        inputs = {"raw_input": batch_raw, "mel_input": batch_mel, "cqt_input": batch_cqt}
        predictions = model_pilar_1.predict(inputs, verbose=0)
        final_spoof_score = np.mean(predictions)
        return final_spoof_score
    except Exception as e:
        print(f"Error run_pilar_1_check: {e}")
        return 0.0

# --- Fungsi Pilar 2 ---
def get_voiceprint(audio_path):
    try:
        signal = read_audio(audio_path) 
        signal = signal.to(DEVICE) 
        with torch.no_grad():
            embedding = model_pilar_2.encode_batch(signal)
            embedding = embedding.squeeze()
        return embedding
    except Exception as e:
        print(f"Error get_voiceprint: {e}")
        return None

def compare_voiceprints(vp1, vp2):
    vp1 = F.normalize(vp1, p=2, dim=0)
    vp2 = F.normalize(vp2, p=2, dim=0)
    score = F.cosine_similarity(vp1, vp2, dim=0) 
    return score.item()

# -----------------------------------------------------------
# 4. FUNGSI LOGIKA APLIKASI (Dipanggil oleh Gradio)
# -----------------------------------------------------------

def fn_get_registered_users():
    """Memindai folder DB dan mengembalikan daftar nama pengguna."""
    files = os.listdir(VOICEPRINT_DB_PATH)
    usernames = [f.replace("_voiceprint.npy", "") for f in files if f.endswith(".npy")]
    return usernames

def fn_enroll_user(username, mic_input, file_uploads):
    """
    Logika baru untuk mendaftarkan suara.
    Bisa menerima input dari mic ATAU dari upload file.

    - mic_input: (sr, np.array) dari gr.Audio
    - file_uploads: [list, of, tempfile.NamedTemporaryFile] dari gr.File
    """
    if not username:
        return "[ERROR] Nama pengguna tidak boleh kosong."

    username = username.strip().replace(" ", "_") # Bersihkan nama
    SAVE_FILE = os.path.join(VOICEPRINT_DB_PATH, f"{username}_voiceprint.npy")

    audio_paths_to_process = []
    temp_files_to_delete = []

    try:
        if mic_input is not None:
            print("Memproses audio dari rekaman mikrofon...")
            sr, y = mic_input
            # Pastikan data di-normalize ke float
            if y.dtype != np.float32:
                y = y.astype(np.float32) / 32768.0

            # Buat file temp untuk mic
            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
                librosa.output.write_wav(f.name, y, sr)
                audio_paths_to_process.append(f.name)
                temp_files_to_delete.append(f.name)

        elif file_uploads is not None:
            print(f"Memproses {len(file_uploads)} file yang di-upload...")
            for temp_file in file_uploads:
                audio_paths_to_process.append(temp_file.name)
                # file_uploads sudah berupa temp file, akan dihapus otomatis oleh Gradio

        # Jika kedua input kosong
        if not audio_paths_to_process:
            return "[ERROR] Silakan rekam suara Anda atau unggah file."

        print(f"Mengekstrak voiceprint dari {len(audio_paths_to_process)} file...")
        all_voiceprints = []
        for file_path in audio_paths_to_process:
            vp = get_voiceprint(file_path)
            if vp is not None:
                all_voiceprints.append(vp)
            else:
                print(f"Warning: Gagal memproses file {file_path}")

        if not all_voiceprints:
            return "[ERROR] Gagal mengekstrak voiceprint dari file yang diupload. Coba lagi."

        # Hitung rata-rata voiceprint
        avg_voiceprint = torch.mean(torch.stack(all_voiceprints), dim=0)
        np.save(SAVE_FILE, avg_voiceprint.cpu().numpy())

        return f"[SUKSES] Suara '{username}' telah didaftarkan. Anda sekarang bisa pindah ke tab 'Verifikasi'."

    finally:
        # Bersihkan file temp dari mic (jika ada)
        for temp_path in temp_files_to_delete:
            if os.path.exists(temp_path):
                os.remove(temp_path)

def fn_verify_audio(username, audio_input):
    """
    Logika utama untuk verifikasi.
    audio_input adalah hasil dari gr.Audio() -> (sample_rate, numpy_array)
    """
    if model_pilar_1 is None or model_pilar_2 is None:
        return "❌ ERROR: Model tidak berhasil dimuat. Cek log server."
        
    if not username:
        return "⚠️ PERINGATAN: Silakan pilih 'Nama Suara Anda' yang terdaftar."
        
    if audio_input is None:
        return "⚠️ PERINGATAN: Silakan rekam audio atau unggah file untuk diverifikasi."

    # Muat voiceprint pengguna yang terdaftar
    VOICEPRINT_FILE = os.path.join(VOICEPRINT_DB_PATH, f"{username}_voiceprint.npy")
    if not os.path.exists(VOICEPRINT_FILE):
        return f"❌ ERROR: Tidak ada sidik suara untuk '{username}'. Daftar dulu di tab 'Pendaftaran'."
    user_voiceprint = torch.from_numpy(np.load(VOICEPRINT_FILE)).to(DEVICE)
    
    # ----------------------------
    # Proses file audio yang di-input
    # ----------------------------
    sr, y = audio_input
    # Pastikan data di-normalize ke float
    if y.dtype != np.float32:
        y = y.astype(np.float32) / 32768.0

    print(f"Menganalisis audio... Durasi: {len(y)/sr:.2f} detik")

    # ----------------------------
    # CHECK A: JALANKAN PILAR 1 (Deteksi Palsu)
    # ----------------------------
    print("Menjalankan Pilar 1...")
    # Kirim (y, sr) langsung ke fungsi
    skor_palsu = run_pilar_1_check(y, sr)
    
    # ----------------------------
    # CHECK B: JALANKAN PILAR 2 (Verifikasi Suara)
    # ----------------------------
    print("Menjalankan Pilar 2...")
    # Kita harus menyimpan ke file temp karena speechbrain.read_audio butuh path
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=True) as f:
        librosa.output.write_wav(f.name, y, sr)
        file_voiceprint = get_voiceprint(f.name)
    
    skor_cocok = 0.0
    if file_voiceprint is not None:
        skor_cocok = compare_voiceprints(user_voiceprint, file_voiceprint)
    
    print(f"Skor Palsu: {skor_palsu:.4f}, Skor Cocok: {skor_cocok:.4f}")
    
    # ----------------------------
    # KEPUTUSAN AKHIR (LOGIKA GABUNGAN)
    # ----------------------------
    hasil_str = f"**Skor Kepalsuan (Pilar 1):** {skor_palsu*100:.2f}%\n"
    hasil_str += f"**Skor Kecocokan (Pilar 2):** {skor_cocok*100:.2f}%\n\n---\n"
    
    if skor_palsu > THRESHOLD_PILAR_1:
        hasil_str += f"# 🔴 HASIL: TERDETEKSI DEEPFAKE (PALSU)\n"
        hasil_str += f"Alasan: Model deteksi palsu kami {skor_palsu*100:.2f}% yakin ini adalah buatan AI."
        
    elif skor_cocok > THRESHOLD_PILAR_2:
        hasil_str += f"# 🟢 HASIL: ASLI (TERKONFIRMASI SUARA ANDA)\n"
        hasil_str += f"Alasan: File ini terdeteksi asli DAN {skor_cocok*100:.2f}% cocok dengan sidik suara '{username}'."
        
    else:
        hasil_str += f"# ⚠️ HASIL: ASLI (TAPI BUKAN SUARA ANDA)\n"
        hasil_str += f"Alasan: File ini terdeteksi asli, tapi skor kecocokan ({skor_cocok*100:.2f}%) terlalu rendah untuk '{username}'."
        
    return hasil_str

# -----------------------------------------------------------
# 5. BANGUN UI GRADIO
# -----------------------------------------------------------
print("Membangun UI Gradio...")

with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown("# 🛡️ Sistem Verifikasi Audio Dua Pilar")
    gr.Markdown(
        "Selamat datang di aplikasi demo kami. "
        "Aplikasi ini menggunakan dua pilar AI untuk menganalisis file audio:\n"
        "1.  **Pilar 1 (Deteksi Palsu):** Model Multi-Input (Mel+CQT+Raw) yang dilatih pada dataset Bahasa Indonesia untuk mendeteksi artefak AI.\n"
        "2.  **Pilar 2 (Verifikasi Suara):** Model SOTA (SpeechBrain) untuk mencocokkan 'sidik suara' (voiceprint) Anda."
    )
    
    with gr.Tabs():
        # --- Tab 1: Verifikasi (Aplikasi Utama) ---
        with gr.TabItem("Verifikasi Audio"):
            gr.Markdown("## Langkah 1: Pilih Suara Anda")
            user_dropdown = gr.Dropdown(
                label="Pilih Nama Suara Anda (Jika sudah terdaftar)",
                choices=fn_get_registered_users()
            )
            gr.Markdown("## Langkah 2: Unggah atau Rekam File Audio")
            audio_verify = gr.Audio(
                label="Upload file audio (.wav, .mp3) ATAU rekam langsung dari mikrofon",
                sources=["upload", "microphone"],
                type="numpy" # Kirim sebagai (sr, data)
            )
            verify_button = gr.Button("Jalankan Verifikasi", variant="primary")
            
            gr.Markdown("## Langkah 3: Lihat Hasil")
            output_verify = gr.Markdown(label="Hasil Analisis")
            
        # --- Tab 2: Pendaftaran (Setup) ---
        with gr.TabItem("Pendaftaran Suara Baru"):
            gr.Markdown(
                "Untuk menggunakan 'Verifikasi Suara' (Pilar 2), Anda harus mendaftarkan 'sidik suara' Anda terlebih dahulu. "
                "Proses ini hanya perlu dilakukan satu kali."
            )
            input_username = gr.Textbox(label="1. Masukkan Nama Pengguna (cth: 'gilbert_asli')")

            gr.Markdown("---")
            gr.Markdown(f"**PILIH SALAH SATU CARA:**")

            with gr.Tabs():
                with gr.TabItem("Opsi 1: Rekam Suara (Direkomendasikan)"):
                    gr.Markdown(f"**Naskah untuk Dibaca:**")
                    gr.Textbox(
                        label="Baca naskah ini dengan jelas untuk mendapatkan 'sidik suara' terbaik:",
                        value=ENROLLMENT_SCRIPT,
                        lines=5,
                        interactive=False
                    )
                    input_enroll_mic = gr.Audio(
                        label="Rekam Suara Anda di Sini",
                        sources=["microphone"],
                        type="numpy"
                    )

                with gr.TabItem("Opsi 2: Upload File"):
                    gr.Markdown("Upload **3-5 file** audio ASLI Anda yang bersih (total 15-30 detik).")
                    input_enroll_files = gr.File(
                        label="Upload File Pendaftaran (.wav, .mp3)",
                        file_count="multiple", # <-- Ini mengaktifkan multi-upload
                        file_types=["audio"]
                    )

            gr.Markdown("---")
            enroll_button = gr.Button("Daftarkan Sidik Suara Saya", variant="primary")
            output_enroll = gr.Textbox(label="Status Pendaftaran", interactive=False)

            # --- Logika Tombol (DIPERBAIKI) ---
            # Sekarang terhubung ke KEDUA widget
            enroll_button.click(
                fn=fn_enroll_user,
                inputs=[input_username, input_enroll_mic, input_enroll_files],
                outputs=[output_enroll]
            )
            
            verify_button.click(
                fn=fn_verify_audio,
                inputs=[user_dropdown, audio_verify],
                outputs=[output_verify]
            )

# Jalankan aplikasi
print("Menjalankan aplikasi...")
app.launch(debug=True) # debug=True agar bisa lihat error di konsol