"""
Maya-1-Voice VLLM Streaming Inference - Standalone Reference Implementation

This is a complete, self-contained example for using Maya-1-Voice TTS model with VLLM and SNAC.
Demonstrates streaming audio generation with sliding window approach for smooth playback.

Requirements:
    pip install vllm transformers torch snac numpy

Usage:
    python vllm_streaming_inference.py

Author: Maya-1-Voice Team
License: MIT
"""

import torch
import numpy as np
import asyncio
from typing import List, Optional, AsyncGenerator
from transformers import AutoTokenizer
from vllm import AsyncLLMEngine, AsyncEngineArgs, SamplingParams
from snac import SNAC


# ============================================================================
# CONSTANTS
# ============================================================================

# Special control tokens
CODE_START_TOKEN_ID = 128257  # Start of Speech (SOS)
CODE_END_TOKEN_ID = 128258    # End of Speech (EOS) - stop token for audio
CODE_TOKEN_OFFSET = 128266    # Start of SNAC codes

# SNAC token range (7 tokens per frame, 4096 codes per level)
SNAC_MIN_ID = 128266
SNAC_MAX_ID = 156937  # 128266 + (7 * 4096) - 1

# SNAC configuration
SNAC_MODEL_NAME = "hubertsiuzdak/snac_24khz"
SNAC_SAMPLE_RATE = 24000
SNAC_TOKENS_PER_FRAME = 7

# Generation parameters
DEFAULT_TEMPERATURE = 0.4
DEFAULT_TOP_P = 0.9
DEFAULT_MAX_TOKENS = 2000
DEFAULT_MIN_TOKENS = 28  # At least 4 SNAC frames
DEFAULT_REPETITION_PENALTY = 1.1


# ============================================================================
# SNAC DECODER
# ============================================================================

class SNACDecoder:
    """
    Decodes SNAC tokens (7-token frames) to audio waveforms.
    
    The unpacking logic converts flat 7-token frames back to hierarchical
    3-level SNAC codes (matching the training preprocessing exactly).
    """
    
    def __init__(self, device: str = "cuda"):
        """Initialize SNAC decoder with 24kHz model."""
        self.device = device
        print(f"🎵 Loading SNAC 24kHz model to {device}...")
        self.snac_model = SNAC.from_pretrained(SNAC_MODEL_NAME).eval().to(device)
        print(f"✅ SNAC decoder initialized")
    
    def unpack_snac_from_7(self, vocab_ids: List[int]) -> List[List[int]]:
        """
        Unpack 7-token SNAC frames to 3 hierarchical levels.
        
        This is the EXACT INVERSE of training preprocessing.
        
        Frame structure (7 tokens per frame):
        [slot0, slot1, slot2, slot3, slot4, slot5, slot6]
        
        Unpacking to [L1, L2, L3]:
        - slot0 → L1[i]       (coarse: 1x rate)
        - slot1 → L2[2*i]     (medium: 2x rate, even)
        - slot2 → L3[4*i+0]   (fine: 4x rate)
        - slot3 → L3[4*i+1]
        - slot4 → L2[2*i+1]   (medium: odd)
        - slot5 → L3[4*i+2]
        - slot6 → L3[4*i+3]
        
        Args:
            vocab_ids: List of SNAC token IDs (128266-156937), length divisible by 7
        
        Returns:
            [L1, L2, L3] where L1=n, L2=2n, L3=4n elements
        """
        # Remove EOS token if present
        if vocab_ids and vocab_ids[-1] == CODE_END_TOKEN_ID:
            vocab_ids = vocab_ids[:-1]
        
        # Ensure complete frames
        frames = len(vocab_ids) // SNAC_TOKENS_PER_FRAME
        vocab_ids = vocab_ids[:frames * SNAC_TOKENS_PER_FRAME]
        
        if frames == 0:
            return [[], [], []]
        
        l1, l2, l3 = [], [], []
        
        for i in range(frames):
            slots = vocab_ids[i*7:(i+1)*7]
            
            # Subtract offset and mod 4096 to get original SNAC codes
            l1.append((slots[0] - CODE_TOKEN_OFFSET) % 4096)
            l2.extend([
                (slots[1] - CODE_TOKEN_OFFSET) % 4096,  # Even
                (slots[4] - CODE_TOKEN_OFFSET) % 4096,  # Odd
            ])
            l3.extend([
                (slots[2] - CODE_TOKEN_OFFSET) % 4096,
                (slots[3] - CODE_TOKEN_OFFSET) % 4096,
                (slots[5] - CODE_TOKEN_OFFSET) % 4096,
                (slots[6] - CODE_TOKEN_OFFSET) % 4096,
            ])
        
        return [l1, l2, l3]
    
    @torch.inference_mode()
    def decode(
        self, 
        snac_tokens: List[int], 
        use_sliding_window: bool = False
    ) -> Optional[np.ndarray]:
        """
        Decode SNAC tokens to audio waveform.
        
        Args:
            snac_tokens: List of SNAC token IDs (7*n tokens)
            use_sliding_window: If True, return only middle 2048 samples 
                               (for smooth streaming without pops/clicks)
        
        Returns:
            Audio waveform as float32 numpy array, 24kHz mono
        """
        if len(snac_tokens) < SNAC_TOKENS_PER_FRAME:
            return None
        
        # Unpack to 3 hierarchical levels
        levels = self.unpack_snac_from_7(snac_tokens)
        
        if not levels[0]:
            return None
        
        # Convert to tensors
        codes = [
            torch.tensor(level, dtype=torch.long, device=self.device).unsqueeze(0)
            for level in levels
        ]
        
        # Decode through SNAC quantizer + decoder
        z_q = self.snac_model.quantizer.from_codes(codes)
        audio = self.snac_model.decoder(z_q)
        
        # Extract audio: [batch, 1, samples] → [samples]
        audio = audio[0, 0].cpu().numpy()
        
        # Sliding window mode: keep middle 2048 samples only
        # This eliminates popping/cracking in streaming by overlapping windows
        if use_sliding_window and len(audio) >= 4096:
            audio = audio[2048:4096]
        
        return audio
    
    def decode_to_bytes(
        self, 
        snac_tokens: List[int], 
        use_sliding_window: bool = False
    ) -> Optional[bytes]:
        """
        Decode SNAC tokens to audio bytes (int16 PCM).
        
        Args:
            snac_tokens: List of SNAC token IDs
            use_sliding_window: Use sliding window for smooth streaming
        
        Returns:
            Audio as bytes (int16 PCM, 24kHz mono)
        """
        audio = self.decode(snac_tokens, use_sliding_window=use_sliding_window)
        
        if audio is None:
            return None
        
        # Convert float32 to int16 PCM
        audio_int16 = (audio * 32767).astype(np.int16)
        return audio_int16.tobytes()


