#!/usr/bin/env python3
"""
Transcribe an MP3 file and show the top 5 token candidates for each position.

Usage:
    python -m whisperlivekit.whisper.transcribe_top_tokens audio.mp3 --model base
"""

import argparse
import sys
from typing import List, Tuple

import numpy as np
import torch
import torch.nn.functional as F

from . import load_model
from .audio import SAMPLE_RATE, load_audio, log_mel_spectrogram, pad_or_trim
from .tokenizer import get_tokenizer


def transcribe_with_top_tokens(
    model_name: str = "base",
    audio_path: str = None,
    language: str = "en",
    top_k: int = 5,
    device: str = None,
):
    """
    Transcribe audio and show top-k token candidates at each decoding step.

    Args:
        model_name: Whisper model size (tiny, base, small, medium, large)
        audio_path: Path to audio file
        language: Language code (e.g., 'en', 'fr')
        top_k: Number of top tokens to show
        device: Device to use (cuda, cpu, or None for auto)
    """
    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"

    print(f"Loading model '{model_name}' on {device}...")
    model = load_model(model_name, device=device)
    model.eval()

    tokenizer = get_tokenizer(
        model.is_multilingual,
        num_languages=model.num_languages,
        language=language,
        task="transcribe",
    )

    print(f"Loading audio from '{audio_path}'...")
    audio = load_audio(audio_path)
    mel = log_mel_spectrogram(audio, model.dims.n_mels, padding=0)

    # Process in 30-second chunks
    content_frames = mel.shape[-1]
    n_frames = 3000  # 30 seconds

    all_tokens = []
    all_top_k_info = []

    seek = 0
    while seek < content_frames:
        segment = pad_or_trim(mel[:, seek:], n_frames)
        segment = segment.to(device)

        # Encode audio
        audio_features = model.encoder(segment.unsqueeze(0))

        # Build initial tokens (SOT sequence)
        tokens = list(tokenizer.sot_sequence)
        tokens_tensor = torch.tensor([tokens], device=device)

        # Initialize KV cache
        kv_cache = {}

        segment_tokens = []
        segment_top_k = []

        # Decode tokens one by one
        max_tokens = model.dims.n_text_ctx // 2
        for step in range(max_tokens):
            # Get logits from decoder
            with torch.no_grad():
                if step == 0:
                    # First forward pass - process all initial tokens
                    logits = model.decoder(tokens_tensor, audio_features, kv_cache=kv_cache)
                    logits = logits[:, -1, :]  # Take last position
                else:
                    # Subsequent passes - only process last token
                    last_token = tokens_tensor[:, -1:]
                    logits = model.decoder(last_token, audio_features, kv_cache=kv_cache)
                    logits = logits[:, -1, :]

            # Apply softmax to get probabilities
            probs = F.softmax(logits, dim=-1)

            # Get top-k tokens and their probabilities
            top_probs, top_indices = probs[0].topk(top_k)

            # Decode token texts
            top_tokens_info = []
            for prob, idx in zip(top_probs.cpu().tolist(), top_indices.cpu().tolist()):
                try:
                    # Handle special tokens
                    if idx >= tokenizer.timestamp_begin:
                        # Timestamp token
                        time_offset = (idx - tokenizer.timestamp_begin) * 0.02
                        token_text = f"<|{time_offset:.2f}|>"
                    elif idx == tokenizer.eot:
                        token_text = "<|endoftext|>"
                    else:
                        token_text = tokenizer.decode([idx])
                except:
                    token_text = f"<token_{idx}>"

                top_tokens_info.append((token_text, prob, idx))

            # Select the best token (greedy)
            next_token = top_indices[0].item()

            # Stop if we hit EOT
            if next_token == tokenizer.eot:
                segment_top_k.append(top_tokens_info)
                break

            # Store info for non-timestamp tokens (actual text)
            if next_token < tokenizer.timestamp_begin:
                segment_tokens.append(next_token)
                segment_top_k.append(top_tokens_info)

            # Append token and continue
            tokens_tensor = torch.cat([
                tokens_tensor,
                torch.tensor([[next_token]], device=device)
            ], dim=1)

        all_tokens.extend(segment_tokens)
        all_top_k_info.extend(segment_top_k)

        # Move to next segment
        seek += n_frames

    return all_tokens, all_top_k_info, tokenizer


def print_results(tokens: List[int], top_k_info: List[List[Tuple]], tokenizer):
    """Print the transcription results with top-k tokens."""
    print("\n" + "=" * 80)
    print("TRANSCRIPTION WITH TOP-5 TOKEN CANDIDATES")
    print("=" * 80)

    # Print full transcription first
    full_text = tokenizer.decode(tokens)
    print(f"\nFull transcription: {full_text}")
    print("\n" + "-" * 80)
    print("Token-by-token breakdown:")
    print("-" * 80)

    for i, (token_id, top_k) in enumerate(zip(tokens, top_k_info)):
        selected_text = tokenizer.decode([token_id])
        print(f"\n[Position {i+1}] Selected: '{selected_text}' (token_id={token_id})")
        print("  Top 5 candidates:")
        for rank, (text, prob, tid) in enumerate(top_k, 1):
            marker = " <-- SELECTED" if tid == token_id else ""
            print(f"    {rank}. '{text}' (prob={prob:.4f}, id={tid}){marker}")


def main():
    parser = argparse.ArgumentParser(
        description="Transcribe audio with top-5 token visibility"
    )
    parser.add_argument(
        "audio",
        type=str,
        help="Path to the audio file (MP3, WAV, etc.)",
    )
    parser.add_argument(
        "--model",
        "-m",
        type=str,
        default="base",
        choices=["tiny", "base", "small", "medium", "large", "large-v2", "large-v3"],
        help="Whisper model size (default: base)",
    )
    parser.add_argument(
        "--language",
        "-l",
        type=str,
        default="en",
        help="Language code (default: en)",
    )
    parser.add_argument(
        "--top-k",
        "-k",
        type=int,
        default=5,
        help="Number of top tokens to show (default: 5)",
    )
    parser.add_argument(
        "--device",
        "-d",
        type=str,
        default=None,
        help="Device to use (cuda, cpu, or auto)",
    )

    args = parser.parse_args()

    tokens, top_k_info, tokenizer = transcribe_with_top_tokens(
        model_name=args.model,
        audio_path=args.audio,
        language=args.language,
        top_k=args.top_k,
        device=args.device,
    )

    print_results(tokens, top_k_info, tokenizer)


if __name__ == "__main__":
    main()
