#!/usr/bin/env python3
"""Profile direct EasyCom Qwen2.5-Omni extraction on one physical GPU.

This is a diagnostic entry point for ``streaming_turn_model_directly.py``.
It accepts the same extraction arguments plus three profiler-only options:

``--gpu-id``
    Physical host GPU ID.  It defaults to 7 and is exposed to the extraction
    process as ``cuda:0``.
``--profile-log-dir``
    Directory for the console log, GPU telemetry CSV, and JSON summary.
``--gpu-sample-seconds``
    ``nvidia-smi`` sampling period in seconds.

The profiler synchronizes CUDA around individual model stages so their times
are attributable.  That synchronization perturbs throughput; use this script
to find bottlenecks, not to produce the final runtime measurement.
"""

from __future__ import annotations

import csv
import functools
import json
import os
import re
import resource
import statistics
import subprocess
import sys
import time
import types
from collections import defaultdict
from datetime import datetime
from pathlib import Path
from typing import Any, Callable, TextIO, TypeVar


def _pop_cli_option(name: str, default: str) -> str:
    """Remove one profiler-only option before the normal parser sees it."""

    for index, argument in enumerate(sys.argv[1:], start=1):
        if argument == name:
            if index + 1 >= len(sys.argv):
                raise SystemExit(f"{name} requires a value")
            value = sys.argv[index + 1]
            del sys.argv[index : index + 2]
            return value
        prefix = f"{name}="
        if argument.startswith(prefix):
            value = argument[len(prefix) :]
            del sys.argv[index]
            return value
    return default


# CUDA visibility must be fixed before importing torch or the Qwen modules.
PHYSICAL_GPU_ID = int(_pop_cli_option("--gpu-id", "7"))
PROFILE_LOG_DIR = Path(
    _pop_cli_option("--profile-log-dir", "logs/bottleneck_gpu7")
).expanduser()
GPU_SAMPLE_SECONDS = float(_pop_cli_option("--gpu-sample-seconds", "1"))
if PHYSICAL_GPU_ID < 0:
    raise SystemExit("--gpu-id must be >= 0")
if GPU_SAMPLE_SECONDS <= 0:
    raise SystemExit("--gpu-sample-seconds must be > 0")
os.environ["CUDA_VISIBLE_DEVICES"] = str(PHYSICAL_GPU_ID)

import torch  # noqa: E402  (must follow CUDA_VISIBLE_DEVICES)

import streaming_turn_model as core  # noqa: E402
import streaming_turn_model_directly as direct  # noqa: E402


T = TypeVar("T")


class TeeStream:
    """Write stdout/stderr to both the terminal and a durable log."""

    def __init__(self, terminal: TextIO, log_file: TextIO) -> None:
        self.terminal = terminal
        self.log_file = log_file
        self.encoding = getattr(terminal, "encoding", "utf-8")

    def write(self, data: str) -> int:
        terminal_count = self.terminal.write(data)
        self.log_file.write(data)
        self.flush()
        return terminal_count

    def flush(self) -> None:
        self.terminal.flush()
        self.log_file.flush()

    def isatty(self) -> bool:
        return self.terminal.isatty()


def _percentile(values: list[float], fraction: float) -> float:
    if not values:
        return 0.0
    ordered = sorted(values)
    index = round((len(ordered) - 1) * fraction)
    return float(ordered[index])


def _sync_cuda() -> None:
    if torch.cuda.is_available() and torch.cuda.is_initialized():
        torch.cuda.synchronize()


class StageProfiler:
    """Accumulate inclusive wall times for named extraction stages."""

    def __init__(self) -> None:
        self.durations: dict[str, list[float]] = defaultdict(list)
        self.metadata: dict[str, Any] = {}

    def measure(
        self,
        stage: str,
        function: Callable[..., T],
        *args: Any,
        cuda_sync: bool = False,
        **kwargs: Any,
    ) -> T:
        if cuda_sync:
            _sync_cuda()
        started = time.perf_counter()
        try:
            return function(*args, **kwargs)
        finally:
            if cuda_sync:
                _sync_cuda()
            self.durations[stage].append(time.perf_counter() - started)

    def rows(self) -> list[dict[str, Any]]:
        extract_time = sum(self.durations.get("extract_total", []))
        rows = []
        for stage, values in self.durations.items():
            total = sum(values)
            rows.append(
                {
                    "stage": stage,
                    "calls": len(values),
                    "total_seconds": total,
                    "average_ms": total / len(values) * 1000.0,
                    "p50_ms": statistics.median(values) * 1000.0,
                    "p95_ms": _percentile(values, 0.95) * 1000.0,
                    "maximum_ms": max(values) * 1000.0,
                    "percent_of_extract": (
                        total / extract_time * 100.0 if extract_time else None
                    ),
                }
            )
        return sorted(rows, key=lambda row: row["total_seconds"], reverse=True)


