import argparse
import gc
import json
import os
import statistics
import time

from vllm import LLM, SamplingParams


def parse_case(raw: str) -> tuple[str, int, int, str]:
    parts = raw.split(":")
    if len(parts) not in (3, 4):
        raise argparse.ArgumentTypeError(
            "case must be LABEL:OFFLOAD_GROUP_SIZE:OFFLOAD_NUM_IN_GROUP[:IMPL]"
        )
    label, group_size, num_in_group = parts[:3]
    impl = parts[3] if len(parts) == 4 else "prefetch"
    if impl not in {"prefetch", "sync", "noop"}:
        raise argparse.ArgumentTypeError("impl must be one of: prefetch, sync, noop")
    return label, int(group_size), int(num_in_group), impl


def configure_impl(impl: str) -> None:
    if impl == "sync":
        os.environ["ASCEND_OFFLOAD_IMPL"] = "sync"
    else:
        os.environ.pop("ASCEND_OFFLOAD_IMPL", None)


parser = argparse.ArgumentParser()
parser.add_argument("--model", required=True)
parser.add_argument("--case", action="append", type=parse_case, required=True)
parser.add_argument("--iterations", type=int, default=10)
parser.add_argument("--max-model-len", type=int, default=256)
parser.add_argument("--gpu-memory-utilization", type=float, default=0.6)
args = parser.parse_args()

prompts = [
    "Briefly define vLLM Ascend.",
    "Say one short sentence about CPU weight offload.",
    "Name one benefit of reducing NPU memory usage.",
]

for label, group_size, num_in_group, impl in args.case:
    configure_impl(impl)
    load_start = time.perf_counter()
    llm = LLM(
        model=args.model,
        trust_remote_code=True,
        dtype="float16",
        max_model_len=args.max_model_len,
        max_num_seqs=1,
        max_num_batched_tokens=args.max_model_len,
        tensor_parallel_size=1,
        enforce_eager=True,
        gpu_memory_utilization=args.gpu_memory_utilization,
        offload_group_size=group_size,
        offload_num_in_group=num_in_group,
    )
    load_seconds = time.perf_counter() - load_start

    sampling = SamplingParams(temperature=0.0, max_tokens=1)
    llm.generate(["warm up"], sampling, use_tqdm=False)

    latencies = []
    for i in range(args.iterations):
        prompt = prompts[i % len(prompts)]
        start = time.perf_counter()
        outputs = llm.generate([prompt], sampling, use_tqdm=False)
        elapsed = time.perf_counter() - start
        if not outputs or not outputs[0].outputs or not outputs[0].outputs[0].token_ids:
            raise RuntimeError(f"empty first-token output for case {label}, iter {i}")
        latencies.append(elapsed)

    result = {
        "label": label,
        "impl": impl,
        "model": args.model,
        "offload_group_size": group_size,
        "offload_num_in_group": num_in_group,
        "iterations": args.iterations,
        "load_seconds": round(load_seconds, 4),
        "ttft_proxy_avg_s": round(statistics.mean(latencies), 4),
        "ttft_proxy_min_s": round(min(latencies), 4),
        "ttft_proxy_max_s": round(max(latencies), 4),
        "ttft_proxy_p50_s": round(statistics.median(latencies), 4),
    }
    print("TTFT_RESULT", json.dumps(result, sort_keys=True), flush=True)

    del llm
    gc.collect()
