#!/usr/bin/env python3
"""Validate a quantized Piko-9b build before publishing it.

Checks loading, degenerate output, text agreement against a reference model,
**image input**, memory, and throughput.  The vision path is checked explicitly
because most quantization tooling calibrates on text only, so a build can pass
every text check while its image handling has quietly broken.

    python scripts/validate_quantized_model.py \
        --quantized ./piko-9b-awq --reference Dexy2/Piko-9b \
        --output reports/quantization_validation.json

Never publish a quantized artefact that has not passed this.
"""

from __future__ import annotations

import argparse
import json
import platform
import time
from pathlib import Path
from typing import Any

import torch

REPO_ROOT = Path(__file__).resolve().parents[1]
ASSETS = REPO_ROOT / "evaluation" / "custom_suite" / "assets"

TEXT_PROBES = [
    ("capital", "What is the capital of France? Name only.", "paris"),
    ("arithmetic", "What is 17 * 23? Number only.", "391"),
    ("code", "Write a Python one-liner that reverses the string s. Code only.", "[::-1]"),
    ("format", "Reply with exactly the word OK and nothing else.", "ok"),
]

IMAGE_PROBES = [
    ("ocr_merchant", "receipt.png", "What is the merchant name? Name only.", "northgate"),
    ("ocr_total", "receipt.png", "What is the TOTAL on this receipt? Number only.", "27.30"),
    ("chart", "bar_chart.png", "Which quarter has the highest bar? Label only.", "q4"),
]


def visible(text: str) -> str:
    return text.rsplit("</think>", 1)[-1].strip() if "</think>" in text else text.strip()


def load(model_id: str, quantization: str | None) -> tuple[Any, Any, dict[str, Any]]:
    from transformers import AutoModelForMultimodalLM, AutoProcessor

    kwargs: dict[str, Any] = {"dtype": torch.bfloat16, "device_map": {"": 0}}
    if quantization == "4bit":
        from transformers import BitsAndBytesConfig

        kwargs["quantization_config"] = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
        )
    elif quantization == "8bit":
        from transformers import BitsAndBytesConfig

        kwargs["quantization_config"] = BitsAndBytesConfig(load_in_8bit=True)

    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    began = time.perf_counter()
    model = AutoModelForMultimodalLM.from_pretrained(model_id, **kwargs)
    model.eval()
    torch.cuda.synchronize()
    stats = {
        "load_seconds": round(time.perf_counter() - began, 1),
        "weights_vram_gb": round(torch.cuda.memory_allocated() / 1024**3, 3),
    }
    return model, AutoProcessor.from_pretrained(model_id), stats


def ask(
    model: Any, processor: Any, prompt: str, image: Path | None, tokens: int
) -> tuple[str, float]:
    content: list[dict[str, Any]] = []
    if image is not None:
        content.append({"type": "image", "url": str(image)})
    content.append({"type": "text", "text": prompt})
    inputs = processor.apply_chat_template(
        [{"role": "user", "content": content}],
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
    ).to(model.device)

    torch.cuda.synchronize()
    began = time.perf_counter()
    with torch.inference_mode():
        output = model.generate(**inputs, max_new_tokens=tokens, do_sample=False)
    torch.cuda.synchronize()
    elapsed = time.perf_counter() - began

    generated = int(output.shape[1] - inputs["input_ids"].shape[1])
    text = processor.decode(
        output[0][inputs["input_ids"].shape[1] :], skip_special_tokens=True
    ).strip()
    return text, round(generated / max(elapsed, 1e-9), 2)


