import argparse
import json
import math
import os
import re
import subprocess
import time
from pathlib import Path
from typing import Any


_MIB_PAIR = re.compile(r"(?P<used>\d+)\s*/\s*(?P<total>\d+)")
_DEVICE_ROW = re.compile(r"^\|\s*(?P<index>\d+)\s+\S+\s*\|")
_DEVICE_DETAIL_ROW = re.compile(
    r"^\|\s*\d+\s*\|\s*"
    r"(?:[0-9A-Fa-f]{4}:)?[0-9A-Fa-f]{2}:[0-9A-Fa-f]{2}\.\d\s*\|"
)


def parse_npu_smi_hbm(output: str, npu_index: int) -> dict[str, int] | None:
    """Extract HBM from a device detail row associated with its preceding NPU."""
    current_device: int | None = None
    for line in output.splitlines():
        device_match = _DEVICE_ROW.match(line)
        if device_match:
            current_device = int(device_match.group("index"))
            continue
        if "Process id" in line or "Process memory" in line:
            current_device = None
            continue
        if current_device != npu_index or not _DEVICE_DETAIL_ROW.match(line):
            continue
        for match in _MIB_PAIR.finditer(line):
            used = int(match.group("used"))
            total = int(match.group("total"))
            if total >= 1024 and used <= total:
                return {"used_mib": used, "total_mib": total}
    return None


def current_rss_mib() -> float | None:
    try:
        for line in Path("/proc/self/status").read_text().splitlines():
            if line.startswith("VmRSS:"):
                return int(line.split()[1]) / 1024.0
    except (OSError, ValueError, IndexError):
        pass
    return None


class ResourceSampler:
    def __init__(self, npu_index: int):
        self.npu_index = npu_index
        self.samples: list[dict[str, Any]] = []
        self.rss_samples: list[float] = []

    def _sample(self, phase: str) -> None:
        rss = current_rss_mib()
        if rss is not None:
            self.rss_samples.append(rss)
        sample: dict[str, Any] = {"timestamp": time.time(), "phase": phase}
        try:
            completed = subprocess.run(
                ["npu-smi", "info"],
                capture_output=True,
                text=True,
                timeout=5,
                check=False,
            )
            raw = completed.stdout or completed.stderr
            parsed = parse_npu_smi_hbm(raw, self.npu_index)
            if parsed is None:
                sample["used_mib"] = None
                sample["total_mib"] = None
                sample["raw"] = raw[:4096]
                sample["returncode"] = completed.returncode
            else:
                sample.update(parsed)
        except (OSError, subprocess.SubprocessError) as exc:
            sample["used_mib"] = None
            sample["total_mib"] = None
            sample["raw"] = str(exc)
        self.samples.append(sample)

    def checkpoint(self, phase: str) -> None:
        self._sample(phase)

    def hbm_summary(self) -> tuple[int | None, int | None, int | None]:
        if not self.samples:
            return None, None, None
        parsed = [sample["used_mib"] for sample in self.samples
                  if sample.get("used_mib") is not None]
        before = self.samples[0].get("used_mib")
        after = self.samples[-1].get("used_mib")
        peak = max(parsed) if parsed else None
        return before, peak, after


def percentile(values: list[float], quantile: float) -> float | None:
    if not values:
        return None
    ordered = sorted(values)
    if len(ordered) == 1:
        return ordered[0]
    position = (len(ordered) - 1) * quantile
    lower = math.floor(position)
    upper = math.ceil(position)
    if lower == upper:
        return ordered[lower]
    return ordered[lower] + (ordered[upper] - ordered[lower]) * (position - lower)


