"""Deployment verification harness for ControlMT v2.3.

Loads the model from HuggingFace (or local path), translates a fixed set of
test pairs in both directions, measures latency, and prints a structured
JSON report. The same script runs against CPU, GPU, int8-CPU, int8-bnb so
we can compare apples-to-apples.

Usage:
    python verify_deployment.py --device cpu
    python verify_deployment.py --device cuda
    python verify_deployment.py --device cpu --dtype bfloat16
    python verify_deployment.py --device cuda --dtype float16
    python verify_deployment.py --device cpu --quant int8-dynamic
    python verify_deployment.py --device cuda --quant bnb-int8
"""

import argparse, json, time, sys, platform
from importlib.metadata import version as _v


TEST_PAIRS = [
    ("kn2en", "ನಾನು ಕನ್ನಡ ಮಾತನಾಡುತ್ತೇನೆ.",                      "I speak Kannada."),
    ("kn2en", "ಬೆಂಗಳೂರಿನಲ್ಲಿ ಮೆಟ್ರೋ ಬಹಳ ಅನುಕೂಲಕರವಾಗಿದೆ.",        "metro convenience in Bangalore"),
    ("kn2en", "ಆಪಲ್ ಹೊಸ ಐಫೋನ್ 17 ಬಿಡುಗಡೆ ಮಾಡಿದೆ.",              "Apple released new iPhone 17"),
    ("en2kn", "I speak Kannada.",                              "ನಾನು ಕನ್ನಡ ಮಾತನಾಡುತ್ತೇನೆ."),
    ("en2kn", "The new metro line opens next month.",          "ಮುಂದಿನ ತಿಂಗಳು ಹೊಸ ಮೆಟ್ರೋ ಲೈನ್"),
    ("en2kn", "Please transfer money to my UPI ID.",           "UPI ಗೆ ಹಣ ವರ್ಗಾಯಿಸಿ"),
]


def env_versions():
    out = {"python": platform.python_version()}
    for pkg in ["torch", "transformers", "sentencepiece", "safetensors",
                "huggingface_hub", "accelerate", "bitsandbytes", "onnxruntime"]:
        try: out[pkg] = _v(pkg)
        except Exception: out[pkg] = None
    return out


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--model", default="anandkaman/controlmt-v2.3")
    p.add_argument("--device", choices=["cpu", "cuda"], default="cpu")
    p.add_argument("--dtype", choices=["fp32", "bfloat16", "float16"], default="fp32",
                   help="dtype to cast model to AFTER loading (no quantization)")
    p.add_argument("--quant", choices=["none", "int8-dynamic", "bnb-int8"], default="none",
                   help="quantization path (overrides --dtype where applicable)")
    p.add_argument("--num_beams", type=int, default=2)
    p.add_argument("--warmup_pairs", type=int, default=1)
    p.add_argument("--out", help="write JSON report to this path")
    args = p.parse_args()

    import torch
    from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

    print(f"=== Loading {args.model} ({args.device}, {args.dtype}, quant={args.quant}) ===")
    t0 = time.time()
    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)

    if args.quant == "bnb-int8":
        # GPU-only int8 via bitsandbytes — must be applied during load
        from transformers import BitsAndBytesConfig
        model = AutoModelForSeq2SeqLM.from_pretrained(
            args.model, trust_remote_code=True,
            quantization_config=BitsAndBytesConfig(load_in_8bit=True),
            device_map="auto",
        )
        # device map handled, skip manual .to()
        model_device = next(model.parameters()).device
    else:
        torch_dtype = {"fp32": torch.float32, "bfloat16": torch.bfloat16, "float16": torch.float16}[args.dtype]
        model = AutoModelForSeq2SeqLM.from_pretrained(
            args.model, trust_remote_code=True,
            torch_dtype=torch_dtype if args.dtype != "fp32" else None,
        )
        if args.quant == "int8-dynamic":
            # CPU dynamic quantization (Linear layers → qint8)
            model = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)
        model = model.to(torch.device(args.device)).eval()
        model_device = next(model.parameters()).device

    load_s = time.time() - t0
    print(f"   loaded in {load_s:.1f}s on {model_device}")

    # Warmup
    for direction, src, _ in TEST_PAIRS[:args.warmup_pairs]:
        try:
            _ = model.translate(src, tokenizer=tokenizer, direction=direction,
                                num_beams=args.num_beams, max_length=200)
        except Exception as e:
            print(f"   warmup failed: {type(e).__name__}: {e}")

    # Benchmark
    rows = []
    for direction, src, hint in TEST_PAIRS:
        t_start = time.time()
        try:
            out = model.translate(src, tokenizer=tokenizer, direction=direction,
                                  num_beams=args.num_beams, max_length=200,
                                  anti_lm_alpha=0.5)
            dt = time.time() - t_start
            rows.append({"direction": direction, "src": src, "hint": hint, "out": out,
                         "latency_s": round(dt, 3)})
            print(f"   [{direction} {dt:5.2f}s] {src[:35]:<35} → {out[:55]}")
        except Exception as e:
            rows.append({"direction": direction, "src": src, "error": f"{type(e).__name__}: {e}"})
            print(f"   [{direction}] ERROR: {type(e).__name__}: {e}")

    # Memory check
    mem_mb = None
    if model_device.type == "cuda":
        mem_mb = round(torch.cuda.memory_allocated(model_device) / 1024**2, 1)
        print(f"   GPU mem (allocated): {mem_mb} MB")

    report = {
        "model": args.model,
        "device": str(model_device),
        "dtype": args.dtype,
        "quant": args.quant,
        "num_beams": args.num_beams,
        "load_s": round(load_s, 2),
        "env": env_versions(),
        "gpu_mem_mb": mem_mb,
        "rows": rows,
        "median_latency_s": round(sorted([r["latency_s"] for r in rows if "latency_s" in r])[len(rows)//2], 3),
    }
    print()
    print(json.dumps({k: report[k] for k in ["device","dtype","quant","load_s","median_latency_s","gpu_mem_mb"]}, indent=2))
    if args.out:
        with open(args.out, "w") as f:
            json.dump(report, f, indent=2, ensure_ascii=False)
        print(f"\n→ wrote {args.out}")

if __name__ == "__main__":
    main()