PROFILER = StageProfiler()


def _patch_function(
    owner: Any,
    name: str,
    stage: str,
    *,
    cuda_sync: bool = False,
) -> Callable[..., Any]:
    original = getattr(owner, name)

    @functools.wraps(original)
    def wrapped(*args: Any, **kwargs: Any) -> Any:
        return PROFILER.measure(
            stage,
            original,
            *args,
            cuda_sync=cuda_sync,
            **kwargs,
        )

    setattr(owner, name, wrapped)
    return original


def _patch_bound_method(
    owner: Any,
    name: str,
    stage: str,
    *,
    cuda_sync: bool,
) -> None:
    original = getattr(owner, name)

    @functools.wraps(original)
    def wrapped(_self: Any, *args: Any, **kwargs: Any) -> Any:
        return PROFILER.measure(
            stage,
            original,
            *args,
            cuda_sync=cuda_sync,
            **kwargs,
        )

    setattr(owner, name, types.MethodType(wrapped, owner))


def _install_profiling_hooks() -> None:
    # One-time input setup.
    _patch_function(direct, "_decode_audio_exact", "audio_decode")
    _patch_function(direct, "_timestamp_array", "video_timestamp_index")
    _patch_function(direct.CausalSourceAV, "__init__", "media_reader_init")

    # Per-state CPU media and processor work.
    _patch_function(direct.CausalSourceAV, "frames", "media_frames")
    _patch_function(direct.CausalSourceAV, "audio_slice", "media_audio_slice")
    _patch_function(core, "_prepare_conversation", "processor")

    # Per-state GPU/cache work.  Leaf GPU calls synchronize for attribution.
    _patch_function(core, "clone_cache", "cache_clone", cuda_sync=True)
    _patch_function(core, "_multimodal_embeddings", "multimodal_merge_total")
    _patch_function(core, "_append_embeddings", "append_total")
    _patch_function(core, "_read_open_state", "readout_total")
    _patch_function(
        core.OpenStreamReplayExtractor,
        "_prepare_pieces",
        "prepare_piece_total",
    )
    _patch_function(
        core.OpenStreamReplayExtractor,
        "_cache_from_intervals",
        "rolling_cache_rebuild",
    )
    _patch_function(
        core.CompletedBlockPieceCache,
        "put",
        "completed_block_cache_store",
    )
    _patch_function(
        core.CompletedBlockPieceCache,
        "get_many",
        "completed_block_cache_lookup",
    )
    _patch_function(torch, "save", "output_save")

    original_extract = core.OpenStreamReplayExtractor.extract

    @functools.wraps(original_extract)
    def extract_wrapped(self: Any, *args: Any, **kwargs: Any) -> Any:
        result = PROFILER.measure(
            "extract_total",
            original_extract,
            self,
            *args,
            cuda_sync=True,
            **kwargs,
        )
        PROFILER.metadata["num_states"] = int(len(result.times))
        return result

    core.OpenStreamReplayExtractor.extract = extract_wrapped

    original_load = core.load_qwen_omni

    @functools.wraps(original_load)
    def load_wrapped(*args: Any, **kwargs: Any) -> Any:
        model, processor = PROFILER.measure(
            "model_load",
            original_load,
            *args,
            cuda_sync=True,
            **kwargs,
        )

        # These are the mutually exclusive heavy GPU leaves inside one state.
        _patch_bound_method(
            model, "get_audio_features", "audio_tower", cuda_sync=True
        )
        _patch_bound_method(
            model, "get_video_features", "video_tower", cuda_sync=True
        )
        _patch_bound_method(model, "get_rope_index", "tmrope", cuda_sync=True)
        _patch_bound_method(
            model.model, "forward", "thinker_forward", cuda_sync=True
        )
        torch.cuda.reset_peak_memory_stats()
        return model, processor

    core.load_qwen_omni = load_wrapped


def _start_gpu_monitor(path: Path) -> tuple[subprocess.Popen[str] | None, TextIO | None]:
    command = [
        "nvidia-smi",
        "-i",
        str(PHYSICAL_GPU_ID),
        "--query-gpu=timestamp,index,name,utilization.gpu,utilization.memory,"
        "memory.used,memory.total,power.draw,temperature.gpu",
        "--format=csv,nounits",
        "-l",
        f"{GPU_SAMPLE_SECONDS:g}",
    ]
    try:
        output = path.open("w", encoding="utf-8", buffering=1)
        process = subprocess.Popen(
            command,
            stdout=output,
            stderr=subprocess.STDOUT,
            text=True,
        )
        return process, output
    except (FileNotFoundError, OSError) as exc:
        print(f"WARNING: could not start nvidia-smi monitor: {exc}", flush=True)
        return None, None