# ============================================================================
# CUSTOM LOGITS PROCESSOR
# ============================================================================

class OnlyAudioAfterSOS:
    """
    Restricts vocabulary to SNAC codes + EOS after SOS token.
    
    This prevents the model from generating text tokens during audio phase,
    which would cause "hallucination" where the model repeats description text
    instead of generating proper audio codes.
    """
    
    def __init__(self):
        self._seen_sos = False
    
    def __call__(
        self,
        prompt_token_ids: List[int],
        generated_token_ids: List[int],
        logits: torch.Tensor,
    ) -> torch.Tensor:
        """
        Apply constraint: after SOS, only allow SNAC codes + EOS.
        
        Args:
            prompt_token_ids: Original prompt token IDs
            generated_token_ids: Tokens generated so far
            logits: Logits for next token [vocab_size]
        
        Returns:
            Modified logits with masked tokens
        """
        # Check if SOS has been generated
        if not self._seen_sos:
            all_token_ids = prompt_token_ids + generated_token_ids
            if CODE_START_TOKEN_ID in all_token_ids:
                self._seen_sos = True
            else:
                return logits  # No constraint yet
        
        # Apply constraint: mask all tokens except SNAC codes + EOS
        mask = torch.full_like(logits, float('-inf'))
        mask[SNAC_MIN_ID:SNAC_MAX_ID + 1] = 0  # Allow SNAC codes
        mask[CODE_END_TOKEN_ID] = 0            # Allow EOS
        
        return logits + mask
    
    def reset(self):
        """Reset state for reuse across generations."""
        self._seen_sos = False


# ============================================================================
# MAYA-1-VOICE MODEL
# ============================================================================

