#!/usr/bin/env python
"""
VibeVoice ASR Batch Inference Demo Script

This script supports batch inference for ASR model and compares results
between batch processing and single-sample processing.
"""

import os
import sys
import torch
import numpy as np
from pathlib import Path
import argparse
import time
import json
import re
from typing import List, Dict, Any, Optional
from functools import wraps

from vibevoice.modular.modeling_vibevoice_asr import VibeVoiceASRForConditionalGeneration
from vibevoice.processor.vibevoice_asr_processor import VibeVoiceASRProcessor


class VibeVoiceASRBatchInference:
    """Batch inference wrapper for VibeVoice ASR model."""
    
    def __init__(
        self, 
        model_path: str, 
        device: str = "cuda", 
        dtype: torch.dtype = torch.bfloat16,
        attn_implementation: str = "sdpa"
    ):
        """
        Initialize the ASR batch inference pipeline.
        
        Args:
            model_path: Path to the pretrained model
            device: Device to run inference on (cuda, mps, xpu, cpu, auto)
            dtype: Data type for model weights
            attn_implementation: Attention implementation to use ('flash_attention_2', 'sdpa', 'eager')
        """
        print(f"Loading VibeVoice ASR model from {model_path}")
        
        # Load processor
        self.processor = VibeVoiceASRProcessor.from_pretrained(
            model_path,
            language_model_pretrained_name="Qwen/Qwen2.5-7B"
        )
        
        # Load model with specified attention implementation
        print(f"Using attention implementation: {attn_implementation}")
        self.model = VibeVoiceASRForConditionalGeneration.from_pretrained(
            model_path,
            dtype=dtype,
            device_map=device if device == "auto" else None,
            attn_implementation=attn_implementation,
            trust_remote_code=True
        )
        
        if device != "auto":
            self.model = self.model.to(device)
        
        self.device = device if device != "auto" else next(self.model.parameters()).device
        self.dtype = dtype
        self.model.eval()
        
        print(f"Model loaded successfully on {self.device}")
    
    def _prepare_generation_config(
        self,
        max_new_tokens: int = 512,
        temperature: float = 0.0,
        top_p: float = 0.9,
        do_sample: bool = True,
        num_beams: int = 1,
    ) -> dict:
        """Prepare generation configuration."""
        config = {
            "max_new_tokens": max_new_tokens,
            "pad_token_id": self.processor.pad_id,
            "eos_token_id": self.processor.tokenizer.eos_token_id,
        }
        
        # Beam search vs sampling
        if num_beams > 1:
            config["num_beams"] = num_beams
            config["do_sample"] = False  # Beam search doesn't use sampling
        else:
            config["do_sample"] = do_sample
            # Only set temperature and top_p when sampling is enabled
            if do_sample:
                config["temperature"] = temperature
                config["top_p"] = top_p
        
        return config
    
    def transcribe_batch(
        self,
        audio_inputs: List,
        max_new_tokens: int = 512,
        temperature: float = 0.0,
        top_p: float = 1.0,
        do_sample: bool = True,
        num_beams: int = 1,
    ) -> List[Dict[str, Any]]:
        """
        Transcribe multiple audio files/arrays in a single batch.
        
        Args:
            audio_inputs: List of audio file paths or (array, sampling_rate) tuples
            max_new_tokens: Maximum tokens to generate
            temperature: Temperature for sampling
            top_p: Top-p for nucleus sampling
            do_sample: Whether to use sampling
            
        Returns:
            List of transcription results
        """
        if len(audio_inputs) == 0:
            return []
        
        batch_size = len(audio_inputs)
        print(f"\nProcessing batch of {batch_size} audio(s)...")
        
        # Process all audio together
        inputs = self.processor(
            audio=audio_inputs,
            sampling_rate=None,
            return_tensors="pt",
            padding=True,
            add_generation_prompt=True
        )
        
        # Move to device
        inputs = {k: v.to(self.device) if isinstance(v, torch.Tensor) else v 
                  for k, v in inputs.items()}
        
        # Print batch info
        print(f"  Input IDs shape: {inputs['input_ids'].shape}")
        print(f"  Speech tensors shape: {inputs['speech_tensors'].shape}")
        print(f"  Attention mask shape: {inputs['attention_mask'].shape}")
        
        # Generate
        generation_config = self._prepare_generation_config(
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            do_sample=do_sample,
            num_beams=num_beams,
        )
        
        start_time = time.time()
        
        with torch.no_grad():
            output_ids = self.model.generate(
                **inputs,
                **generation_config
            )
        
        generation_time = time.time() - start_time
        
        # Decode outputs for each sample in the batch
        results = []
        input_length = inputs['input_ids'].shape[1]
        
        for i, audio_input in enumerate(audio_inputs):
            # Get generated tokens for this sample (excluding input tokens)
            generated_ids = output_ids[i, input_length:]
            
            # Remove padding tokens from the end
            # Find the first eos_token or pad_token
            eos_positions = (generated_ids == self.processor.tokenizer.eos_token_id).nonzero(as_tuple=True)[0]
            if len(eos_positions) > 0:
                generated_ids = generated_ids[:eos_positions[0] + 1]
            
            generated_text = self.processor.decode(generated_ids, skip_special_tokens=True)
            
            # Parse structured output
            try:
                transcription_segments = self.processor.post_process_transcription(generated_text)
            except Exception as e:
                print(f"Warning: Failed to parse structured output: {e}")
                transcription_segments = []
            
            # Get file name based on input type
            if isinstance(audio_input, str):
                file_name = audio_input
            elif isinstance(audio_input, dict) and 'id' in audio_input:
                file_name = audio_input['id']
            else:
                file_name = f"audio_{i}"
            
            results.append({
                "file": file_name,
                "raw_text": generated_text,
                "segments": transcription_segments,
                "generation_time": generation_time / batch_size,
            })
        
        print(f"  Total generation time: {generation_time:.2f}s")
        print(f"  Average time per sample: {generation_time/batch_size:.2f}s")
        
        return results
    
    def transcribe_with_batching(
        self,
        audio_inputs: List,
        batch_size: int = 4,
        max_new_tokens: int = 512,
        temperature: float = 0.0,
        top_p: float = 1.0,
        do_sample: bool = True,
        num_beams: int = 1,
    ) -> List[Dict[str, Any]]:
        """
        Transcribe multiple audio files/arrays with automatic batching.
        
        Args:
            audio_inputs: List of audio file paths or (array, sampling_rate) tuples
            batch_size: Number of samples per batch
            max_new_tokens: Maximum tokens to generate
            temperature: Temperature for sampling
            top_p: Top-p for nucleus sampling
            do_sample: Whether to use sampling
            
        Returns:
            List of transcription results
        """
        all_results = []
        
        # Process in batches
        for i in range(0, len(audio_inputs), batch_size):
            batch_inputs = audio_inputs[i:i + batch_size]
            print(f"\n{'='*60}")
            print(f"Processing batch {i//batch_size + 1}/{(len(audio_inputs) + batch_size - 1)//batch_size}")
            
            batch_results = self.transcribe_batch(
                batch_inputs,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_p=top_p,
                do_sample=do_sample,
                num_beams=num_beams,
            )
            all_results.extend(batch_results)
        
        return all_results


