import gradio as gr
from quantguard_engine import (
    MODEL_PROFILES,
    QUANTIZER_REGISTRY,
    DATASET_REGISTRY,
    ATTACK_REGISTRY,
    DEFENSE_REGISTRY,
    run_group_quantization_experiment,
    build_comparison_dataframe,
    plot_quantguard_diagnostics,
    inspect_hf_model_repo,
    generate_plug_and_play_code,
)


def run_sandbox(model_key, quantizer_key, attack_key, defense_key, dataset_key, bits, group_size):
    bits = int(bits)
    group_size = int(group_size)
    res_unprotected = run_group_quantization_experiment(
        model_key=model_key,
        quantizer_key=quantizer_key,
        attack_key=attack_key,
        defense_key="none",
        dataset_key=dataset_key,
        bits=bits,
        group_size=group_size,
    )
    res_selected = run_group_quantization_experiment(
        model_key=model_key,
        quantizer_key=quantizer_key,
        attack_key=attack_key,
        defense_key=defense_key,
        dataset_key=dataset_key,
        bits=bits,
        group_size=group_size,
    )
    df_comp = build_comparison_dataframe(
        model_key=model_key,
        quantizer_key=quantizer_key,
        attack_key=attack_key,
        dataset_key=dataset_key,
        bits=bits,
        group_size=group_size,
    )
    fig = plot_quantguard_diagnostics(res_unprotected, res_selected, df_comp)

    asr_drop = round(res_unprotected["Backdoor_ASR_%"] - res_selected["Backdoor_ASR_%"], 1)
    flip_drop = round(res_unprotected["ConditionalFlip_%"] - res_selected["ConditionalFlip_%"], 1)
    zero_drop = round(res_unprotected["Zero_Collapse_%"] - res_selected["Zero_Collapse_%"], 1)
    vram_saved = round(res_unprotected["VRAM_GB"] / res_selected["VRAM_GB"], 1) if res_selected["VRAM_GB"] else 0

    def _chip(ok, text):
        color, bg = ("#166534", "#dcfce7") if ok else ("#991b1b", "#fee2e2")
        return f'<span style="background:{bg}; color:{color}; border:1px solid {color}22; border-radius:999px; padding:2px 10px; font-weight:700; white-space:nowrap;">{text}</span>'

    highlights_html = """
    <div style="background:#f8fafc; border:1px solid #e2e8f0; border-radius:10px; padding:10px 14px; margin:8px 0;">
      <b>✨ Điểm nổi bật khi bật phòng thủ này (so với không bảo vệ):</b>
      <div style="display:flex; flex-wrap:wrap; gap:8px; margin-top:8px;">
        {asr_chip} {flip_chip} {zero_chip} {vram_chip}
      </div>
    </div>
    """.format(
        asr_chip=_chip(asr_drop > 0, f"🛡️ ASR giảm {asr_drop} điểm %"),
        flip_chip=_chip(flip_drop > 0, f"⚖️ Safety-flip giảm {flip_drop} điểm %"),
        zero_chip=_chip(True, f"📉 Zero-collapse {res_selected['Zero_Collapse_%']}%"),
        vram_chip=_chip(vram_saved >= 3.5, f"💾 Nén ~{vram_saved}× so với FP16"),
    )

    summary_md = f"""
### 🧪 Kết quả MÔ PHỎNG (`{res_selected['Model']}` | `{res_selected['Bits']}`)
> ⚠️ **Minh bạch:** đây là mô phỏng trên tensor tổng hợp (ma trận ngẫu nhiên, seed cố định) —
> **không phải** kết quả chạy trên model thật. Các chỉ số PPL/CACC/ASR minh hoạ cơ chế hoạt động,
> không phải benchmark.
{highlights_html}
- **Thuật toán Lượng tử hóa (`QuantizerBackend`):** `{res_selected['Quantizer']}`
- **Kịch bản Tấn công (`AttackPlugin`):** `{res_selected['Attack']}`
- **Cơ chế Phòng thủ đang chọn (`DefensePlugin`):** `{res_selected['Defense']}`
- **Bộ dữ liệu Đánh giá (`DatasetAdapter`):** `{res_selected['Dataset']}` *(chỉ dùng làm nhãn mô phỏng — không có dữ liệu thật nào được tải)*

| Chỉ số (Mô phỏng) | Khi KHÔNG Bảo vệ (`None`) | Khi Bật `{defense_key}` | Ghi chú |
| :--- | :---: | :---: | :--- |
| **Max Group Outlier Ratio** | `{res_unprotected['Max_Outlier_Ratio']}x` | **`{res_selected['Max_Outlier_Ratio']}x`** | > 15x = nhóm có outlier bất thường |
| **Perplexity (`PPL` ↓)** | `{res_unprotected['PPL']}` | **`{res_selected['PPL']}`** | Mô phỏng chất lượng ngôn ngữ |
| **Clean Accuracy (`CACC` % ↑)** | `{res_unprotected['Clean_Acc_%']}%` | **`{res_selected['Clean_Acc_%']}%`** | Mô phỏng độ chính xác trên tác vụ sạch |
| **Tỷ lệ Lật An toàn (`ConditionalFlip` % ↓)** | `{res_unprotected['ConditionalFlip_%']}%` | **`{res_selected['ConditionalFlip_%']}%`** | Mô phỏng lỗi mất từ chối an toàn |
| **Tỷ lệ Kích hoạt Cửa sau (`Backdoor ASR` % ↓)** | `{res_unprotected['Backdoor_ASR_%']}%` | **`{res_selected['Backdoor_ASR_%']}%`** | Mô phỏng nguy cơ cửa sau lượng tử hóa |
| **Tỷ lệ Sụp đổ Trọng số về 0 (`Zero-Collapse` % ↓)** | `{res_unprotected['Zero_Collapse_%']}%` | **`{res_selected['Zero_Collapse_%']}%`** | Đo thật trên tensor mô phỏng |
| **Bộ nhớ GPU ước tính (`VRAM`)** | `{res_unprotected['VRAM_GB']} GB` | **`{res_selected['VRAM_GB']} GB`** | Công thức `bits/16 × FP16`, không đo thật |

🔧 *Chẩn đoán mẫu (preset mô phỏng của model đã chọn):* `{res_selected['Failure_Mode']}` → Khuyến nghị: `{res_selected['Recommended_Fix']}`
"""
    cli_code, py_code = generate_plug_and_play_code(
        model_key, quantizer_key, attack_key, defense_key, dataset_key, bits, group_size
    )
    return summary_md, fig, df_comp, cli_code, py_code


