#!/usr/bin/env python3
"""Run the public ModelSentry Space against a fixed cross-section of Hub repositories."""

from __future__ import annotations

import argparse
import json
import statistics
import time
import urllib.error
import urllib.request
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


BASE_URL = "https://seedofevil-modelsentry-hf-auditor.hf.space"
TARGETS = [
    # SeedOfEvil public inventory.
    {"category": "seed", "target": "model:SeedOfEvil/Qwen-Rapid-AIO-v18-NSFW-diffusers"},
    {"category": "seed", "target": "space:SeedOfEvil/Pro-Realism-Edit-Studio"},
    {"category": "seed", "target": "space:SeedOfEvil/Pro-Realism-Edit-Studio-Qwen-2511"},
    {"category": "seed", "target": "space:SeedOfEvil/Wan2.1-T2V-1.3B-Local"},
    {"category": "seed", "target": "space:SeedOfEvil/TRELLIS-2-Local-Privacy-Mirror"},
    {"category": "seed", "target": "space:SeedOfEvil/ReplayForge"},
    {"category": "seed", "target": "space:SeedOfEvil/ModelSentry-HF-Auditor"},
    # Popular models spanning embeddings, text generation, time series, and encoders.
    {"category": "popular_model", "target": "model:sentence-transformers/all-MiniLM-L6-v2"},
    {"category": "popular_model", "target": "model:google-bert/bert-base-uncased"},
    {"category": "popular_model", "target": "model:BAAI/bge-m3"},
    {"category": "popular_model", "target": "model:Qwen/Qwen3-0.6B"},
    {"category": "popular_model", "target": "model:google-t5/t5-small"},
    {"category": "popular_model", "target": "model:amazon/chronos-2"},
    # Popular Spaces with varied repository size and application structure.
    {"category": "popular_space", "target": "space:open-llm-leaderboard/open_llm_leaderboard"},
    {"category": "popular_space", "target": "space:mteb/leaderboard"},
    {"category": "popular_space", "target": "space:microsoft/TRELLIS"},
    {"category": "popular_space", "target": "space:facebook/MusicGen"},
    {"category": "popular_space", "target": "space:enzostvs/deepsite"},
    {"category": "popular_space", "target": "space:jbilcke-hf/ai-comic-factory"},
    # Recent low-visibility controls discovered on the benchmark date.
    {"category": "low_visibility_model", "target": "model:stage-babylm/llama-64-1L"},
    {"category": "low_visibility_model", "target": "model:ganeshkota/ganeshkotaslu-sentiment-analysis"},
    {"category": "low_visibility_model", "target": "model:ganeshkota/ganeshkotaslu-incident-triage-copilot"},
    {"category": "low_visibility_space", "target": "space:kamalika19/modelpractical"},
    {"category": "low_visibility_space", "target": "space:mbeigih/paddy-leaf-disease-classifier"},
    {"category": "low_visibility_space", "target": "space:tylerjharden/adapt1-fle-research-preview"},
]

NEGATIVE_CASES = [
    {"name": "external_host", "target": "https://example.com/owner/repository"},
    {"name": "missing_repository", "target": "model:ModelSentryBenchmark/does-not-exist"},
]


def request_json(url: str, data: dict[str, Any] | None = None, timeout: int = 150) -> Any:
    body = json.dumps(data).encode("utf-8") if data is not None else None
    headers = {"User-Agent": "ModelSentry-Live-Benchmark/1.0"}
    if body is not None:
        headers["Content-Type"] = "application/json"
    request = urllib.request.Request(url, data=body, headers=headers, method="POST" if body else "GET")
    with urllib.request.urlopen(request, timeout=timeout) as response:
        return json.loads(response.read().decode("utf-8"))


def call_scan(target: str) -> tuple[str, float]:
    started = time.perf_counter()
    event = request_json(f"{BASE_URL}/gradio_api/call/run_scan", {"data": [target]})
    event_id = event["event_id"]
    request = urllib.request.Request(
        f"{BASE_URL}/gradio_api/call/run_scan/{event_id}",
        headers={"User-Agent": "ModelSentry-Live-Benchmark/1.0"},
    )
    with urllib.request.urlopen(request, timeout=180) as response:
        payload = response.read().decode("utf-8")
    return payload, time.perf_counter() - started


def parse_sse(payload: str) -> tuple[str, Any]:
    event = ""
    for line in payload.splitlines():
        if line.startswith("event: "):
            event = line.removeprefix("event: ").strip()
        elif line.startswith("data: "):
            return event, json.loads(line.removeprefix("data: "))
    raise ValueError("No SSE data payload was returned")