def _stop_gpu_monitor(
    process: subprocess.Popen[str] | None,
    output: TextIO | None,
) -> None:
    if process is not None and process.poll() is None:
        process.terminate()
        try:
            process.wait(timeout=5)
        except subprocess.TimeoutExpired:
            process.kill()
            process.wait(timeout=5)
    if output is not None:
        output.close()


def _numeric(value: str | None) -> float | None:
    if not value:
        return None
    match = re.search(r"[-+]?[0-9]*\.?[0-9]+", value)
    return float(match.group(0)) if match else None


def _gpu_telemetry_summary(path: Path) -> dict[str, Any]:
    if not path.is_file():
        return {"samples": 0}
    metrics: dict[str, list[float]] = defaultdict(list)
    with path.open("r", encoding="utf-8", errors="replace") as handle:
        reader = csv.DictReader(handle)
        for row in reader:
            for key, value in row.items():
                normalized = key.strip().lower()
                number = _numeric(value)
                if number is None:
                    continue
                if normalized.startswith("utilization.gpu"):
                    metrics["gpu_util_percent"].append(number)
                elif normalized.startswith("utilization.memory"):
                    metrics["memory_util_percent"].append(number)
                elif normalized.startswith("memory.used"):
                    metrics["memory_used_mib"].append(number)
                elif normalized.startswith("power.draw"):
                    metrics["power_watts"].append(number)
                elif normalized.startswith("temperature.gpu"):
                    metrics["temperature_c"].append(number)
    result: dict[str, Any] = {
        "samples": max((len(values) for values in metrics.values()), default=0)
    }
    for name, values in metrics.items():
        result[name] = {
            "mean": statistics.fmean(values),
            "p50": statistics.median(values),
            "p95": _percentile(values, 0.95),
            "max": max(values),
        }
    return result


def _print_summary(
    *,
    rows: list[dict[str, Any]],
    gpu_summary: dict[str, Any],
    overall_seconds: float,
) -> None:
    extract_seconds = sum(PROFILER.durations.get("extract_total", []))
    num_states = int(PROFILER.metadata.get("num_states", 0))
    print("\n=== Qwen Omni bottleneck profile (inclusive timings) ===")
    print(
        "Nested rows intentionally overlap; compare leaf rows such as "
        "processor/audio_tower/video_tower/thinker_forward/cache_clone."
    )
    print(
        f"overall={overall_seconds:.3f}s extract={extract_seconds:.3f}s "
        f"states={num_states} "
        f"throughput={(num_states / extract_seconds if extract_seconds else 0):.3f} state/s"
    )
    print(
        f"{'stage':<27} {'calls':>7} {'total_s':>10} {'avg_ms':>10} "
        f"{'p95_ms':>10} {'extract%':>10}"
    )
    for row in rows:
        percent = row["percent_of_extract"]
        percent_text = "-" if percent is None else f"{percent:.1f}"
        print(
            f"{row['stage']:<27} {row['calls']:>7d} "
            f"{row['total_seconds']:>10.3f} {row['average_ms']:>10.3f} "
            f"{row['p95_ms']:>10.3f} {percent_text:>10}"
        )

    gpu_util = gpu_summary.get("gpu_util_percent")
    memory_used = gpu_summary.get("memory_used_mib")
    if gpu_util:
        print(
            "GPU telemetry: "
            f"samples={gpu_summary['samples']} "
            f"util mean/p95/max={gpu_util['mean']:.1f}/"
            f"{gpu_util['p95']:.1f}/{gpu_util['max']:.1f}%"
        )
    if memory_used:
        print(
            "GPU memory MiB mean/p95/max="
            f"{memory_used['mean']:.0f}/{memory_used['p95']:.0f}/"
            f"{memory_used['max']:.0f}"
        )

    leaf_stages = (
        "media_frames",
        "processor",
        "audio_tower",
        "video_tower",
        "tmrope",
        "thinker_forward",
        "cache_clone",
    )
    totals = {
        stage: sum(PROFILER.durations.get(stage, [])) for stage in leaf_stages
    }
    bottleneck = max(totals, key=totals.get)
    advice = {
        "media_frames": "bulk-decode/cache the selected 2-FPS resized frames",
        "processor": "prefetch or microbatch the ten causal prefixes per 2-s block",
        "audio_tower": "microbatch the ten 0.2-s prefixes in each 2-s block",
        "video_tower": "cache unique 2-FPS visual inputs and batch equal shapes",
        "tmrope": "cache position IDs for the ten repeated prefix shapes",
        "thinker_forward": "microbatch independent streams/prefix branches",
        "cache_clone": "replace full KV clones with append-and-crop branch reuse",
    }
    print(
        f"Primary measured leaf bottleneck: {bottleneck} "
        f"({totals[bottleneck]:.3f}s). Recommended next optimization: "
        f"{advice[bottleneck]}."
    )


