"""ModelBrew Dataset Cleaner — Hugging Face Space.

Upload a fine-tuning dataset (.jsonl/.json/.csv), get a scored issue report
and a cleaned export with critical rows dropped.
"""
import tempfile
from pathlib import Path

import gradio as gr
import pandas as pd

from modelbrew_cleaner import Severity, clean_file, export_clean, issue_summary

MAX_ROWS = 5000
SAMPLE = Path(__file__).parent / "sample.jsonl"

SEV_ORDER = {"critical": 0, "warning": 1, "suggestion": 2}


def analyze(path: str):
    rows = clean_file(path)
    if len(rows) > MAX_ROWS:
        raise gr.Error(
            f"This demo caps at {MAX_ROWS:,} rows (got {len(rows):,}). "
            f"Run locally with `pip install modelbrew-cleaner` for unlimited size."
        )
    records = [
        {
            "row": r.row_index,
            "severity": i.severity.value,
            "check": i.code,
            "message": i.message,
            "auto-fixable": bool(i.auto_fixable),
        }
        for r in rows
        for i in r.issues
    ]
    records.sort(key=lambda x: (SEV_ORDER.get(x["severity"], 9), x["row"]))
    df = pd.DataFrame(records, columns=["row", "severity", "check", "message", "auto-fixable"])

    s = issue_summary(rows)
    n_critical_rows = sum(
        1 for r in rows if any(i.severity == Severity.critical for i in r.issues)
    )
    summary_md = (
        f"### {len(rows):,} rows scanned\n"
        f"- 🔴 **{s['critical']} critical** issues ({n_critical_rows} rows dropped in the cleaned export)\n"
        f"- 🟡 **{s['warning']} warnings**\n"
        f"- 🔵 **{s['suggestion']} suggestions**\n"
    )

    cleaned = export_clean(rows)
    out = tempfile.NamedTemporaryFile(
        mode="w", suffix=".cleaned.jsonl", delete=False, encoding="utf-8"
    )
    out.write(cleaned)
    out.close()
    return summary_md, df, out.name


def analyze_upload(file):
    if file is None:
        raise gr.Error("Upload a .jsonl, .json, or .csv file — or click 'Try the sample'.")
    return analyze(file.name if hasattr(file, "name") else str(file))


def analyze_sample():
    return analyze(str(SAMPLE))


with gr.Blocks(title="ModelBrew Dataset Cleaner") as demo:
    gr.Markdown(
        "# 🧹 ModelBrew Dataset Cleaner\n"
        "90+ quality checks for fine-tuning datasets: PII with checksum validation, "
        "exact/near duplicates, prompt-injection & jailbreak patterns, label errors, "
        "truncated responses, and more. Critical rows are dropped from the cleaned export.\n\n"
        "`pip install modelbrew-cleaner` to run locally on datasets of any size."
    )
    with gr.Row():
        upload = gr.File(label="Dataset (.jsonl / .json / .csv)", file_types=[".jsonl", ".json", ".csv"])
    with gr.Row():
        run_btn = gr.Button("Clean my dataset", variant="primary")
        sample_btn = gr.Button("Try the sample (deliberately dirty)")
    summary = gr.Markdown()
    issues = gr.Dataframe(label="Issues found", interactive=False, wrap=True)
    cleaned_file = gr.File(label="Cleaned dataset (critical rows dropped)")

    run_btn.click(analyze_upload, inputs=upload, outputs=[summary, issues, cleaned_file])
    sample_btn.click(analyze_sample, inputs=None, outputs=[summary, issues, cleaned_file])

    gr.Markdown(
        "---\n"
        "Built by [ModelBrew](https://modelbrew.ai) — we work on fine-tuning without "
        "catastrophic forgetting (patent-pending CRMA adapters), and honest measurement "
        "starts with clean data. Read: *Your forgetting benchmark is lying to you* "
        "({{HF_ARTICLE_URL}}) · [pip package]({{PYPI_URL}})"
    )

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