#!/usr/bin/env python3
"""ASR Evaluation Script

Evaluates ASR model performance by calculating WER on audio samples.

Example call:
    CUDA_VISILBE_DEVICES=6 python scripts/asr_eval.py --n_songs 100 --sweep-cfg 1,1.5,2,3 --output sweep/eval.jsonl
"""

import argparse
import json
import os
import re
from pathlib import Path
from typing import Iterator, Literal, Optional
from collections import defaultdict
import tempfile
import io

import torch
import modal
from jiwer import wer
from openai import OpenAI
from google import genai
from google.genai import types

from suno_utils.audio import Audio
from suno_utils.diffusion.generation import encode_semantic
from suno_utils.gpt.bct.bct import Block, BlockType
from suno_utils.gpt.bct.bct_generation_simple import (
    BCTGenerationConfig,
    BlockSequence,
    generate_block,
)
from suno_utils.gpt.generation import GPT, load_model
from suno_utils.tasks.mert_25 import preload_models as preload_semantic_models_
from suno_utils.utils.s3 import download_s3_file_if_needed

import langcodes

# Constants
MAX_DURATION = 300  # 5 minutes
MIN_DURATION = 5  # 5 seconds

# Global model cache
_cached_model_container = None
_cached_cluster_model = None

TextBlockType = BlockType(
    name="text",
    is_causal=True,
)
NonCausalSemanticBlockType = BlockType(
    name="continuous_semantic",
    is_causal=False,
)


def get_lang_category(lang_str):
    lang_str = lang_str.lower()
    try:
        # Normalize to ISO 639-1 code (e.g., "english" -> "en", "japanese" -> "ja")
        lang_code = langcodes.find(lang_str).language
    except:
        # If can't normalize, use as-is
        lang_code = lang_str

    print(f"lang_code: {lang_code}")
    if lang_code in ["zh", "ja", "th", "lo", "my"]:  # Chinese, Japanese, Thai, Lao, Burmese
        return "space-less"  # space-less languages
    else:
        return "space-delimited"


def calculate_wer(target_text: str, generated_text: str, lang: str) -> float:
    """Calculate Word Error Rate (WER) between target and generated text.

    Normalizes text by removing brackets, punctuation, and lowercasing.

    Args:
        target_text: Reference/ground truth text
        generated_text: Hypothesis/generated text

    Returns:
        WER as a float (0.0 = perfect match, higher = more errors)
    """

    def normalize(text: str, lang: str) -> str:
        # Insert a space between every character for space-less languages
        if get_lang_category(lang) == "space-less":
            text = " ".join(text)
        # Remove everything inside square brackets
        text = re.sub(r"\[.*?\]", "", text)
        # Remove everything inside parentheses
        text = re.sub(r"\(.*?\)", "", text)
        # Lowercase
        text = text.lower()
        # Remove filler words: hmm, mm, mhm, mmm, uh, um (as whole words)
        text = re.sub(r"\b(hmm|mm|mhm|mmm|uh|um)\b", "", text)
        # Remove punctuation
        text = re.sub(r"[^\w\s]", "", text)
        # Collapse multiple spaces
        text = re.sub(r"\s+", " ", text)
        # Strip leading/trailing whitespace
        text = text.strip()
        return text

    target_clean = normalize(target_text, lang)
    generated_clean = normalize(generated_text, lang)

    return wer(target_clean, generated_clean)


def read_jsonl_filtered(
    filepath: Path | str,
    limit: int | None = None,
) -> Iterator[dict]:
    """Read a JSONL file line-by-line, filtering by duration.

    Args:
        filepath: Path to the JSONL file
        limit: Optional maximum number of entries to yield

    Yields:
        Parsed JSON objects that pass the filter
    """
    count = 0
    with open(filepath, "r") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue

            data = json.loads(line)

            if data.get("lang") is None:
                print(f"Skipping {data.get('id')} because lang is None")
                continue
            if data["duration_s"] > MAX_DURATION or data["duration_s"] < MIN_DURATION:
                continue

            yield data

            if limit is not None:
                count += 1
                if count >= limit:
                    break