def print_result(result: Dict[str, Any]):
    """Pretty print a single transcription result."""
    print(f"\nFile: {result['file']}")
    print(f"Generation Time: {result['generation_time']:.2f}s")
    print(f"\n--- Raw Output ---")
    print(result['raw_text'][:500] + "..." if len(result['raw_text']) > 500 else result['raw_text'])
    
    if result['segments']:
        print(f"\n--- Structured Output ({len(result['segments'])} segments) ---")
        for seg in result['segments'][:50]:  # Show first 50 segments
            print(f"[{seg.get('start_time', 'N/A')} - {seg.get('end_time', 'N/A')}] "
                  f"Speaker {seg.get('speaker_id', 'N/A')}: {seg.get('text', '')}...")
        if len(result['segments']) > 50:
            print(f"  ... and {len(result['segments']) - 50} more segments")


def load_dataset_and_concatenate(
    dataset_name: str,
    split: str,
    max_duration: float,
    num_audios: int,
    target_sr: int = 24000
) -> Optional[List[np.ndarray]]:
    """
    Load a HuggingFace dataset and concatenate audio samples into long audio chunks.
    (Note, just for demo purpose, not for benchmark evaluation)
    
    Args:
        dataset_name: HuggingFace dataset name (e.g., 'openslr/librispeech_asr')
        split: Dataset split to use (e.g., 'test', 'test.other')
        max_duration: Maximum duration in seconds for each concatenated audio
        num_audios: Number of concatenated audios to create
        target_sr: Target sample rate (default: 24000)
    
    Returns:
        List of concatenated audio arrays, or None if loading fails
    """
    try:
        from datasets import load_dataset
        import torchcodec # just for decode audio in datasets
    except ImportError:
        print("Please install it with: pip install datasets torchcodec")
        return None        
    
    print(f"\nLoading dataset: {dataset_name} (split: {split})")
    print(f"Will create {num_audios} concatenated audio(s), each up to {max_duration:.1f}s ({max_duration/3600:.2f} hours)")
    
    try:
        # Use streaming to avoid downloading the entire dataset
        dataset = load_dataset(dataset_name, split=split, streaming=True)
        print(f"Dataset loaded in streaming mode")
        
        concatenated_audios = []  # List of concatenated audio metadata
        
        # Create multiple concatenated audios based on num_audios
        current_chunks = []
        current_duration = 0.0
        current_samples_used = 0
        sample_idx = 0
        
        for sample in dataset:
            if len(concatenated_audios) >= num_audios:
                break
                
            if 'audio' not in sample:
                continue
            
            audio_data = sample['audio']
            audio_array = audio_data['array']
            sr = audio_data['sampling_rate']
            
            # Resample if needed
            if sr != target_sr:
                duration = len(audio_array) / sr
                new_length = int(duration * target_sr)
                audio_array = np.interp(
                    np.linspace(0, len(audio_array) - 1, new_length),
                    np.arange(len(audio_array)),
                    audio_array
                )
            
            chunk_duration = len(audio_array) / target_sr
            
            # Check if adding this chunk exceeds max_duration
            if current_duration + chunk_duration > max_duration:
                remaining_duration = max_duration - current_duration
                if remaining_duration > 0.5:  # Only add if > 0.5s remaining
                    samples_to_take = int(remaining_duration * target_sr)
                    current_chunks.append(audio_array[:samples_to_take])
                    current_duration += remaining_duration
                    current_samples_used += 1
                
                # Save current concatenated audio and start a new one
                if current_chunks:
                    concatenated_audios.append({
                        'array': np.concatenate(current_chunks),
                        'duration': current_duration,
                        'samples_used': current_samples_used,
                    })
                    print(f"  Created audio {len(concatenated_audios)}: {current_duration:.1f}s from {current_samples_used} samples")
                
                # Reset for next concatenated audio
                current_chunks = []
                current_duration = 0.0
                current_samples_used = 0
                
                if len(concatenated_audios) >= num_audios:
                    break
            
            current_chunks.append(audio_array)
            current_duration += chunk_duration
            current_samples_used += 1
            
            sample_idx += 1
            if sample_idx % 100 == 0:
                print(f"  Processed {sample_idx} samples...")
        
        # Don't forget the last batch if it has content
        if current_chunks and len(concatenated_audios) < num_audios:
            concatenated_audios.append({
                'array': np.concatenate(current_chunks),
                'duration': current_duration,
                'samples_used': current_samples_used,
            })
            print(f"  Created audio {len(concatenated_audios)}: {current_duration:.1f}s from {current_samples_used} samples")
        
        if not concatenated_audios:
            print("Warning: No audio samples found in dataset")
            return None
        
        # Extract arrays and print summary
        result = [a['array'] for a in concatenated_audios]
        total_duration = sum(a['duration'] for a in concatenated_audios)
        total_samples = sum(a['samples_used'] for a in concatenated_audios)
        print(f"\nCreated {len(result)} concatenated audio(s), total {total_duration:.1f}s ({total_duration/60:.1f} min) from {total_samples} samples")
        
        return result
        
    except Exception as e:
        print(f"Error loading dataset: {e}")
        import traceback
        traceback.print_exc()
        return None


