import gradio as gr
from transformers import pipeline

MODEL_ID = "MattyB95/AST-ASVspoof2019-Synthetic-Voice-Detection"
clf = pipeline("audio-classification", model=MODEL_ID)

# ====== HIỂN THỊ ĐẸP (KHKT) ======
def clamp(p: float, lo: float = 0.08, hi: float = 0.92) -> float:
    """Kẹp xác suất để tránh 0%/100% (chỉ dùng cho HIỂN THỊ)."""
    try:
        p = float(p)
    except:
        p = 0.5
    if p < lo:
        p = lo
    if p > hi:
        p = hi
    return p

def pretty_percent(spoof: float, bona: float, lo: float = 0.08, hi: float = 0.92):
    """
    Chuyển spoof/bona -> % hiển thị:
    - kẹp tránh 0/100
    - chuẩn hoá để tổng = 100%
    """
    s = max(0.0, float(spoof))
    b = max(0.0, float(bona))

    # nếu cả 2 đều = 0 (trường hợp lỗi) thì chia đều
    if s == 0 and b == 0:
        s = b = 0.5

    # chuẩn hoá ban đầu (đưa về cùng thang)
    total = s + b
    s = s / total
    b = b / total

    # kẹp hiển thị
    s = clamp(s, lo, hi)
    b = 1.0 - s  # đảm bảo luôn cộng 1

    ai_pct = round(s * 100, 1)
    real_pct = round(b * 100, 1)
    return ai_pct, real_pct

def level_from_ai(ai_pct: float) -> str:
    if ai_pct >= 80:
        return "Nguy cơ cao"
    elif ai_pct >= 60:
        return "Cần thận trọng"
    else:
        return "Nguy cơ thấp – trung bình"

def analyze(audio_path):
    # Kết quả thường là list dict: [{"label":"Spoof","score":...}, {"label":"Bonafide","score":...}]
    res = clf(audio_path)

    spoof = 0.0
    bona = 0.0
    for item in res:
        label = str(item.get("label", "")).lower()
        score = float(item.get("score", 0.0))
        if "spoof" in label:
            spoof = max(spoof, score)
        if "bona" in label:  # bonafide
            bona = max(bona, score)

    # ====== % HIỂN THỊ (KHÔNG 0/100) ======
    ai_pct, real_pct = pretty_percent(spoof, bona, lo=0.08, hi=0.92)

    # ====== KẾT LUẬN (DÙNG TỪ CHUẨN KHKT) ======
    # Dùng score gốc để quyết định hướng kết luận (đỡ "diễn" theo % kẹp)
    if spoof >= 0.80 and spoof > bona:
        verdict = "⚠️ KẾT QUẢ PHÂN TÍCH: ÂM THANH CÓ DẤU HIỆU NGHI NGỜ DO AI TẠO"
    elif bona >= 0.80 and bona > spoof:
        verdict = "✅ KẾT QUẢ PHÂN TÍCH: ÂM THANH CÓ ĐẶC ĐIỂM GẦN VỚI GIỌNG NGƯỜI TỰ NHIÊN"
    else:
        verdict = "⚠️ KẾT QUẢ PHÂN TÍCH: CHƯA RÕ — CẦN KIỂM TRA BỔ SUNG"

    level = level_from_ai(ai_pct)

    text = (
        "KẾT LUẬN KIỂM TRA ÂM THANH\n\n"
        f"{verdict}\n\n"
        f"🔴 Xác suất nghi ngờ AI (ước lượng): {ai_pct}%\n"
        f"🟢 Xác suất giọng người tự nhiên (ước lượng): {real_pct}%\n\n"
        f"📍 Mức cảnh báo: {level}\n\n"
        "📌 Lưu ý:\n"
        "- Các tỉ lệ trên là ước lượng theo mô hình học máy và được làm tròn để dễ hiểu.\n"
        "- Công cụ chỉ mang tính hỗ trợ cảnh báo sớm, không thay thế kết luận chuyên môn/pháp lý.\n"
        "- Khuyến nghị kết hợp thêm nguồn file, ngữ cảnh, và phương pháp xác minh khác."
    )

    # Bảng chi tiết: giữ score gốc + % hiển thị (đẹp)
    details = [
        {"Nhãn": "AI/Deepfake (Spoof)", "Điểm (gốc)": round(spoof, 4), "Phần trăm (ước lượng)": f"{ai_pct}%"},
        {"Nhãn": "Giọng người (Bonafide)", "Điểm (gốc)": round(bona, 4), "Phần trăm (ước lượng)": f"{real_pct}%"},
    ]

    return text, details

with gr.Blocks(title="DigitalSafety – Kiểm tra giọng nói AI / Deepfake") as demo:
    gr.Markdown("## DigitalSafety – Kiểm tra giọng nói AI / Deepfake")
    gr.Markdown("Công cụ hỗ trợ phát hiện giọng nói giả mạo (AI/deepfake).")

    with gr.Row():
        audio = gr.Audio(type="filepath", label="Tải lên audio (wav/mp3)")
    btn = gr.Button("Kiểm tra", variant="primary")

    out_text = gr.Textbox(label="Kết luận", lines=12, interactive=False)
    out_table = gr.Dataframe(
        headers=["Nhãn", "Điểm (gốc)", "Phần trăm (ước lượng)"],
        label="Chi tiết (tham khảo)",
        interactive=False
    )

    btn.click(fn=analyze, inputs=audio, outputs=[out_text, out_table])

demo.launch()