def asr_custom(audio: Audio, gpt_model_path: Optional[str] = None, cfg_scale: float = 1.0) -> str:
    """Perform ASR on audio using custom model.

    Loads the model internally on first call.

    Args:
        audio: Audio object to transcribe
        gpt_model_path: Path to GPT model checkpoint
        cfg_scale: Classifier-free guidance scale

    Returns:
        Generated text transcription
    """
    global _cached_model_container, _cached_cluster_model

    # Load models only on first call
    if _cached_model_container is None:
        print("Loading models for the first time...")
        gpt_model_path = (
            gpt_model_path or "/app2/suno/checkpoints/2025-09-25_18-16-42/last_ckpt_infer.pt"
        )
        tokenizer_path = "s3://suno-data/georg/models/tokenizers/tokenizer_60k.json"

        _cached_model_container = load_model(
            ckpt_path=download_s3_file_if_needed(gpt_model_path),
            tokenizer_path=download_s3_file_if_needed(tokenizer_path),
        )

        _cached_cluster_model = preload_semantic_models_(
            checkpoint_filepath="s3://suno-data/georg/models/semantic/mert_25.pt",
            centroids_filepath="s3://suno-data/georg/models/semantic/mert_25_2x4k.npy",
        )["cluster_model"].cpu()
        print("Models loaded and cached.")

    model_container = _cached_model_container
    model: GPT = model_container["model"]
    cfg = model.config

    codes = torch.from_numpy(
        encode_semantic(audio, pad_to_chunksize=True, batch_size=48, do_clustering=False).T
    ).unsqueeze(0)

    sem_block = Block(
        NonCausalSemanticBlockType,
        inputs={
            "continuous_semantic_input": codes.to(torch.bfloat16),
        },
    )
    prompt_text = "[no tags]"
    text_codes = [cfg.text_infer_token] + model_container["tokenizer"].encode(prompt_text)

    text_block = Block(
        TextBlockType,
        inputs={
            "text_input": torch.tensor(text_codes).reshape(1, 1, -1),
        },
    )
    blocks = BlockSequence([sem_block, text_block])
    no_sem_blocks = BlockSequence([text_block])
    gconf = BCTGenerationConfig(
        [(cfg_scale, blocks), (1.0 - cfg_scale, no_sem_blocks)],  # classifier-free guidance
        max_autoregressive_steps=1000,
        eos_token=model.config.text_pad_token,
        top_k=1,
    )

    block = generate_block(model, gconf)
    text_codes = block.inputs["text_input"][0, 0, : len(block)]
    tok = model_container["tokenizer"]

    text_output = tok.decode(text_codes, clean_up_tokenization_spaces=True)

    return text_output


def asr_api_dummy(audio: Audio) -> str:
    """Dummy ASR function for API-based models.

    Args:
        audio: Audio object to transcribe

    Returns:
        Placeholder text
    """
    return "meow"


def asr_whisper(audio: Audio, model: Optional[str] = None) -> str:
    """ASR using OpenAI Whisper API.

    Args:
        audio: Audio object to transcribe

    Returns:
        Transcribed text
    """
    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

    wav_path = "tmp_whisper.wav"
    audio.write_wav(wav_path)

    try:
        with open(wav_path, "rb") as f:
            transcript = client.audio.transcriptions.create(
                file=f,
                model=model or "gpt-4o-mini-transcribe",
                response_format="text",
            )
    finally:
        os.remove(wav_path)

    return transcript


def asr_gemini(audio: Audio, model: Optional[str] = None) -> str:
    """ASR using Google Gemini API.

    Args:
        audio: Audio object to transcribe

    Returns:
        Transcribed text
    """
    # Load Google credentials
    creds_path = os.getenv(
        "GOOGLE_APPLICATION_CREDENTIALS", "/home/vibert/data/credentials/google-3.json"
    )
    os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = creds_path

    # Initialize Gemini client with Vertex AI
    client = genai.Client(vertexai=True, project="ml-gemini-455703", location="us-central1")

    wav_path = "tmp_gemini.wav"
    audio.write_wav(wav_path)

    try:
        with open(wav_path, "rb") as f:
            audio_bytes = f.read()

        response = client.models.generate_content(
            model=model or "gemini-2.5-flash",
            contents=[
                types.Part(inline_data=types.Blob(data=audio_bytes, mime_type="audio/wav")),
                types.Part(
                    text="Please transcribe all the lyrics or speech in this audio file. Write them in their original script rather than using any romanized forms such as pinyin or romanization."
                ),
            ],
        )
        return response.text
    finally:
        os.remove(wav_path)


