from __future__ import annotations

import argparse
import gc
import json
from pathlib import Path

import mlx.core as mx
import numpy as np
import torch
from huggingface_hub import snapshot_download
from transformers import AutoModelForCausalLM, AutoProcessor, AutoTokenizer

from ark_asr_mlx import ArkASR
from ark_asr_mlx.conversion import DEFAULT_REVISION, DEFAULT_SOURCE


def _cosine_similarity(left: np.ndarray, right: np.ndarray) -> float:
    left_flat = left.astype(np.float64).reshape(-1)
    right_flat = right.astype(np.float64).reshape(-1)
    denominator = np.linalg.norm(left_flat) * np.linalg.norm(right_flat)
    return float(np.dot(left_flat, right_flat) / denominator)


def _source_outputs(source_path: Path, audio_path: Path) -> dict[str, object]:
    device = "mps" if torch.backends.mps.is_available() else "cpu"
    dtype = torch.bfloat16 if device == "mps" else torch.float32
    processor = AutoProcessor.from_pretrained(source_path, trust_remote_code=True)
    tokenizer = AutoTokenizer.from_pretrained(
        source_path,
        trust_remote_code=True,
        fix_mistral_regex=True,
    )
    model = AutoModelForCausalLM.from_pretrained(
        source_path,
        trust_remote_code=True,
        dtype=dtype,
        attn_implementation="sdpa",
    ).to(device)
    model.eval()

    conversation = [
        {
            "role": "user",
            "content": [
                {"type": "audio", "path": str(audio_path)},
                {"type": "text", "text": "Please transcribe this audio."},
            ],
        }
    ]
    inputs = processor.apply_chat_template(
        conversation,
        add_generation_prompt=True,
        return_tensors="pt",
        sampling_rate=16_000,
        audio_padding="longest",
        text_kwargs={"padding": "longest"},
        audio_max_length=30 * 16_000,
    )
    raw_input_features = inputs["audios"].float().numpy()
    inputs = inputs.to(device)
    inputs["audios"] = inputs["audios"].to(dtype=dtype)
    eos_ids = tokenizer.eos_token_id
    keep_ids = {eos_ids} if isinstance(eos_ids, int) else set(eos_ids or [])
    bad_ids = set(tokenizer.all_special_ids) - keep_ids
    bad_ids.update(
        token_id
        for token, token_id in tokenizer.get_added_vocab().items()
        if token.startswith("<") and token.endswith(">") and token_id not in keep_ids
    )
    bad_words_ids = [[token_id] for token_id in sorted(bad_ids)]

    with torch.inference_mode():
        adapted = model.audio_encoder(inputs["audios"])
        outputs = model(**inputs, use_cache=False)
        generated = model.generate(
            **inputs,
            do_sample=False,
            max_new_tokens=256,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
            bad_words_ids=bad_words_ids,
        )

    generated_ids = generated[:, inputs.input_ids.shape[1] :]
    text = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
    result = {
        "input_ids": inputs.input_ids.detach().cpu().numpy(),
        "input_features": raw_input_features,
        "adapted": adapted.float().detach().cpu().numpy(),
        "last_logits": outputs.logits[:, -1, :].float().detach().cpu().numpy(),
        "token_ids": generated_ids.detach().cpu().numpy()[0].tolist(),
        "text": text,
    }
    del generated, outputs, adapted, inputs, model, processor, tokenizer
    gc.collect()
    if device == "mps":
        torch.mps.empty_cache()
    return result


def validate(model_path: Path, audio_path: Path) -> dict[str, object]:
    source_path = Path(
        snapshot_download(repo_id=DEFAULT_SOURCE, revision=DEFAULT_REVISION)
    )
    source = _source_outputs(source_path, audio_path)

    asr = ArkASR.from_pretrained(model_path)
    processed = asr.processor.process(audio_path)
    mlx_ids = mx.array(processed.input_ids)
    mlx_features = mx.array(processed.input_features).astype(
        asr.model.audio_encoder.whisper.conv1.weight.dtype
    )
    adapted = asr.model.audio_encoder(mlx_features)
    embeddings = asr.model.prepare_prompt_embeddings(
        mlx_ids,
        mlx_features,
        processed.audio_start,
        processed.audio_count,
    )
    logits = asr.model(mlx_ids, input_embeddings=embeddings)
    mx.eval(adapted, logits)
    mlx_result = asr.transcribe(audio_path)

    source_tokens = list(source["token_ids"])
    if source_tokens and source_tokens[-1] == asr.model.config.eos_token_id:
        source_tokens.pop()
    metrics = {
        "input_ids_equal": np.array_equal(processed.input_ids, source["input_ids"]),
        "input_features_max_abs": float(
            np.max(
                np.abs(
                    np.array(mlx_features.astype(mx.float32))
                    - source["input_features"]
                )
            )
        ),
        "adapter_cosine": _cosine_similarity(
            np.array(adapted.astype(mx.float32)),
            source["adapted"],
        ),
        "last_logits_cosine": _cosine_similarity(
            np.array(logits[:, -1, :].astype(mx.float32)),
            source["last_logits"],
        ),
        "token_ids_equal": list(mlx_result.token_ids) == source_tokens,
        "text_equal": mlx_result.text == source["text"],
        "mlx_text": mlx_result.text,
        "pytorch_text": source["text"],
    }
    if not metrics["input_ids_equal"]:
        raise AssertionError(f"Input IDs differ: {metrics}")
    if metrics["input_features_max_abs"] > 1e-6:
        raise AssertionError(f"Input features differ: {metrics}")
    if metrics["adapter_cosine"] < 0.999:
        raise AssertionError(f"Adapter parity failed: {metrics}")
    if metrics["last_logits_cosine"] < 0.999:
        raise AssertionError(f"Logit parity failed: {metrics}")
    if not metrics["token_ids_equal"] or not metrics["text_equal"]:
        raise AssertionError(f"Generation parity failed: {metrics}")
    return metrics


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(validate(args.model, args.audio), indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()