"""Measure standalone packed loading and a short System One request in a fresh process.

Use an empty HF_HOME and HF_HUB_OFFLINE=1. The base loaders are explicitly
blocked. RSS alone misses Metal buffers, so also sample physical footprint.
"""
import argparse
import json
import platform
import sys
import time
from pathlib import Path

import mlx.core as mx

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from kev.api import SystemOneRequest, to_record
from kev.checkpoint import Checkpoint, LoadOptions
from kev.model import SERVE_MAX_BRANCH, SERVE_MAX_STATE
from scripts.mlx_long_states import RSSPeak, footprint


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--run", required=True)
    parser.add_argument("--out", required=True)
    parser.add_argument("--mlx-limit-gb", type=float, default=3, help="MLX memory guideline, not a hard RAM cap")
    parser.add_argument("--filler-repeats", type=int, default=0, help="append neutral filler to measure longer states; this is not an accuracy benchmark")
    args = parser.parse_args()
    import kev.mlx_model
    import huggingface_hub

    def forbidden(*unused, **kwargs):
        raise AssertionError("standalone profiling must never call a base loader or download")

    kev.mlx_model.load_base = kev.mlx_model.load_full = kev.mlx_model.merge_lora = forbidden
    huggingface_hub.snapshot_download = forbidden
    mx.set_memory_limit(int(args.mlx_limit_gb * 1e9))
    mx.reset_peak_memory()
    start = time.perf_counter()
    with RSSPeak() as load_peak:
        ck = Checkpoint(args.run)
        tok, model = ck.load("mps", LoadOptions())
    result = {"run": str(Path(args.run).resolve()), "backend": model.backend, "dtype": model.dtype,
              "platform": platform.platform(), "original_loaders_blocked": True,
              "mlx_limit_bytes": int(args.mlx_limit_gb * 1e9), "mlx_limit_is_hard_cap": False, "load_seconds": time.perf_counter() - start,
              "mlx_load_peak_bytes": mx.get_peak_memory(), "mlx_after_load_bytes": mx.get_active_memory(),
              "load_footprint_peak_bytes": max(load_peak.footprint, footprint()), "load_rss_peak_bytes": load_peak.peak}
    request = SystemOneRequest(state="I was charged twice for order 8812. Please refund the duplicate charge." + " This is a routine account note." * args.filler_repeats,
                               questions={"team": {"type": "choice", "instructions": "Which team should handle this ticket?",
                                                   "criteria": {"billing": "Charges and payments", "shipping": "Delivery and delays"}}})
    record, _ = to_record(request)
    enc = model.encode(tok, record, max_state=SERVE_MAX_STATE, max_branch=SERVE_MAX_BRANCH, strict=True)
    start = time.perf_counter()
    with RSSPeak() as infer_peak:
        probs = model.probs(enc)
    result.update(state_tokens=enc["seg"].count(0), input_tokens=len(enc["ids"]), inference_seconds=time.perf_counter() - start,
                  probabilities={"billing": float(probs[0][0]), "shipping": float(probs[0][1])},
                  mlx_inference_peak_bytes=mx.get_peak_memory(), inference_footprint_peak_bytes=max(infer_peak.footprint, footprint()),
                  inference_rss_peak_bytes=infer_peak.peak)
    output = Path(args.out)
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()