"""MOSS-Audio MLX runner for /srt eval — soundfile (no numba) + ASR transcription prompt + time markers.
Usage: python run_moss.py --model <dir> --audio clip.wav [--prompt "..."] [--max-tokens 4096] [--temp 0.0]
Works for both 4B-Thinking and 8B-hybrid bundles (audio int4 auto-detected).
"""
from __future__ import annotations
import argparse, sys, time
from pathlib import Path

def main():
    p = argparse.ArgumentParser()
    p.add_argument("--model", required=True, help="bundle dir (has mlx_llm/ mlx_audio/ scripts/)")
    p.add_argument("--audio", required=True)
    p.add_argument("--prompt", default="Please transcribe this audio.")
    p.add_argument("--max-tokens", type=int, default=4096)
    p.add_argument("--temp", type=float, default=0.0)
    p.add_argument("--repetition-penalty", type=float, default=1.02)
    p.add_argument("--no-time-marker", action="store_true")
    args = p.parse_args()

    HERE = Path(args.model).resolve()
    sys.path.insert(0, str(HERE / "scripts"))
    import soundfile as sf
    import mlx.core as mx
    import numpy as np
    from mlx_lm import load as mlx_load
    from mlx_lm.generate import generate_step
    from mlx_lm.sample_utils import make_sampler, make_logits_processors
    from moss_audio_mlx_bridge_v3 import load_mlx_audio_path, run_mlx_audio_pipeline, install_deepstack_hooks
    from moss_audio_mel_mlx import build_mel_and_input_ids

    ad_w = mx.load(str(HERE / "mlx_audio/audio_adapter.safetensors"))
    int4_audio = "down_proj.scales" in ad_w
    llm_hidden = (ad_w["down_proj.scales"].shape[0] if int4_audio else ad_w["down_proj.weight"].shape[0])
    size_tag = "4B" if llm_hidden == 2560 else "8B"
    print(f"[detect] {size_tag} bundle, audio_int4={int4_audio}, prompt={args.prompt!r}, time_marker={not args.no_time_marker}", flush=True)

    t0 = time.perf_counter()
    mlx_model, tok = mlx_load(str(HERE / "mlx_llm"))
    print(f"[load] LLM {time.perf_counter()-t0:.1f}s peak={mx.get_peak_memory()/1e9:.2f}GB", flush=True)
    t0 = time.perf_counter()
    encoder, adapter, mergers = load_mlx_audio_path(HERE / "mlx_audio", int4=int4_audio)
    print(f"[load] audio path {time.perf_counter()-t0:.1f}s", flush=True)

    y, sr = sf.read(args.audio)
    y = np.asarray(y, dtype=np.float32)
    if y.ndim > 1:
        y = y.mean(axis=1)
    assert sr == 16000, f"expected 16kHz, got {sr}"
    print(f"[audio] {args.audio} ({len(y)/16000:.1f}s)", flush=True)

    t0 = time.perf_counter()
    mel, lens, input_ids_mx, audio_token_id = build_mel_and_input_ids(
        y, tok, prompt=args.prompt, enable_time_marker=not args.no_time_marker)
    primary, ds_embeds = run_mlx_audio_pipeline(encoder, adapter, mergers, mel, lens)
    primary = primary.astype(mx.bfloat16)
    ds_embeds = [d.astype(mx.bfloat16) for d in ds_embeds]
    mx.eval(primary, *ds_embeds)
    print(f"[encode] {time.perf_counter()-t0:.2f}s primary={primary.shape}", flush=True)
    del encoder, adapter, mergers, mel, lens
    import gc; gc.collect(); mx.clear_cache()

    audio_mask = input_ids_mx == audio_token_id
    audio_positions = np.where(np.array(audio_mask[0]))[0]
    text_embeds = mlx_model.model.embed_tokens(input_ids_mx)
    text_np = np.array(text_embeds.astype(mx.float32))
    primary_np = np.array(primary.astype(mx.float32))
    text_np[0, audio_positions, :] = primary_np[0, :, :]
    merged = mx.array(text_np).astype(mx.bfloat16)
    ds_flat = [d[0] for d in ds_embeds]
    install_deepstack_hooks(mlx_model, ds_flat, audio_positions)

    sampler = make_sampler(temp=args.temp, top_p=1.0, top_k=(0 if args.temp == 0 else 50))
    gen_kwargs = dict(prompt=input_ids_mx[0], model=mlx_model,
                      input_embeddings=merged[0], max_tokens=args.max_tokens, sampler=sampler)
    if args.repetition_penalty:
        gen_kwargs["logits_processors"] = make_logits_processors(
            repetition_penalty=args.repetition_penalty, repetition_context_size=20)

    t0 = time.perf_counter(); generated = []; ttft = None
    for tok_id, _ in generate_step(**gen_kwargs):
        if ttft is None:
            ttft = time.perf_counter() - t0
        generated.append(int(tok_id))
        if tok_id == tok.eos_token_id:
            break
    elapsed = time.perf_counter() - t0
    n = len(generated)
    decode_s = max(elapsed - (ttft or 0), 1e-6)
    print(f"[gen] {n}tok in {elapsed:.1f}s ttft={ttft:.2f}s decode={max(n-1,0)/decode_s:.1f}t/s "
          f"rtf={(len(y)/16000)/elapsed:.2f}x peak={mx.get_peak_memory()/1e9:.2f}GB", flush=True)
    print(f"\n=== OUTPUT ===\n{tok.decode(generated)}\n=== END ===", flush=True)

if __name__ == "__main__":
    main()