import argparse
import asyncio
import logging
import time
import psutil
import os
import numpy as np
from openmusic_analysis.settings import Settings
from openmusic_analysis.worker.service import WorkerService

logging.basicConfig(level=logging.INFO)
log = logging.getLogger(__name__)

async def run_benchmark(audio_file: str, iterations: int):
    # Enable unpinned models so benchmark can run with main revision locally
    os.environ["OWLSONIC_ALLOW_UNPINNED_MODELS"] = "true"
    settings = Settings.from_env()
    service = WorkerService(settings)
    
    log.info("Starting benchmark...")
    
    t0 = time.perf_counter()
    await service.startup()
    startup_ms = (time.perf_counter() - t0) * 1000
    log.info(f"Startup complete in {startup_ms:.2f} ms")
    
    process = psutil.Process()
    rss_mb = process.memory_info().rss / (1024 * 1024)
    log.info(f"Memory RSS after startup: {rss_mb:.2f} MB")
    
    for i in range(iterations):
        log.info(f"--- Iteration {i+1}/{iterations} ---")
        try:
            t0 = time.perf_counter()
            embed_res = await service.audio_embeddings(audio_file)
            full_ms = (time.perf_counter() - t0) * 1000
            
            log.info(f"Embeddings API time: {full_ms:.2f} ms")
            log.info(f"  Probe time: {embed_res.timings.probe_ms:.2f} ms")
            log.info(f"  Decode time: {embed_res.timings.decode_ms:.2f} ms")
            log.info(f"  Inference time: {embed_res.timings.model_inference_ms:.2f} ms")
            
            # Pass to emotions
            t0 = time.perf_counter()
            
            temporal_dicts = None
            if embed_res.audio_temporal and embed_res.temporal_positions:
                temporal_dicts = []
                for t_emb, pos in zip(embed_res.audio_temporal, embed_res.temporal_positions):
                    temporal_dicts.append({
                        "start_ms": pos["start_ms"],
                        "end_ms": pos["end_ms"],
                        "embedding": t_emb
                    })
                    
            emo_res = await service.emotions_from_audio_embeddings(
                embedding_version=embed_res.metadata.audio_embedding_version,
                global_embedding=embed_res.audio_global,
                temporal_embeddings=temporal_dicts
            )
            emo_ms = (time.perf_counter() - t0) * 1000
            
            log.info(f"Emotions API time: {emo_ms:.2f} ms")
            log.info(f"  Inference time: {emo_res.timings.model_inference_ms:.2f} ms")
            
        except Exception as e:
            log.error(f"Benchmark iteration {i+1} failed: {e}")
            
        if (i + 1) % 10 == 0 or i == iterations - 1:
            rss_mb = process.memory_info().rss / (1024 * 1024)
            log.info(f"Memory RSS at iteration {i+1}: {rss_mb:.2f} MB")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("audio", help="Path to audio file")
    parser.add_argument("--iters", type=int, default=1, help="Number of iterations")
    args = parser.parse_args()
    
    asyncio.run(run_benchmark(args.audio, args.iters))