def scan_target(entry: dict[str, str]) -> dict[str, Any]:
    target = entry["target"]
    record: dict[str, Any] = {"category": entry["category"], "target": target}
    try:
        raw, elapsed = call_scan(target)
        event, data = parse_sse(raw)
        record.update({"elapsed_seconds": round(elapsed, 3), "event": event})
        if event != "complete" or not isinstance(data, list):
            record.update({"success": False, "error": data})
            return record
        report_text = next(
            item for item in data
            if isinstance(item, str) and '"schema": "modelsentry.scan.v4"' in item
        )
        report = json.loads(report_text)
        sbom_file = next(
            item for item in data
            if isinstance(item, dict) and str(item.get("orig_name", "")).endswith(".cdx.json")
        )
        sbom = request_json(sbom_file["url"])
        record.update({
            "success": True,
            "revision": report.get("revision"),
            "repo_type": report.get("repo_type"),
            "summary": report.get("summary"),
            "finding_groups": len(report.get("findings") or []),
            "finding_occurrences": report.get("finding_occurrence_count", 0),
            "finding_details": [
                {
                    "rule_id": item.get("rule_id"),
                    "severity": item.get("severity"),
                    "status": item.get("status"),
                    "title": item.get("title"),
                    "occurrences": len(item.get("occurrences") or []),
                }
                for item in report.get("findings") or []
            ],
            "packages": len(report.get("package_inventory") or []),
            "vulnerabilities": len(report.get("vulnerabilities") or []),
            "vulnerability_details": [
                {
                    "id": item.get("id"),
                    "severity": item.get("severity"),
                    "package": item.get("package"),
                    "version": item.get("version"),
                    "fixed_versions": item.get("fixed_versions") or [],
                }
                for item in report.get("vulnerabilities") or []
            ],
            "osv": report.get("osv"),
            "coverage": (report.get("artifacts") or {}).get("static_coverage"),
            "inspected_files": len(report.get("inspected_files") or []),
            "skipped_files": report.get("skipped_files", 0),
            "sbom_valid": (
                sbom.get("bomFormat") == "CycloneDX"
                and sbom.get("specVersion") == "1.6"
                and isinstance(sbom.get("components"), list)
            ),
            "sbom_components": len(sbom.get("components") or []),
        })
    except Exception as exc:  # benchmark must preserve each failure and continue
        record.update({"success": False, "error": f"{type(exc).__name__}: {exc}"})
    return record


def check_negative_case(entry: dict[str, str]) -> dict[str, Any]:
    try:
        raw, elapsed = call_scan(entry["target"])
        event, data = parse_sse(raw)
        return {
            **entry,
            "elapsed_seconds": round(elapsed, 3),
            "event": event,
            "rejected": event == "error",
            "response": data,
        }
    except Exception as exc:
        return {**entry, "rejected": False, "error": f"{type(exc).__name__}: {exc}"}


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", required=True, type=Path)
    args = parser.parse_args()

    config = request_json(f"{BASE_URL}/config")
    run_scan = next(item for item in config["dependencies"] if item.get("api_name") == "run_scan")
    ui_contract = {
        "title": config.get("title"),
        "gradio_version": config.get("version"),
        "run_scan_outputs": len(run_scan.get("outputs") or []),
        "dropdowns": sum(item.get("type") == "dropdown" for item in config.get("components") or []),
        "dataframes": sum(item.get("type") == "dataframe" for item in config.get("components") or []),
        "explicit_scrollbar_css": ".gradio-container::-webkit-scrollbar" in (config.get("css") or ""),
    }

    results = []
    for index, entry in enumerate(TARGETS, start=1):
        record = scan_target(entry)
        results.append(record)
        print(
            f"[{index:02d}/{len(TARGETS)}] {entry['target']} "
            f"success={record['success']} elapsed={record.get('elapsed_seconds', 'n/a')}",
            flush=True,
        )

    negative_results = [check_negative_case(entry) for entry in NEGATIVE_CASES]
    successful_times = [item["elapsed_seconds"] for item in results if item.get("success")]
    document = {
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "space": BASE_URL,
        "ui_contract": ui_contract,
        "targets": results,
        "negative_cases": negative_results,
        "aggregate": {
            "total": len(results),
            "successful": sum(bool(item.get("success")) for item in results),
            "failed": sum(not item.get("success") for item in results),
            "median_seconds": round(statistics.median(successful_times), 3) if successful_times else None,
            "mean_seconds": round(statistics.mean(successful_times), 3) if successful_times else None,
            "max_seconds": max(successful_times, default=None),
            "sbom_valid": sum(bool(item.get("sbom_valid")) for item in results),
            "packages": sum(int(item.get("packages", 0)) for item in results),
            "vulnerabilities": sum(int(item.get("vulnerabilities", 0)) for item in results),
            "finding_groups": sum(int(item.get("finding_groups", 0)) for item in results),
        },
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(document, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(document["aggregate"], sort_keys=True), flush=True)


if __name__ == "__main__":
    main()