class Maya1VoiceModel:
    """
    Maya-1-Voice TTS Model with VLLM inference engine.
    
    Handles model loading, tokenizer initialization, and VLLM engine setup.
    """
    
    def __init__(
        self,
        model_path: str,
        dtype: str = "bfloat16",
        max_model_len: int = 8192,
        gpu_memory_utilization: float = 0.85,
    ):
        """
        Initialize Maya-1-Voice model with VLLM.
        
        Args:
            model_path: Path to model checkpoint (local or HuggingFace)
            dtype: Model precision (bfloat16 recommended)
            max_model_len: Maximum sequence length
            gpu_memory_utilization: GPU memory fraction to use (0.0-1.0)
        """
        self.model_path = model_path
        
        print(f"🚀 Initializing Maya-1-Voice Model")
        print(f"📁 Model: {model_path}")
        print(f"🔢 Dtype: {dtype}")
        
        # Load tokenizer (must be from checkpoint with emotion tags)
        print(f"📝 Loading tokenizer...")
        self.tokenizer = AutoTokenizer.from_pretrained(
            model_path,
            trust_remote_code=True,
        )
        print(f"✅ Tokenizer loaded: {len(self.tokenizer)} tokens")
        
        # Initialize VLLM async engine
        print(f"🔧 Initializing VLLM engine...")
        engine_args = AsyncEngineArgs(
            model=model_path,
            tokenizer=model_path,
            dtype=dtype,
            max_model_len=max_model_len,
            gpu_memory_utilization=gpu_memory_utilization,
            trust_remote_code=True,
        )
        
        self.engine = AsyncLLMEngine.from_engine_args(engine_args)
        print(f"✅ VLLM engine ready")
    
    def build_prompt(self, description: str, text: str) -> str:
        """
        Build prompt in Maya-1-Voice format using chat template.
        
        Format: Chat template with <description="..."> text as content
        
        The model expects:
        1. Description of voice/character
        2. Text to synthesize (optionally with <emotion> tags)
        
        Args:
            description: Voice description 
                Example: "Realistic male voice in the 30s age with american accent. 
                         Normal pitch, warm timbre, conversational pacing."
            text: Text to synthesize
                Example: "Hello world! <excited> This is amazing!"
        
        Returns:
            Formatted prompt string using chat template
        """
        content = f'<description="{description}"> {text}'
        messages = [{"role": "user", "content": content}]
        return self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)


# ============================================================================
# STREAMING PIPELINE
# ============================================================================

class Maya1VoiceStreamingPipeline:
    """
    Streaming TTS pipeline using sliding window approach.
    
    This generates smooth audio by:
    1. Streaming tokens from VLLM as they're generated
    2. Every 7 tokens, decoding the last 28 tokens (4 frames) - sliding window
    3. Keeping only middle 2048 samples from each decode
    4. Creating natural overlap between chunks for artifact-free playback
    """
    
    def __init__(self, model: Maya1VoiceModel, snac_decoder: SNACDecoder):
        """Initialize streaming pipeline."""
        self.model = model
        self.snac_decoder = snac_decoder
        print(f"🌊 Maya-1-Voice Streaming Pipeline initialized")
    
    async def generate_speech_stream(
        self,
        description: str,
        text: str,
        temperature: float = DEFAULT_TEMPERATURE,
        top_p: float = DEFAULT_TOP_P,
        max_tokens: int = DEFAULT_MAX_TOKENS,
        repetition_penalty: float = DEFAULT_REPETITION_PENALTY,
    ) -> AsyncGenerator[bytes, None]:
        """
        Generate speech audio with streaming.
        
        Args:
            description: Voice/character description
            text: Text to synthesize (with optional <emotion> tags)
            temperature: Sampling temperature (lower = more stable)
            top_p: Nucleus sampling
            max_tokens: Max SNAC tokens to generate
            repetition_penalty: Prevent repetition loops
        
        Yields:
            Audio chunks as bytes (int16 PCM, 24kHz mono)
        """
        print(f"\n🌊 Starting streaming generation")
        print(f"📝 Description: {description[:80]}...")
        print(f"💬 Text: {text}")
        
        # Build prompt
        prompt = self.model.build_prompt(description, text)
        
        # Configure sampling (removed custom logits processor for V1 compatibility)
        sampling_params = SamplingParams(
            temperature=temperature,
            top_p=top_p,
            max_tokens=max_tokens,
            min_tokens=DEFAULT_MIN_TOKENS,
            repetition_penalty=repetition_penalty,
            stop_token_ids=[CODE_END_TOKEN_ID],  # Stop on audio EOS
        )
        
        print(f"🎲 Sampling: temp={temperature}, top_p={top_p}, max_tokens={max_tokens}")
        
        # Token buffer for sliding window
        token_buffer = []
        total_tokens = 0
        total_chunks = 0
        
        # Generate with VLLM
        import uuid
        import time
        request_id = f"maya1voice-{uuid.uuid4().hex[:8]}-{int(time.time() * 1000000)}"
        
        results_generator = self.model.engine.generate(
            prompt=prompt,
            sampling_params=sampling_params,
            request_id=request_id,
        )
        
        # Stream tokens with sliding window decoding
        async for request_output in results_generator:
            generated_ids = request_output.outputs[0].token_ids
            
            # Process only new tokens
            new_tokens = generated_ids[total_tokens:]
            total_tokens = len(generated_ids)
            
            # Filter and buffer SNAC tokens only
            for token_id in new_tokens:
                if SNAC_MIN_ID <= token_id <= SNAC_MAX_ID:
                    token_buffer.append(token_id)
                    
                    # Sliding window: process every 7 tokens when buffer > 27
                    # Take last 28 tokens (4 frames) for smooth overlap
                    if len(token_buffer) % 7 == 0 and len(token_buffer) > 27:
                        window_tokens = token_buffer[-28:]
                        
                        # Decode with sliding window (returns middle 2048 samples)
                        audio_bytes = self.snac_decoder.decode_to_bytes(
                            window_tokens, 
                            use_sliding_window=True
                        )
                        
                        if audio_bytes:
                            total_chunks += 1
                            if total_chunks == 1:
                                print(f"🎵 First chunk decoded ({len(audio_bytes)} bytes)")
                            yield audio_bytes
        
        print(f"✅ Streaming complete: {total_tokens} tokens → {total_chunks} chunks")