def main() -> None:
    PROFILE_LOG_DIR.mkdir(parents=True, exist_ok=True)
    stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    console_log = PROFILE_LOG_DIR / f"profile_gpu{PHYSICAL_GPU_ID}_{stamp}.log"
    gpu_log = PROFILE_LOG_DIR / f"gpu_gpu{PHYSICAL_GPU_ID}_{stamp}.csv"
    json_log = PROFILE_LOG_DIR / f"profile_gpu{PHYSICAL_GPU_ID}_{stamp}.json"

    original_stdout = sys.stdout
    original_stderr = sys.stderr
    status = "ok"
    error_text: str | None = None
    started = time.perf_counter()
    usage_started = resource.getrusage(resource.RUSAGE_SELF)

    with console_log.open("w", encoding="utf-8", buffering=1) as log_handle:
        sys.stdout = TeeStream(original_stdout, log_handle)
        sys.stderr = TeeStream(original_stderr, log_handle)
        monitor: subprocess.Popen[str] | None = None
        monitor_output: TextIO | None = None
        try:
            print(f"Physical GPU: {PHYSICAL_GPU_ID} -> process cuda:0")
            print(f"CUDA_VISIBLE_DEVICES={os.environ['CUDA_VISIBLE_DEVICES']}")
            print(f"Console log: {console_log.resolve()}")
            print(f"GPU log: {gpu_log.resolve()}")
            print(f"JSON summary: {json_log.resolve()}")
            print("Extraction argv:", " ".join(sys.argv[1:]))
            print(
                "NOTE: CUDA synchronization is enabled for attribution; "
                "this is not a production throughput benchmark."
            )
            if torch.cuda.is_available():
                print(f"Visible GPU: {torch.cuda.get_device_name(0)}")
                print(f"PyTorch: {torch.__version__}, CUDA: {torch.version.cuda}")
            else:
                raise RuntimeError("CUDA is not available after GPU selection")

            monitor, monitor_output = _start_gpu_monitor(gpu_log)
            _install_profiling_hooks()
            direct.main()
        except BaseException as exc:
            status = "error"
            error_text = f"{type(exc).__name__}: {exc}"
            raise
        finally:
            _sync_cuda()
            _stop_gpu_monitor(monitor, monitor_output)
            overall_seconds = time.perf_counter() - started
            usage_finished = resource.getrusage(resource.RUSAGE_SELF)
            gpu_summary = _gpu_telemetry_summary(gpu_log)
            rows = PROFILER.rows()
            PROFILER.metadata.update(
                {
                    "status": status,
                    "error": error_text,
                    "physical_gpu_id": PHYSICAL_GPU_ID,
                    "visible_cuda_device": 0,
                    "gpu_name": torch.cuda.get_device_name(0),
                    "torch_version": torch.__version__,
                    "torch_cuda_version": torch.version.cuda,
                    "overall_seconds": overall_seconds,
                    "cpu_user_seconds": (
                        usage_finished.ru_utime - usage_started.ru_utime
                    ),
                    "cpu_system_seconds": (
                        usage_finished.ru_stime - usage_started.ru_stime
                    ),
                    "max_rss_kib": usage_finished.ru_maxrss,
                    "peak_cuda_allocated_mib": (
                        torch.cuda.max_memory_allocated() / 1024**2
                    ),
                    "peak_cuda_reserved_mib": (
                        torch.cuda.max_memory_reserved() / 1024**2
                    ),
                    "console_log": str(console_log.resolve()),
                    "gpu_log": str(gpu_log.resolve()),
                }
            )
            _print_summary(
                rows=rows,
                gpu_summary=gpu_summary,
                overall_seconds=overall_seconds,
            )
            print(
                "Process resources: "
                f"user_cpu={PROFILER.metadata['cpu_user_seconds']:.2f}s "
                f"system_cpu={PROFILER.metadata['cpu_system_seconds']:.2f}s "
                f"max_rss={PROFILER.metadata['max_rss_kib'] / 1024:.1f}MiB "
                f"peak_cuda_allocated="
                f"{PROFILER.metadata['peak_cuda_allocated_mib']:.1f}MiB"
            )
            with json_log.open("w", encoding="utf-8") as handle:
                json.dump(
                    {
                        "metadata": PROFILER.metadata,
                        "gpu_telemetry": gpu_summary,
                        "stages": rows,
                        "argv": sys.argv[1:],
                    },
                    handle,
                    indent=2,
                )
            print(f"Saved profile summary to {json_log.resolve()}")
            sys.stdout = original_stdout
            sys.stderr = original_stderr


if __name__ == "__main__":
    main()