def asr_qwen_omni(audio: Audio, model: Optional[str] = None, s3_id: Optional[str] = None) -> str:
    """ASR using Qwen-Omni Instruct via Modal API.

    Args:
        audio: Audio object to transcribe
        model: Model deployment name (default: "qwen-omni-captioner-dev")
        s3_id: S3 file ID if audio is already in S3 (studio/uploads/{s3_id})

    Returns:
        Transcribed text
    """
    # Get Modal stub for Qwen-Omni Instruct
    deployment = model or "qwen-omni-captioner-dev"
    QwenOmniInstruct = modal.Cls.from_name(deployment, "QwenOmniInstruct")
    qwen_stub = QwenOmniInstruct()

    # If no S3 ID provided, we need to upload the audio to S3
    if s3_id is None:
        import boto3
        import uuid

        # Generate a unique ID for S3 upload
        audio_id = str(uuid.uuid4())

        # Upload audio to S3
        wav_path = f"tmp_qwen_{audio_id}.mp3"
        audio.write_mp3(wav_path)

        try:
            s3_client = boto3.client("s3")
            s3_path = f"studio/uploads/{audio_id}.mp3"
            with open(wav_path, "rb") as f:
                s3_client.upload_fileobj(f, "suno-data-uploads", s3_path)

            s3_id = f"{audio_id}.mp3"
        finally:
            os.remove(wav_path)

    # Create QueueItem JSON with S3 ID
    qwen_args = json.dumps(
        {
            "id": s3_id,
            "metadata": {},
        }
    )

    # Use the same prompt as Gemini for consistency
    text_prompt = "Please transcribe all the lyrics or speech in this audio file. Write them in their original script rather than using any romanized forms such as pinyin or romanization."

    # Call Modal API
    result = qwen_stub.instruct_audio.remote(qwen_args, text_prompt=text_prompt)
    return result


def perform_asr(
    audio: Audio, model_type: str, model: Optional[str] = None, cfg_scale: float = 1.0
) -> str:
    """Dispatch ASR call based on model type.

    Args:
        audio: Audio object to transcribe
        model: Model to use
        model: Model name or path to a model checkpoint
        cfg_scale: Classifier-free guidance scale

    Returns:
        Generated text transcription
    """
    if model_type == "custom":
        return asr_custom(audio, gpt_model_path=model, cfg_scale=cfg_scale)
    elif model_type == "whisper":
        return asr_whisper(audio, model=model)
    elif model_type == "gemini":
        return asr_gemini(audio, model=model)
    elif model_type == "qwen-omni":
        return asr_qwen_omni(audio, model=model)
    else:
        return asr_api_dummy(audio)