def build_prompts(tokenizer: Any, prompt_count: int, input_tokens: int) -> list[str]:
    seed_text = (
        "Analyze CPU weight offloading for deterministic large language model "
        "inference on Ascend hardware and summarize the practical tradeoffs. "
    )
    prompts = []
    for index in range(prompt_count):
        token_ids = tokenizer.encode(
            f"Request {index}: {seed_text}", add_special_tokens=False
        )
        if not token_ids:
            raise RuntimeError("tokenizer returned no tokens for benchmark prompt")
        repeats = math.ceil(input_tokens / len(token_ids))
        prompts.append(tokenizer.decode((token_ids * repeats)[:input_tokens]))
    return prompts


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", required=True)
    parser.add_argument("--label", required=True)
    parser.add_argument("--offload-group-size", type=int, required=True)
    parser.add_argument("--offload-num-in-group", type=int, required=True)
    parser.add_argument("--prompt-count", type=int, default=8)
    parser.add_argument("--input-tokens", type=int, default=128)
    parser.add_argument("--max-tokens", type=int, default=64)
    parser.add_argument("--max-model-len", type=int, default=512)
    parser.add_argument("--warmup-rounds", type=int, default=1)
    parser.add_argument("--measured-rounds", type=int, default=3)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--result-file", type=Path, required=True)
    parser.add_argument("--npu-index", type=int, default=7)
    parser.add_argument("--gpu-memory-utilization", type=float, default=0.5)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    if min(args.prompt_count, args.input_tokens, args.max_tokens,
           args.measured_rounds) <= 0 or args.warmup_rounds < 0:
        raise ValueError("workload counts must be positive and warmup rounds non-negative")
    if not 0 < args.gpu_memory_utilization <= 1:
        raise ValueError("gpu memory utilization must be in (0, 1]")

    from vllm import LLM, SamplingParams

    rss_before = current_rss_mib()
    sampler = ResourceSampler(args.npu_index)
    sampler.checkpoint("before_load")
    completed = False
    try:
        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=args.prompt_count,
            max_num_batched_tokens=args.max_model_len * args.prompt_count,
            tensor_parallel_size=1,
            enforce_eager=True,
            gpu_memory_utilization=args.gpu_memory_utilization,
            seed=args.seed,
            offload_group_size=args.offload_group_size,
            offload_num_in_group=args.offload_num_in_group,
        )
        load_seconds = time.perf_counter() - load_start
        prompts = build_prompts(llm.get_tokenizer(), args.prompt_count,
                                args.input_tokens)
        sampling = SamplingParams(
            temperature=0,
            max_tokens=args.max_tokens,
            seed=args.seed,
        )

        for _ in range(args.warmup_rounds):
            llm.generate(prompts, sampling)

        round_latencies: list[float] = []
        output_tokens = 0
        generation_start = time.perf_counter()
        for _ in range(args.measured_rounds):
            round_start = time.perf_counter()
            outputs = llm.generate(prompts, sampling)
            round_latencies.append(time.perf_counter() - round_start)
            output_tokens += sum(
                len(request.outputs[0].token_ids) for request in outputs
            )
        generate_seconds = time.perf_counter() - generation_start
        completed = True
    finally:
        if completed:
            sampler.checkpoint("after_generation")

    rss_after = current_rss_mib()
    rss_values = [value for value in [rss_before, *sampler.rss_samples, rss_after]
                  if value is not None]
    hbm_before, hbm_peak, hbm_after = sampler.hbm_summary()
    request_count = args.prompt_count * args.measured_rounds
    result = {
        "label": args.label,
        "model": args.model,
        "seed": args.seed,
        "offload_group_size": args.offload_group_size,
        "offload_num_in_group": args.offload_num_in_group,
        "prompt_count": args.prompt_count,
        "input_tokens": args.input_tokens,
        "max_tokens": args.max_tokens,
        "max_model_len": args.max_model_len,
        "gpu_memory_utilization": args.gpu_memory_utilization,
        "warmup_rounds": args.warmup_rounds,
        "measured_rounds": args.measured_rounds,
        "npu_index": args.npu_index,
        "load_seconds": load_seconds,
        "generate_seconds": generate_seconds,
        "request_count": request_count,
        "output_tokens": output_tokens,
        "output_tokens_per_second": output_tokens / generate_seconds,
        "requests_per_second": request_count / generate_seconds,
        "round_latencies_seconds": round_latencies,
        "latency_p50_seconds": percentile(round_latencies, 0.50),
        "latency_p95_seconds": percentile(round_latencies, 0.95),
        "process_rss_before_mib": rss_before,
        "process_rss_peak_mib": max(rss_values) if rss_values else None,
        "process_rss_after_mib": rss_after,
        "npu_hbm_before_mib": hbm_before,
        "npu_hbm_peak_mib": hbm_peak,
        "npu_hbm_after_mib": hbm_after,
        "npu_hbm_samples": sampler.samples,
        "pid": os.getpid(),
        "recorded_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
    }
    args.result_file.parent.mkdir(parents=True, exist_ok=True)
    with args.result_file.open("a", encoding="utf-8") as handle:
        handle.write(json.dumps(result, sort_keys=True) + "\n")
    print("BENCH_RESULT", json.dumps(result, sort_keys=True), flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