with gr.Blocks(title="QuantGuard - Plug-and-Play LLM Quantization Security", theme=gr.themes.Soft()) as demo:
    # ─── HERO: giá trị cốt lõi & điểm mạnh ───
    gr.HTML(
        """
        <div style="text-align:center; padding:18px 10px 6px 10px;">
          <h1 style="margin:0; font-size:1.9em;">🛡️ QuantGuard</h1>
          <p style="margin:6px 0 0 0; font-size:1.05em; color:#475569;">
            Kiểm toán & bảo vệ <b>Plug-and-Play</b> cho LLM lượng tử hóa —
            quét trọng số <b>thật</b> trên Hugging Face Hub chỉ trong vài giây,
            không cần tải nguyên model.
          </p>
        </div>
        <div style="display:flex; flex-wrap:wrap; gap:10px; justify-content:center; margin:14px 0 4px 0;">
          <div style="flex:1; min-width:170px; max-width:230px; background:#ecfdf5; border:1px solid #10b981; border-radius:10px; padding:10px 12px;">
            <div style="font-size:1.5em;">📐</div>
            <b>Đo trên trọng số thật</b>
            <div style="font-size:0.85em; color:#065f46; margin-top:4px;">
              Tab 2 đọc header + vài tensor nhỏ qua HTTP Range — đo Outlier Ratio & Zero-Collapse thực tế, không mô phỏng.
            </div>
          </div>
          <div style="flex:1; min-width:170px; max-width:230px; background:#eff6ff; border:1px solid #3b82f6; border-radius:10px; padding:10px 12px;">
            <div style="font-size:1.5em;">⚡</div>
            <b>Siêu nhẹ & nhanh</b>
            <div style="font-size:0.85em; color:#1e40af; margin-top:4px;">
              Tải vài MB thay vì cả model (~GB). Chạy trên CPU-free Space, không cần GPU.
            </div>
          </div>
          <div style="flex:1; min-width:170px; max-width:230px; background:#fefce8; border:1px solid #eab308; border-radius:10px; padding:10px 12px;">
            <div style="font-size:1.5em;">🧩</div>
            <b>Plug-and-Play 5 slot</b>
            <div style="font-size:0.85em; color:#854d0e; margin-top:4px;">
              Hoán đổi tự do: Model × Quantizer × Attack × Defense × Dataset — so sánh 5 chiến lược phòng thủ trong 1 lần chạy.
            </div>
          </div>
          <div style="flex:1; min-width:170px; max-width:230px; background:#faf5ff; border:1px solid #a855f7; border-radius:10px; padding:10px 12px;">
            <div style="font-size:1.5em;">📦</div>
            <b>Coverage ~76% Hub</b>
            <div style="font-size:0.85em; color:#6b21a8; margin-top:4px;">
              ~327k model text-generation full-precision audit được thật; ~10k AWQ/GPTQ ra verdict từ metadata.
            </div>
          </div>
        </div>
        <p style="text-align:center; font-size:0.85em; color:#64748b; margin:8px 0 0 0;">
          Tham chiếu cơ chế group-quantization theo <a href="https://github.com/mit-han-lab/llm-awq" target="_blank">mit-han-lab/llm-awq</a> và AWQ (<a href="https://arxiv.org/abs/2306.00978" target="_blank">Lin et al., MLSys 2024</a>)
          · Lưu ý minh bạch: Tab 1 là <b>mô phỏng tensor tổng hợp</b>, Tab 2 là <b>trọng số thật</b> — xem chi tiết ở tab 📖 Hướng dẫn.
        </p>
        """
    )

    with gr.Tabs():
        # TAB 0: HƯỚNG DẪN SỬ DỤNG (DEMO)
        with gr.TabItem("📖 Hướng dẫn sử dụng (Demo)"):
            gr.Markdown(
                """
                ## 🚀 Bắt đầu trong 60 giây — 3 bước

                ### Bước 1 — 🔍 Quét trọng số thật (điểm mạnh nhất của QuantGuard)
                1. Mở tab **"🔍 2. Quét Model trên Hugging Face Hub (Trọng số thật)"**
                2. Giữ nguyên model mặc định `Qwen/Qwen2.5-0.5B-Instruct` (hoặc nhập model bất kỳ, ví dụ `HuggingFaceTB/SmolLM2-135M-Instruct`)
                3. Bấm **🔎 Quét Trọng số Thật** — sau vài giây bạn thấy:
                   - **Max Group Outlier Ratio** đo trên tensor thật (ngưỡng heuristic `> 15x` = cần xem xét thủ công)
                   - **Zero-Collapse Rate** thực tế dưới RTN INT4 (g128)
                   - Verdict **✅ PASS / ⚠️ REVIEW** dựa trên số đo thật, không phải giả lập

                > 💡 *Vì sao nhanh?* QuantGuard chỉ tải safetensors header + vài tensor 2D nhỏ (HTTP Range),
                > vài MB thay vì cả model — model 7B cũng quét được trong chốc lát.

                ### Bước 2 — 🔬 Hiểu cơ chế qua Sandbox mô phỏng
                1. Mở tab **"🔬 1. Plug-and-Play Security Sandbox"**
                2. Bấm **🚀 Chạy Mô phỏng** (đã chọn sẵn cấu hình đề xuất) — kết quả hiện ngay:
                   - Heatmap **trước/sau** khi bật defense: thấy rõ trọng số sụp về `0` (Zero-Collapse)
                   - Biểu đồ so sánh **5 Defense Plugins** — quan sát `quantguard_full` giảm ASR mạnh nhất
                3. Thử hoán đổi **AttackPlugin** (ví dụ `agentq_multicodebook`) hoặc **Số Bit** (`3`) rồi chạy lại —
                   đây là cách khám phá "slot" plug-and-play.

                > ⚠️ Tab này là **mô phỏng trên ma trận ngẫu nhiên (seed cố định)** minh hoạ cơ chế AWQ —
                > các con số PPL/CACC/ASR mang tính minh hoạ, KHÔNG phải benchmark trên model thật.

                ### Bước 3 — 💻 Xuất code chạy được
                1. Mở tab **"💻 3. Xuất Lệnh CLI & Python SDK"**
                2. Copy lệnh `pip install` + `huggingface-cli download` (tải engine 1-file từ Space này)
                3. Chạy script Python sinh sẵn — gọi trực tiếp `inspect_hf_model_repo()` để audit
                   trọng số thật ngay trên máy bạn, offline khỏi Space.

                ---

                ### 🗺️ Model nào dùng được?

                | Nhóm | Định dạng | Số lượng trên Hub | QuantGuard |
                | :--- | :--- | :---: | :--- |
                | 🟢 Full-precision safetensors (F16/BF16) | `.safetensors` thường | **≈ 327k (~76%)** | ✅ Audit trọng số thật đầy đủ |
                | 🟡 AWQ / GPTQ (int4 đóng gói) | `.safetensors` đóng gói int32 | **≈ 10k (~2.3%)** | ⚠️ Verdict từ metadata (tags, config) |
                | 🔴 GGUF thuần | `.gguf` | **≈ 41k (~9.6%)** | ❌ Chưa hỗ trợ parser |

                💡 **Khuyên dùng thử ngay với:** `Qwen/Qwen2.5-0.5B-Instruct` · `HuggingFaceTB/SmolLM2-135M-Instruct` · `Qwen/Qwen2.5-1.5B-Instruct`
                """
            )

        # TAB 1: PLUG-AND-PLAY SANDBOX (SYNTHETIC SIMULATION)
        with gr.TabItem("🔬 1. Plug-and-Play Security Sandbox (Mô phỏng tổng hợp)"):
            gr.Markdown(
                "> 🧪 **Tab này chạy mô phỏng trên ma trận ngẫu nhiên (seed cố định)** — các số liệu minh hoạ "
                "cơ chế group-quantization của AWQ và tác động của từng Defense, không phải kết quả đo trên model thật."
            )
            with gr.Row():
                with gr.Column(scale=1):
                    model_dd = gr.Dropdown(
                        choices=list(MODEL_PROFILES.keys()),
                        value="Qwen/Qwen2.5-1.5B-Instruct",
                        label="1. Chọn Mô hình Mục tiêu (preset mô phỏng)",
                    )
                    quant_dd = gr.Dropdown(
                        choices=[(v["name"], k) for k, v in QUANTIZER_REGISTRY.items()],
                        value="awq",
                        label="2. Chọn Thuật toán Lượng tử hóa (QuantizerBackend)",
                    )
                    dataset_dd = gr.Dropdown(
                        choices=[(v["name"], k) for k, v in DATASET_REGISTRY.items()],
                        value="advbench_wikitext2",
                        label="3. Chọn Bộ dữ liệu Đánh giá (DatasetAdapter - nhãn mô phỏng)",
                    )
                    attack_dd = gr.Dropdown(
                        choices=[(v["name"], k) for k, v in ATTACK_REGISTRY.items()],
                        value="outlier_injection",
                        label="4. Chọn Kiểu Tấn công / Rủi ro (AttackPlugin)",
                    )
                    defense_dd = gr.Dropdown(
                        choices=[(v["name"], k) for k, v in DEFENSE_REGISTRY.items()],
                        value="quantguard_full",
                        label="5. Chọn Kiểu Phòng thủ (DefensePlugin)",
                    )
                    with gr.Row():
                        bits_radio = gr.Radio(choices=[3, 4, 8], value=4, label="Số Bit Lượng tử hóa (W-bit)")
                        group_radio = gr.Radio(choices=[32, 64, 128], value=128, label="Kích thước Nhóm (Group Size G)")
                    run_btn = gr.Button("🚀 Chạy Mô phỏng Plug-and-Play", variant="primary")

                with gr.Column(scale=2):
                    summary_out = gr.Markdown()
                    plot_out = gr.Plot(label="Trực quan hóa Sụp đổ Trọng số (Zero-Collapse) & So sánh 5 Cơ chế Phòng thủ")

            gr.Markdown("### 📋 Bảng Ma trận Đối sánh Toàn bộ 5 Defense Plugins trên cùng Cấu hình (mô phỏng)")
            df_out = gr.Dataframe(interactive=False)

        # TAB 2: LIVE HUGGING FACE REPO SCANNER (REAL WEIGHTS)
        with gr.TabItem("🔍 2. Quét Model trên Hugging Face Hub (Trọng số thật) & Bộ 20 Prompt PCR"):
            gr.Markdown(
                """
                Nhập bất kỳ **Model ID trên Hugging Face Hub** (ví dụ: `Qwen/Qwen2.5-1.5B-Instruct`, `Qwen/Qwen2.5-0.5B-Instruct`, `Qwen/Qwen2.5-3B-Instruct`)
                để quét **trọng số THẬT**: tool chỉ tải header safetensors và một vài tensor 2D nhỏ qua **HTTP Range requests**
                (không tải nguyên model), sau đó đo **Group Outlier Ratio** ($\\max(|W_g|)/\\text{median}(|W_g|)$) và
                **Zero-Collapse rate** thực tế dưới RTN INT4 (g128) trên chính trọng số đó.

                > ℹ️ Ngưỡng `> 15x` là heuristic tham khảo — outlier channel tự nhiên trong LLM cũng có thể rất lớn;
                > kết quả vượt ngưỡng nghĩa là **cần xem xét thủ công**, không phải kết luận chắc chắn có backdoor.
                """
            )
            with gr.Row():
                hf_repo_input = gr.Textbox(
                    value="Qwen/Qwen2.5-0.5B-Instruct",
                    label="Hugging Face Model Repo ID",
                    scale=3,
                )
                sim_attack_chk = gr.Checkbox(
                    value=False,
                    label="Kèm DEMO mô phỏng Outlier-Injection (dữ liệu tổng hợp, KHÔNG phải trọng số của repo)",
                    scale=2,
                )
                scan_btn = gr.Button("🔎 Quét Trọng số Thật trên Hugging Face", variant="primary", scale=1)

            hf_report_out = gr.Markdown()
            pcr_table_out = gr.Dataframe(label="Bộ 20 Câu hỏi Hiệu chuẩn PCR (PCR Calibration Prompts)")

        # TAB 3: CODE & CLI EXPORT
        with gr.TabItem("💻 3. Xuất Lệnh CLI & Python SDK (chạy được thật)"):
            gr.Markdown(
                """
                Đoạn mã dưới đây được sinh tự động theo đúng cấu hình Plug-and-Play bạn đã chọn ở Tab 1.

                **Code này chạy được thật**: nó cài đúng dependencies của engine, tải module
                `quantguard_engine.py` từ Space [banhcarrot/quantguard-simulation](https://huggingface.co/spaces/banhcarrot/quantguard-simulation)
                trên Hub, rồi gọi trực tiếp các hàm có thật:
                `inspect_hf_model_repo` (audit trọng số thật) và `run_group_quantization_experiment` (mô phỏng).
                Không có package `quantguard` bên ngoài nào — mọi thứ nằm trong 1 file engine.
                """
            )
            cli_box = gr.Code(language="shell", label="Command-Line Interface (CLI)")
            py_box = gr.Code(language="python", label="Python Plug-and-Play API")

    run_btn.click(
        fn=run_sandbox,
        inputs=[model_dd, quant_dd, attack_dd, defense_dd, dataset_dd, bits_radio, group_radio],
        outputs=[summary_out, plot_out, df_out, cli_box, py_box],
    )
    scan_btn.click(
        fn=inspect_hf_model_repo,
        inputs=[hf_repo_input, sim_attack_chk],
        outputs=[hf_report_out, pcr_table_out],
    )
    demo.load(
        fn=run_sandbox,
        inputs=[model_dd, quant_dd, attack_dd, defense_dd, dataset_dd, bits_radio, group_radio],
        outputs=[summary_out, plot_out, df_out, cli_box, py_box],
    )
    demo.load(
        fn=inspect_hf_model_repo,
        inputs=[hf_repo_input, sim_attack_chk],
        outputs=[hf_report_out, pcr_table_out],
    )

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