#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Any

import psutil

APP_DIR = Path(__file__).resolve().parent
if str(APP_DIR) not in sys.path:
    sys.path.insert(0, str(APP_DIR))

from asr_onnx_runtime import OnnxCacheAsrEngine  # noqa: E402


def nvidia_total_used_mb() -> tuple[float | None, list[float]]:
    try:
        out = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"],
            text=True,
            timeout=2,
        )
        values = [float(line.strip()) for line in out.splitlines() if line.strip()]
        return sum(values), values
    except Exception:
        return None, []


class RequestSampler:
    def __init__(self, process: psutil.Process, interval: float = 0.01) -> None:
        self.process = process
        self.interval = float(interval)
        self.peak_rss_bytes = process.memory_info().rss
        self._stop = threading.Event()
        self._thread = threading.Thread(target=self._run, daemon=True)

    def _run(self) -> None:
        while not self._stop.wait(self.interval):
            try:
                self.peak_rss_bytes = max(self.peak_rss_bytes, self.process.memory_info().rss)
            except psutil.Error:
                break

    def __enter__(self) -> "RequestSampler":
        self.peak_rss_bytes = self.process.memory_info().rss
        self._thread.start()
        return self

    def __exit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
        self._stop.set()
        self._thread.join(timeout=1.0)
        try:
            self.peak_rss_bytes = max(self.peak_rss_bytes, self.process.memory_info().rss)
        except psutil.Error:
            pass


def measure_one(
    *,
    bundle_dir: Path,
    audio_path: Path,
    cache_precision: str,
    audio_precision: str,
    max_new_tokens: int,
    provider: str,
    threads: int | None,
) -> dict[str, Any]:
    process = psutil.Process()
    gpu_before_total, gpu_before = nvidia_total_used_mb()
    started = time.perf_counter()
    engine = OnnxCacheAsrEngine(
        bundle_dir,
        provider=provider,
        intra_op_num_threads=threads,
        cache_precision=cache_precision,
        audio_precision=audio_precision,
    )
    static_rss = process.memory_info().rss
    static_gpu_total, static_gpu = nvidia_total_used_mb()
    load_seconds = time.perf_counter() - started

    with RequestSampler(process) as sampler:
        result = engine.transcribe(audio_path.read_bytes(), language="Chinese", max_new_tokens=max_new_tokens)
    post_rss = process.memory_info().rss
    post_gpu_total, post_gpu = nvidia_total_used_mb()

    return {
        "cache_precision": cache_precision,
        "audio_precision": audio_precision,
        "selected_cache_precision": engine.cache_precision,
        "selected_audio_precision": engine.audio_precision,
        "static_rss_bytes": int(static_rss),
        "infer_peak_rss_bytes": int(sampler.peak_rss_bytes),
        "post_rss_bytes": int(post_rss),
        "load_seconds": float(load_seconds),
        "elapsed_seconds": float(result.get("elapsed_seconds") or 0.0),
        "text": result.get("text"),
        "generated_tokens": result.get("generated_tokens"),
        "hit_stop": result.get("hit_stop"),
        "gpu_total_used_mb_before": gpu_before_total,
        "gpu_total_used_mb_static": static_gpu_total,
        "gpu_total_used_mb_post": post_gpu_total,
        "gpu_used_mb_before": gpu_before,
        "gpu_used_mb_static": static_gpu,
        "gpu_used_mb_post": post_gpu,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="Measure Audio8 ASR ONNX precision memory in a fresh process.")
    parser.add_argument("--bundle_dir", type=Path, default=APP_DIR / "model_bundle")
    parser.add_argument(
        "--audio",
        type=Path,
        required=True,
        help="Audio file path for the measurement.",
    )
    parser.add_argument("--cache_precision", choices=["fp32", "int8", "int4"], required=True)
    parser.add_argument("--audio_precision", choices=["fp32", "int8"], required=True)
    parser.add_argument("--max_new_tokens", type=int, default=64)
    parser.add_argument("--provider", default="CPUExecutionProvider")
    parser.add_argument("--threads", type=int, default=0)
    args = parser.parse_args()
    print(
        json.dumps(
            measure_one(
                bundle_dir=args.bundle_dir,
                audio_path=args.audio,
                cache_precision=args.cache_precision,
                audio_precision=args.audio_precision,
                max_new_tokens=args.max_new_tokens,
                provider=args.provider,
                threads=args.threads if args.threads > 0 else None,
            ),
            ensure_ascii=False,
        )
    )


if __name__ == "__main__":
    main()