# ============================================================================
# MAIN EXAMPLE
# ============================================================================

async def main():
    """
    Example usage of Maya-1-Voice streaming inference.
    
    This demonstrates:
    1. Model initialization
    2. SNAC decoder setup
    3. Streaming generation
    4. Audio chunk handling
    """
    
    # Configuration
    MODEL_PATH = "/home/ubuntu/veena_temp/maya-1-voice"  # Local model path
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
    
    print("=" * 80)
    print("Maya-1-Voice VLLM Streaming Inference Example")
    print("=" * 80)
    
    # Initialize model
    model = Maya1VoiceModel(
        model_path=MODEL_PATH,
        dtype="bfloat16",
        max_model_len=8192,
        gpu_memory_utilization=0.8,  # Reduced for available GPU memory (12GB free)
    )
    
    # Initialize SNAC decoder
    snac_decoder = SNACDecoder(device=DEVICE)
    
    # Create pipeline
    pipeline = Maya1VoiceStreamingPipeline(model, snac_decoder)
    
    # Example 1: Professional voice
    description = (
        "Realistic male voice in the 30s age with american accent. "
        "Normal pitch, warm timbre, conversational pacing, neutral tone delivery at med intensity."
    )
    text = "Hello! This is a test of the Maya-1-Voice text-to-speech system."
    
    print(f"\n{'='*80}")
    print("Example 1: Professional Voice")
    print(f"{'='*80}")
    
    audio_chunks = []
    async for chunk in pipeline.generate_speech_stream(
        description=description,
        text=text,
        temperature=0.4,
        max_tokens=500,
    ):
        audio_chunks.append(chunk)
        print(f"📦 Received chunk {len(audio_chunks)}: {len(chunk)} bytes")
    
    # Combine chunks
    full_audio = b''.join(audio_chunks)
    print(f"\n✅ Total audio: {len(full_audio)} bytes ({len(full_audio)//2} samples, {len(full_audio)/2/24000:.2f}s)")
    
    # Save audio (optional)
    try:
        import wave
        output_file = "output_example1.wav"
        with wave.open(output_file, 'wb') as wav:
            wav.setnchannels(1)  # Mono
            wav.setsampwidth(2)  # 16-bit
            wav.setframerate(24000)  # 24kHz
            wav.writeframes(full_audio)
        print(f"💾 Saved to {output_file}")
    except ImportError:
        print(f"⚠️  Install 'wave' module to save audio files")
    
    # Example 2: Character voice with emotions
    print(f"\n{'='*80}")
    print("Example 2: Character Voice with Emotions")
    print(f"{'='*80}")
    
    description = (
        "Creative, dark_villain character. Male voice in their 40s with british accent. "
        "Low pitch, gravelly timbre, slow pacing, angry tone at high intensity."
    )
    text = "The darkness isn't coming... <angry> it's already here!"
    
    audio_chunks = []
    async for chunk in pipeline.generate_speech_stream(
        description=description,
        text=text,
        temperature=0.5,
        max_tokens=800,
    ):
        audio_chunks.append(chunk)
        print(f"📦 Received chunk {len(audio_chunks)}: {len(chunk)} bytes")
    
    full_audio = b''.join(audio_chunks)
    print(f"\n✅ Total audio: {len(full_audio)} bytes ({len(full_audio)//2} samples, {len(full_audio)/2/24000:.2f}s)")
    
    # Save audio
    try:
        import wave
        output_file = "output_example2.wav"
        with wave.open(output_file, 'wb') as wav:
            wav.setnchannels(1)
            wav.setsampwidth(2)
            wav.setframerate(24000)
            wav.writeframes(full_audio)
        print(f"💾 Saved to {output_file}")
    except ImportError:
        pass
    
    print(f"\n{'='*80}")
    print("🎉 Examples complete!")
    print(f"{'='*80}")


if __name__ == "__main__":
    # Run async main
    asyncio.run(main())