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

import argparse
import asyncio
import json
import resource
import sys
import time
from pathlib import Path

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

from openmusic_analysis.application import build_service


async def main() -> None:
    parser = argparse.ArgumentParser(
        description="Measure model load and per-representation latency"
    )
    parser.add_argument("audio", type=Path)
    parser.add_argument("--lyrics", type=Path)
    parser.add_argument(
        "--representations",
        nargs="+",
        default=None,
        help="Load and measure only these representations (emotion models are opt-in).",
    )
    args = parser.parse_args()
    service = build_service()
    lyrics = args.lyrics.read_text(encoding="utf-8") if args.lyrics else None
    measurements = {}
    started = time.perf_counter()
    if args.representations is None:
        representations = ["audio.global", "audio.temporal"]
        if lyrics is not None:
            representations.append("lyrics.global")
    else:
        representations = list(dict.fromkeys(args.representations))
    await service.load_models(representations)
    measurements["model_startup_seconds"] = time.perf_counter() - started
    for representation in representations:
        started = time.perf_counter()
        await service.analyze(
            args.audio,
            lyrics=lyrics,
            requested_representations=[representation],
            track_id=None,
            content_identity=None,
        )
        measurements[f"{representation}_seconds"] = time.perf_counter() - started
    peak_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    # macOS reports bytes; Linux (including the Docker image) reports KiB.
    measurements["peak_process_rss_bytes"] = int(
        peak_rss if sys.platform == "darwin" else peak_rss * 1024
    )
    try:
        import torch

        if torch.cuda.is_available():
            measurements["peak_cuda_memory_bytes"] = torch.cuda.max_memory_allocated()
    except ImportError:
        pass
    print(json.dumps(measurements, indent=2))


if __name__ == "__main__":
    asyncio.run(main())