import os
import subprocess
import tempfile
import pandas as pd
import gradio as gr
from huggingface_hub import HfApi, hf_hub_download
import matplotlib.pyplot as plt

MAX_MODELS = 4

# --------------------------------------------------
# HF
# --------------------------------------------------

import shutil

print("perplexity:", shutil.which("llama-perplexity"))
print("quantize:", shutil.which("llama-quantize"))
print("cli:", shutil.which("llama-cli"))




api = HfApi()
import os
import subprocess



def list_ggufs(repo_id):
    try:
        files = api.list_repo_files(repo_id)

        ggufs = sorted(
            [
                f
                for f in files
                if f.lower().endswith(".gguf")
            ]
        )

        return (
            gr.update(choices=ggufs),
            gr.update(choices=ggufs),
            gr.update(choices=ggufs),
            gr.update(choices=ggufs),
            f"Found {len(ggufs)} GGUF files"
        )

    except Exception as e:
        return (
            gr.update(choices=[]),
            gr.update(choices=[]),
            gr.update(choices=[]),
            gr.update(choices=[]),
            f"Error: {e}"
        )


# --------------------------------------------------
# PPL
# --------------------------------------------------

def run_perplexity(model_path, dataset_path):

    cmd = [
        "llama-perplexity",
        "-m",
        model_path,
        "-f",
        dataset_path,
    ]

    result = subprocess.run(
        cmd,
        capture_output=True,
        text=True,
    )

    output = result.stdout + result.stderr

    ppl = None

    for line in output.splitlines():

        if "perplexity" in line.lower():

            parts = line.split()

            for token in reversed(parts):

                try:
                    ppl = float(token)
                    break
                except:
                    pass

    if ppl is None:
        raise RuntimeError(
            f"Unable to parse perplexity\n{output}"
        )

    return ppl


# --------------------------------------------------
# BENCHMARK
# --------------------------------------------------

def benchmark(
    repo_id,
    gguf1,
    gguf2,
    gguf3,
    gguf4,
):

    selected = [
        x
        for x in [gguf1, gguf2, gguf3, gguf4]
        if x
    ]

    if len(selected) == 0:
        raise gr.Error(
            "Select at least one GGUF"
        )

    dataset_path = "wikitext.txt"

    rows = []

    for gguf in selected:

        model_path = hf_hub_download(
            repo_id=repo_id,
            filename=gguf,
        )

        size_gb = (
            os.path.getsize(model_path)
            / 1024**3
        )

        ppl = run_perplexity(
            model_path,
            dataset_path,
        )

        rows.append(
            {
                "Model": os.path.basename(gguf),
                "Perplexity": round(ppl, 4),
                "Size GB": round(size_gb, 2),
            }
        )

    df = pd.DataFrame(rows)

    best_ppl = df["Perplexity"].min()

    df["Quality Score"] = (
        best_ppl
        / df["Perplexity"]
        * 100
    ).round(2)

    df = df.sort_values(
        "Perplexity"
    )

    fig = plt.figure(
        figsize=(8, 5)
    )

    plt.scatter(
        df["Size GB"],
        df["Perplexity"],
    )

    for _, row in df.iterrows():

        plt.annotate(
            row["Model"],
            (
                row["Size GB"],
                row["Perplexity"],
            ),
        )

    plt.xlabel("Size (GB)")
    plt.ylabel("Perplexity")
    plt.title(
        "GGUF Quality vs Size"
    )

    return df, fig


# --------------------------------------------------
# UI
# --------------------------------------------------

with gr.Blocks(
    title="GGUF Perplexity Benchmark"
) as demo:

    gr.Markdown(
        """
# 📊 GGUF Perplexity Benchmark

Compare GGUF variants using llama-perplexity.
"""
    )

    repo_id = gr.Textbox(
        label="HF Repository",
        placeholder="rahul7star/my-gguf",
        value="unsloth/gemma-4-E2B-it-GGUF"
    )

    scan_btn = gr.Button(
        "Scan Repository"
    )

    status = gr.Textbox(
        label="Status"
    )

    with gr.Row():

        gguf1 = gr.Dropdown(
            label="Variant 1"
        )

        gguf2 = gr.Dropdown(
            label="Variant 2"
        )

    with gr.Row():

        gguf3 = gr.Dropdown(
            label="Variant 3"
        )

        gguf4 = gr.Dropdown(
            label="Variant 4"
        )

    run_btn = gr.Button(
        "Run Benchmark",
        variant="primary"
    )

    results = gr.Dataframe(
        label="Results"
    )

    chart = gr.Plot(
        label="Quality vs Size"
    )

    scan_btn.click(
        fn=list_ggufs,
        inputs=repo_id,
        outputs=[
            gguf1,
            gguf2,
            gguf3,
            gguf4,
            status,
        ],
    )

    run_btn.click(
        fn=benchmark,
        inputs=[
            repo_id,
            gguf1,
            gguf2,
            gguf3,
            gguf4,
        ],
        outputs=[
            results,
            chart,
        ],
    )

demo.launch()