def probe_model(model: Any, processor: Any, tokens: int) -> dict[str, Any]:
    results: dict[str, Any] = {"text": [], "image": [], "throughput_tok_s": []}

    for name, prompt, expected in TEXT_PROBES:
        try:
            text, rate = ask(model, processor, prompt, None, tokens)
            answer = visible(text)
            results["text"].append(
                {
                    "probe": name,
                    "passed": expected in answer.casefold(),
                    "degenerate": len(set(answer.replace(" ", ""))) < 5,
                    "response": answer[:400],
                }
            )
            results["throughput_tok_s"].append(rate)
        except Exception as exc:  # noqa: BLE001
            results["text"].append({"probe": name, "passed": False, "error": str(exc)})

    for name, asset, prompt, expected in IMAGE_PROBES:
        image = ASSETS / asset
        if not image.is_file():
            results["image"].append(
                {"probe": name, "passed": False, "error": "asset missing; run build_assets.py"}
            )
            continue
        try:
            text, _ = ask(model, processor, prompt, image, tokens)
            answer = visible(text)
            results["image"].append(
                {
                    "probe": name,
                    "passed": expected in answer.casefold(),
                    "response": answer[:400],
                }
            )
        except Exception as exc:  # noqa: BLE001
            results["image"].append({"probe": name, "passed": False, "error": str(exc)})

    return results


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--quantized", required=True)
    parser.add_argument("--quantized-mode", default=None, choices=[None, "4bit", "8bit"])
    parser.add_argument("--reference", default=None, help="Optional full-precision comparison")
    parser.add_argument("--reference-mode", default=None, choices=[None, "4bit", "8bit"])
    parser.add_argument("--max-new-tokens", type=int, default=128)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    if not torch.cuda.is_available():
        raise SystemExit("CUDA required.")

    import transformers

    report: dict[str, Any] = {
        "environment": {
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
            "python": platform.python_version(),
            "platform": platform.platform(),
            "torch": torch.__version__,
            "transformers": transformers.__version__,
            "gpu": torch.cuda.get_device_name(0),
        },
        "candidate": {"model": args.quantized, "mode": args.quantized_mode},
        "reference": None,
        "verdict": None,
    }

    print(f"Loading candidate {args.quantized}", flush=True)
    model, processor, stats = load(args.quantized, args.quantized_mode)
    report["candidate"].update(stats)
    report["candidate"]["probes"] = probe_model(model, processor, args.max_new_tokens)

    del model
    torch.cuda.empty_cache()

    if args.reference:
        print(f"Loading reference {args.reference}", flush=True)
        ref_model, ref_processor, ref_stats = load(args.reference, args.reference_mode)
        report["reference"] = {"model": args.reference, "mode": args.reference_mode, **ref_stats}
        report["reference"]["probes"] = probe_model(
            ref_processor and ref_model, ref_processor, args.max_new_tokens
        )
        del ref_model
        torch.cuda.empty_cache()

    probes = report["candidate"]["probes"]
    degenerate = any(p.get("degenerate") for p in probes["text"])
    text_passed = sum(p["passed"] for p in probes["text"])
    image_passed = sum(p["passed"] for p in probes["image"])

    if degenerate:
        verdict = "REJECT: degenerate output — model is broken or split across devices"
    elif text_passed < len(TEXT_PROBES) - 1:
        verdict = f"REJECT: text quality degraded ({text_passed}/{len(TEXT_PROBES)})"
    elif image_passed == 0:
        verdict = "REJECT: vision path broken — every image probe failed"
    elif image_passed < len(IMAGE_PROBES):
        verdict = (
            f"CAUTION: vision partially degraded ({image_passed}/{len(IMAGE_PROBES)}); "
            "document this before publishing"
        )
    else:
        verdict = (
            f"PASS: text {text_passed}/{len(TEXT_PROBES)}, image {image_passed}/{len(IMAGE_PROBES)}"
        )

    report["verdict"] = verdict
    report["summary"] = {
        "text_passed": text_passed,
        "text_total": len(TEXT_PROBES),
        "image_passed": image_passed,
        "image_total": len(IMAGE_PROBES),
        "degenerate": degenerate,
        "median_throughput_tok_s": (
            sorted(probes["throughput_tok_s"])[len(probes["throughput_tok_s"]) // 2]
            if probes["throughput_tok_s"]
            else None
        ),
    }

    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")

    print(f"\n{verdict}")
    print(f"wrote {args.output}")
    raise SystemExit(0 if verdict.startswith(("PASS", "CAUTION")) else 1)


if __name__ == "__main__":
    main()