from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

import mlx.core as mx
import numpy as np

from ark_asr_mlx import ArkASR
from ark_asr_mlx.generation import _next_token


def benchmark(model_path: Path, audio_path: Path) -> dict[str, object]:
    mx.reset_peak_memory()
    started = time.perf_counter()
    asr = ArkASR.from_pretrained(model_path)
    load_seconds = time.perf_counter() - started

    process_started = time.perf_counter()
    inputs = asr.processor.process(audio_path)
    processing_seconds = time.perf_counter() - process_started

    input_ids = mx.array(inputs.input_ids)
    input_features = mx.array(inputs.input_features).astype(
        asr.model.audio_encoder.whisper.conv1.weight.dtype
    )
    generation_started = time.perf_counter()
    embeddings = asr.model.prepare_prompt_embeddings(
        input_ids,
        input_features,
        inputs.audio_start,
        inputs.audio_count,
    )
    cache = asr.model.make_cache()
    logits = asr.model(input_ids, cache=cache, input_embeddings=embeddings)
    mx.eval(logits)

    suppressed = asr.processor.suppressed_token_ids(asr.model.config.eos_token_id)
    mask = np.zeros(asr.model.config.vocab_size, dtype=np.bool_)
    mask[suppressed] = True
    suppression_mask = mx.array(mask)
    first_token = _next_token(logits, suppression_mask)
    first_token_seconds = time.perf_counter() - generation_started

    generated: list[int] = []
    token = first_token
    while len(generated) < 256 and token != asr.model.config.eos_token_id:
        generated.append(token)
        logits = asr.model(mx.array([[token]], dtype=mx.int32), cache=cache)
        mx.eval(logits)
        token = _next_token(logits, suppression_mask)
    generation_seconds = time.perf_counter() - generation_started
    text = asr.processor.tokenizer.decode(generated, skip_special_tokens=True).strip()
    return {
        "audio": str(audio_path),
        "load_seconds": load_seconds,
        "processing_seconds": processing_seconds,
        "first_token_seconds": first_token_seconds,
        "generation_seconds": generation_seconds,
        "generation_tokens": len(generated),
        "tokens_per_second": len(generated) / generation_seconds,
        "peak_memory_gb": mx.get_peak_memory() / 1_000_000_000,
        "text": text,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("audio", type=Path)
    parser.add_argument(
        "--model",
        type=Path,
        default=Path("models/ARK-ASR-0.6B-bf16"),
    )
    args = parser.parse_args()
    print(json.dumps(benchmark(args.model, args.audio), indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()