def main():
    parser = argparse.ArgumentParser(description="VibeVoice ASR Batch Inference Demo")
    parser.add_argument(
        "--model_path", 
        type=str, 
        default="",
        help="Path to the model checkpoint"
    )
    parser.add_argument(
        "--audio_files", 
        type=str, 
        nargs='+',
        required=False,
        help="Paths to audio files for transcription"
    )
    parser.add_argument(
        "--audio_dir",
        type=str,
        required=False,
        help="Directory containing audio files for batch transcription"
    )
    parser.add_argument(
        "--dataset",
        type=str,
        required=False,
        help="HuggingFace dataset name (e.g., 'openslr/librispeech_asr')"
    )
    parser.add_argument(
        "--split",
        type=str,
        default="test",
        help="Dataset split to use (e.g., 'test', 'test.other', 'test.clean')"
    )
    parser.add_argument(
        "--max_duration",
        type=float,
        default=3600.0,
        help="Maximum duration in seconds for concatenated dataset audio (default: 3600 = 1 hour)"
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        default=2,
        help="Batch size for processing multiple files"
    )
    parser.add_argument(
        "--device", 
        type=str, 
        default="cuda" if torch.cuda.is_available() else ("xpu" if torch.backends.xpu.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu") ),
        choices=["cuda", "cpu", "mps","xpu", "auto"],
        help="Device to run inference on"
    )
    parser.add_argument(
        "--max_new_tokens",
        type=int,
        default=32768,
        help="Maximum number of tokens to generate"
    )
    parser.add_argument(
        "--temperature",
        type=float,
        default=0.0,
        help="Temperature for sampling (0 = greedy decoding)"
    )
    parser.add_argument(
        "--top_p",
        type=float,
        default=1.0,
        help="Top-p for nucleus sampling"
    )
    parser.add_argument(
        "--num_beams",
        type=int,
        default=1,
        help="Number of beams for beam search. Use 1 for greedy/sampling"
    )
    parser.add_argument(
        "--attn_implementation",
        type=str,
        default="auto",
        choices=["flash_attention_2", "sdpa", "eager", "auto"],
        help="Attention implementation to use. 'auto' will select the best available for your device (flash_attention_2 for CUDA, sdpa for MPS/CPU/XPU)"
    )
    
    args = parser.parse_args()
    
    # Auto-detect best attention implementation based on device
    if args.attn_implementation == "auto":
        if args.device == "cuda" and torch.cuda.is_available():
            try:
                import flash_attn
                args.attn_implementation = "flash_attention_2"
            except ImportError:
                print("flash_attn not installed, falling back to sdpa")
                args.attn_implementation = "sdpa"
        else:
            # MPS/XPU/CPU don't support flash_attention_2
            args.attn_implementation = "sdpa"
        print(f"Auto-detected attention implementation: {args.attn_implementation}")
    
    # Collect audio files
    audio_files = []
    concatenated_audio = None  # For storing concatenated dataset audio
    
    if args.audio_files:
        audio_files.extend(args.audio_files)
    
    if args.audio_dir:
        import glob
        for ext in ["*.wav", "*.mp3", "*.flac", "*.mp4", "*.m4a", "*.webm"]:
            audio_files.extend(glob.glob(os.path.join(args.audio_dir, ext)))
    
    if args.dataset:
        concatenated_audio = load_dataset_and_concatenate(
            dataset_name=args.dataset,
            split=args.split,
            max_duration=args.max_duration,
            num_audios=args.batch_size,
        )
        if concatenated_audio is None:
            return
    
    if len(audio_files) == 0 and concatenated_audio is None:
        print("No audio files provided. Please specify --audio_files, --audio_dir, or --dataset.")
        return
    
    if audio_files:
        print(f"\nAudio files to process ({len(audio_files)}):")
        for f in audio_files:
            print(f"  - {f}")
    
    if concatenated_audio:
        print(f"\nConcatenated dataset audios: {len(concatenated_audio)} audio(s)")
    
    # Initialize model
    # Handle MPS device and dtype
    if args.device == "mps":
        model_dtype = torch.float32  # MPS works better with float32
    elif args.device == "xpu":
        model_dtype = torch.float32
    elif args.device == "cpu":
        model_dtype = torch.float32
    else:
        model_dtype = torch.bfloat16
    
    asr = VibeVoiceASRBatchInference(
        model_path=args.model_path,
        device=args.device,
        dtype=model_dtype,
        attn_implementation=args.attn_implementation
    )
    
    # If temperature is 0, use greedy decoding (no sampling)
    do_sample = args.temperature > 0
    
    # Combine all audio inputs
    all_audio_inputs = audio_files + (concatenated_audio or [])
    
    print("\n" + "="*80)
    print(f"Processing {len(all_audio_inputs)} audio(s)")
    print("="*80)
    
    all_results = asr.transcribe_with_batching(
        all_audio_inputs,
        batch_size=args.batch_size,
        max_new_tokens=args.max_new_tokens,
        temperature=args.temperature,
        top_p=args.top_p,
        do_sample=do_sample,
        num_beams=args.num_beams,
    )
    
    # Print results
    print("\n" + "="*80)
    print("Results")
    print("="*80)
    for result in all_results:
        print("\n" + "-"*60)
        print_result(result)


if __name__ == "__main__":
    main()