def main():
    import time

    parser = argparse.ArgumentParser(description="Evaluate ASR model and calculate WER")
    parser.add_argument(
        "--model",
        type=str,
        default="custom",
        help="Model type to use for ASR (custom for local model, anything else for API)",
    )
    parser.add_argument("--n_songs", type=int, required=True, help="Number of songs to evaluate")
    parser.add_argument("--output", type=str, required=True, help="Output JSONL file path for results")
    parser.add_argument(
        "--input",
        type=str,
        default="/app2/suno/data/asr/metas_val.jsonl",
        help="Input JSONL file with audio metadata",
    )
    parser.add_argument(
        "--model-name",
        type=str,
        default=None,
        help="Either a model name or a path to a model checkpoint",
    )
    parser.add_argument(
        "--sweep-cfg",
        type=str,
        default=None,
        help="Comma-separated CFG values to sweep (e.g., '1,1.5,2,3'). If not provided, uses default CFG=1",
    )

    args = parser.parse_args()

    # Parse CFG sweep values
    if args.sweep_cfg:
        cfg_values = [float(x.strip()) for x in args.sweep_cfg.split(",")]
    else:
        cfg_values = [1.0]  # Default CFG

    print(f"Using model: {args.model}")
    print(f"CFG sweep values: {cfg_values}")

    print(f"Reading entries from: {args.input}")
    entries = list(read_jsonl_filtered(args.input, limit=args.n_songs))
    print(f"Processing {len(entries)} entries")

    # languages = defaultdict(int)
    # for entry in entries:
    #     languages[entry["lang"]] += 1
    # print(f"Languages: {languages}")

    # Prepare output directory and filename
    output_path = Path(args.output)
    output_dir = output_path.parent
    output_filename = output_path.name
    output_dir.mkdir(parents=True, exist_ok=True)

    # Store results for each CFG
    all_cfg_results = {}

    # Loop over CFG values
    for cfg_scale in cfg_values:
        print(f"\n{'=' * 60}")
        print(f"Running evaluation with CFG={cfg_scale}")
        print(f"{'=' * 60}\n")

        # Create output path for this CFG
        if args.sweep_cfg:
            cfg_output_path = output_dir / f"cfg-{cfg_scale}-{output_filename}"
        else:
            cfg_output_path = output_path

        results = []
        per_item_times = []
        start_time = time.time()
        for i, entry in enumerate(entries, 1):
            print(f"\nProcessing {i}/{len(entries)}: {entry.get('id', 'unknown')}")
            item_start = time.time()

            # Load audio and generate transcription
            audio = Audio.from_file(entry["local_filepath"])
            generated_text = perform_asr(audio, args.model, model=args.model_name, cfg_scale=cfg_scale)

            # Calculate WER
            reference_text = entry.get("text", "")
            lang = entry.get("lang")
            assert lang is not None
            wer_score = calculate_wer(reference_text, generated_text, lang)

            # Store results
            result = {
                "id": entry.get("id"),
                "local_filepath": entry.get("local_filepath"),
                "duration_s": entry.get("duration_s"),
                "lang": lang,
                "reference_text": reference_text,
                "generated_text": generated_text,
                "wer": wer_score,
            }
            results.append(result)

            item_end = time.time()
            elapsed = item_end - item_start
            per_item_times.append(elapsed)

            print(f"WER: {wer_score:.4f}")
            print(f"Reference: {reference_text[:100]}...")
            print(f"Generated: {generated_text[:100]}...")
            print(f"Time for this item: {elapsed:.2f} seconds")

            if i % 100 == 0:
                print(f"\nWriting intermediate results to: {cfg_output_path}")
                with open(cfg_output_path, "w", encoding="utf-8") as f:
                    for result in results:
                        json.dump(result, f, ensure_ascii=False)
                        f.write("\n")

        end_time = time.time()
        total_time = end_time - start_time
        avg_time = total_time / len(results) if results else 0.0

        # Write results to output file
        print(f"\nWriting results to: {cfg_output_path}")
        with open(cfg_output_path, "w", encoding="utf-8") as f:
            for result in results:
                json.dump(result, f, ensure_ascii=False)
                f.write("\n")

        # Calculate and print summary statistics
        wer_values = torch.tensor([r["wer"] for r in results])
        print("\n" + "=" * 50)
        if args.sweep_cfg:
            print(f"Summary Statistics for CFG={cfg_scale}:")
        else:
            print("Summary Statistics:")
        print(f"Mean WER: {wer_values.mean():.4f}")
        print(f"Median WER: {wer_values.median():.4f}")
        print(f"50th percentile: {wer_values.quantile(0.5):.4f}")
        print(f"90th percentile: {wer_values.quantile(0.9):.4f}")
        print(f"95th percentile: {wer_values.quantile(0.95):.4f}")
        print("=" * 50)
        print(f"Total time: {total_time:.2f} seconds")
        print(f"Average time per item: {avg_time:.2f} seconds")

        # Store statistics for this CFG
        all_cfg_results[cfg_scale] = {
            "mean": wer_values.mean().item(),
            "median": wer_values.median().item(),
            "p50": wer_values.quantile(0.5).item(),
            "p90": wer_values.quantile(0.9).item(),
            "p95": wer_values.quantile(0.95).item(),
        }

    # Print and save summary of all CFG sweeps
    if args.sweep_cfg:
        print("\n\n" + "=" * 60)
        print("SUMMARY OF ALL CFG SWEEPS")
        print("=" * 60)

        for cfg_scale in cfg_values:
            stats = all_cfg_results[cfg_scale]
            print(f"\nCFG={cfg_scale}:")
            print(f"  Mean WER:   {stats['mean']:.4f}")
            print(f"  Median WER: {stats['median']:.4f}")
            print(f"  P50:        {stats['p50']:.4f}")
            print(f"  P90:        {stats['p90']:.4f}")
            print(f"  P95:        {stats['p95']:.4f}")

        print("\n" + "=" * 60)

        # Save summary to file
        summary_path = output_dir / "results"
        print(f"\nSaving summary to: {summary_path}")
        with open(summary_path, "w", encoding="utf-8") as f:
            for cfg_scale in cfg_values:
                stats = all_cfg_results[cfg_scale]
                f.write(f"CFG={cfg_scale}:\n")
                f.write(f"  Mean WER:   {stats['mean']:.4f}\n")
                f.write(f"  Median WER: {stats['median']:.4f}\n")
                f.write(f"  P50:        {stats['p50']:.4f}\n")
                f.write(f"  P90:        {stats['p90']:.4f}\n")
                f.write(f"  P95:        {stats['p95']:.4f}\n")
                f.write("\n")


if __name__ == "__main__":
    main()