import time
import os
import glob
import sys
import numpy as np
import mlx.core as mx
import librosa
from transformers import AutoTokenizer

# Import the model and setup helpers from generate
sys.path.insert(0, "/Users/santosh/Desktop/ASR-Playground/mlx_canary")
from generate import load_llm, load_encoder, get_embed_tokens, extract_mel
from mlx_lm.models.cache import make_prompt_cache

def run_benchmark(audio_path: str, qwen_base_dir: str):
    print(f"Loading audio file: {audio_path}")
    
    # 1. Measure audio duration
    y, sr = librosa.load(audio_path, sr=None)
    duration = len(y) / sr
    print(f"Audio duration: {duration:.2f} seconds")
    
    # Load models
    print("\nLoading models into memory...")
    t_start_load = time.perf_counter()
    canary = load_encoder()
    llm = load_llm(qwen_base_dir)
    tokenizer = AutoTokenizer.from_pretrained(qwen_base_dir)
    embed_tokens = get_embed_tokens(llm)
    # Warmup
    print("\nRunning warm-up pass...")
    warmup_features = extract_mel(audio_path)
    warmup_embeds, _ = canary.encode_audio(warmup_features)
    mx.eval(warmup_embeds)
    
    # Warmup LLM prefill compilation
    warmup_t1 = [151644, 872, 198, 3167, 3114, 279, 2701, 25, 220]
    warmup_t2 = [151645, 198, 151644, 77091, 198]
    warmup_emb1 = embed_tokens(mx.array([warmup_t1]))
    warmup_emb2 = embed_tokens(mx.array([warmup_t2]))
    warmup_embeds_bf16 = warmup_embeds.astype(embed_tokens.weight.dtype)
    warmup_full_embeds = mx.concatenate([warmup_emb1, warmup_embeds_bf16, warmup_emb2], axis=1)
    mx.eval(warmup_full_embeds)
    
    warmup_cache = make_prompt_cache(llm)
    warmup_dummy = mx.zeros((1, warmup_full_embeds.shape[1]), dtype=mx.int32)
    warmup_logits = llm(warmup_dummy, cache=warmup_cache, input_embeddings=warmup_full_embeds)
    warmup_next = mx.argmax(warmup_logits[:, -1, :], axis=-1)
    mx.eval(warmup_next)
    
    # Warmup LLM decode loop compilation (first few steps)
    warmup_token_id = warmup_next.item()
    for _ in range(5):
        warmup_logits = llm(mx.array([[warmup_token_id]]), cache=warmup_cache)
        warmup_next = mx.argmax(warmup_logits[:, -1, :], axis=-1)
        mx.eval(warmup_next)
        warmup_token_id = warmup_next.item()
    
    # 2. Benchmark Audio Feature Extraction + Conformer Encoding
    print("\nBenchmarking Audio Encoder...")
    t_start_enc = time.perf_counter()
    audio_features = extract_mel(audio_path)
    audio_embeds, _ = canary.encode_audio(audio_features)
    # Evaluate MLX graph to block until computation completes
    mx.eval(audio_embeds)
    t_end_enc = time.perf_counter()
    enc_time = t_end_enc - t_start_enc
    print(f"Feature Extraction + Encoder execution time: {enc_time:.4f} seconds")
    
    # 3. Prompt Setup
    t1 = [151644, 872, 198, 3167, 3114, 279, 2701, 25, 220]
    t2 = [151645, 198, 151644, 77091, 198]
    emb1 = embed_tokens(mx.array([t1]))
    emb2 = embed_tokens(mx.array([t2]))
    audio_embeds_bf16 = audio_embeds.astype(embed_tokens.weight.dtype)
    full_embeds = mx.concatenate([emb1, audio_embeds_bf16, emb2], axis=1)
    mx.eval(full_embeds)
    
    # 4. Benchmark LLM Decoding (Autoregressive Generation)
    print("\nBenchmarking LLM Text Generation...")
    cache = make_prompt_cache(llm)
    dummy = mx.zeros((1, full_embeds.shape[1]), dtype=mx.int32)
    
    # Prefill/First token time
    t_start_prefill = time.perf_counter()
    logits = llm(dummy, cache=cache, input_embeddings=full_embeds)
    y = mx.argmax(logits[:, -1, :], axis=-1)
    mx.async_eval(y)
    mx.eval(y)
    token_id = y.item()
    prefill_time = time.perf_counter() - t_start_prefill
    
    def step(token):
        logits = llm(token, cache=cache)
        return mx.argmax(logits[:, -1, :], axis=-1)

    generated = []
    im_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
    
    # Decode loop benchmark
    t_start_decode = time.perf_counter()
    for _ in range(200):
        if token_id in (tokenizer.eos_token_id, im_end_id):
            break
        generated.append(token_id)
        
        # Enqueue next token computation on GPU
        next_y = step(mx.array([[token_id]]))
        mx.async_eval(next_y)
        
        # Fetch current token ID (blocks CPU while next_y runs in parallel on GPU)
        token_id = next_y.item()
    decode_time = time.perf_counter() - t_start_decode
    
    total_gen_time = prefill_time + decode_time
    num_tokens = len(generated)
    
    # Calculate separate metrics
    prompt_tokens = full_embeds.shape[1]
    prefill_tps = prompt_tokens / prefill_time
    
    # Generated tokens excluding the first token (which was generated by prefill)
    decode_tokens = max(0, num_tokens - 1)
    decode_tps = decode_tokens / decode_time if decode_time > 0 else 0.0
    
    rtf = (enc_time + total_gen_time) / duration
    result_text = tokenizer.decode(generated)
    
    print("\n" + "="*50)
    print("                 BENCHMARK RESULTS                ")
    print("="*50)
    print(f"Transcription: '{result_text}'")
    print(f"Audio Duration          : {duration:.2f}s")
    print(f"Prompt Size             : {prompt_tokens} tokens")
    print(f"Time to First Token     : {prefill_time:.4f}s")
    print(f"Prompt Processing Speed : {prefill_tps:.2f} tok/s")
    print(f"Decode Loop Time        : {decode_time:.4f}s")
    print(f"Tokens Decoded          : {decode_tokens}")
    print(f"Decode Generation Speed : {decode_tps:.2f} tok/s (excl. TTFT)")
    print(f"Total Gen Time          : {total_gen_time:.4f}s")
    print(f"Total Processing        : {enc_time + total_gen_time:.4f}s")
    print(f"Real-Time Factor        : {rtf:.4f}x ({(1/rtf):.1f}x faster than real-time)")
    print("="*50)

if __name__ == "__main__":
    audio = sys.argv[1] if len(sys.argv) > 1 else "/Users/santosh/Desktop/ASR-Playground/test.wav"
    
    candidates = glob.glob(os.path.expanduser(
        "~/.cache/huggingface/hub/models--Qwen--Qwen3-1.7B/snapshots/*/"))
    if not candidates:
        print("ERROR: Qwen3-1.7B not found.")
        sys.exit(1)
    qwen_dir = candidates[0].rstrip("/")
    
    run_benchmark(audio, qwen_dir)