"""
Orpheus TTS - Continuous Batching Test v3
==========================================

Uses AsyncLLMEngine with correct token decoding based on Axolotl's preprocessing.
Token format per 7-token frame:
  [layer0, layer1_a, layer2_a, layer2_b, layer1_b, layer2_c, layer2_d]
  Where:
  - layer0: 128266 + value
  - layer1_a: 128266 + 4096 + value
  - layer2_a: 128266 + 2*4096 + value
  - layer2_b: 128266 + 3*4096 + value
  - layer1_b: 128266 + 4*4096 + value
  - layer2_c: 128266 + 5*4096 + value
  - layer2_d: 128266 + 6*4096 + value
"""

import os
import sys
import time
import wave
import asyncio
import numpy as np

os.environ["VLLM_ATTENTION_BACKEND"] = "FLASH_ATTN"

import torch
from transformers import AutoTokenizer
from vllm import AsyncLLMEngine, AsyncEngineArgs, SamplingParams
from snac import SNAC

# Orpheus special tokens
START_TOKEN = 128259
END_TOKENS = [128009, 128260, 128261, 128257]
STOP_TOKEN = 128258
AUDIO_TOKEN_BASE = 128266


def decode_tokens_to_audio(token_ids, snac_model):
    """Decode Orpheus tokens to audio using SNAC with correct layer offsets."""
    # Filter audio tokens and decode by layer
    audio_frames = []
    
    for t in token_ids:
        if isinstance(t, str):
            continue
        if t >= AUDIO_TOKEN_BASE:
            # Determine which layer this token belongs to
            offset = t - AUDIO_TOKEN_BASE
            layer = offset // 4096
            value = offset % 4096
            audio_frames.append((layer, value))
    
    if len(audio_frames) < 7:
        return None
    
    # Group into 7-token frames and extract codes for each layer
    num_complete_frames = len(audio_frames) // 7
    if num_complete_frames == 0:
        return None
    
    codes_0 = []  # layer 0
    codes_1 = []  # layer 1
    codes_2 = []  # layer 2
    
    for i in range(num_complete_frames):
        base = i * 7
        # Frame format: [l0, l1_a, l2_a, l2_b, l1_b, l2_c, l2_d]
        codes_0.append(audio_frames[base][1])      # layer 0 value
        codes_1.append(audio_frames[base + 1][1])  # layer 1 first
        codes_1.append(audio_frames[base + 4][1])  # layer 1 second
        codes_2.append(audio_frames[base + 2][1])  # layer 2 first
        codes_2.append(audio_frames[base + 3][1])  # layer 2 second
        codes_2.append(audio_frames[base + 5][1])  # layer 2 third
        codes_2.append(audio_frames[base + 6][1])  # layer 2 fourth
    
    try:
        # Convert to tensors with correct shape
        with torch.no_grad():
            codes = [
                torch.tensor(codes_0, dtype=torch.int64).unsqueeze(0).to("cuda"),
                torch.tensor(codes_1, dtype=torch.int64).unsqueeze(0).to("cuda"),
                torch.tensor(codes_2, dtype=torch.int64).unsqueeze(0).to("cuda"),
            ]
            audio = snac_model.decode(codes)
        
        return audio.squeeze().cpu().numpy()
    except Exception as e:
        print(f"    Decode error: {e}")
        return None


async def run_tests():
    """Run continuous batching tests."""
    print("=" * 60)
    print("ORPHEUS TTS - CONTINUOUS BATCHING TEST v3")
    print("=" * 60)
    
    # Load tokenizer
    print("\n[1] Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained("canopylabs/orpheus-3b-0.1-ft")
    
    # Create AsyncLLMEngine
    print("\n[2] Loading vLLM AsyncLLMEngine...")
    start_load = time.time()
    
    engine_args = AsyncEngineArgs(
        model="canopylabs/orpheus-3b-0.1-ft",
        dtype="bfloat16",
        max_model_len=4096,
        gpu_memory_utilization=0.9,
        max_num_seqs=8,  # Continuous batching
        enable_chunked_prefill=True,
        enable_prefix_caching=True,
        enforce_eager=False,
    )
    
    engine = AsyncLLMEngine.from_engine_args(engine_args)
    load_time = time.time() - start_load
    print(f"    vLLM loaded in {load_time:.2f}s")
    
    # Load SNAC decoder
    print("\n[3] Loading SNAC decoder...")
    snac = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").eval().to("cuda")
    print("    SNAC loaded!")
    
    # Sampling params
    sampling_params = SamplingParams(
        temperature=0.2,
        top_p=0.9,
        max_tokens=4096,
        stop_token_ids=[STOP_TOKEN],
        repetition_penalty=1.1,
    )
    
    def format_prompt(text: str, voice: str = "tara") -> str:
        """Format prompt with Orpheus special tokens."""
        adapted_prompt = f"{voice}: {text}"
        prompt_tokens = tokenizer(adapted_prompt, return_tensors="pt")
        start_token = torch.tensor([[START_TOKEN]], dtype=torch.int64)
        end_tokens = torch.tensor([END_TOKENS], dtype=torch.int64)
        all_input_ids = torch.cat([start_token, prompt_tokens.input_ids, end_tokens], dim=1)
        prompt_string = tokenizer.decode(all_input_ids[0])
        return prompt_string
    
    async def generate_speech(text: str, voice: str = "tara", request_id: str = None):
        """Generate speech for a single request."""
        prompt_string = format_prompt(text, voice)
        request_id = request_id or f"req_{time.time()}"
        
        start = time.time()
        token_ids = []
        
        async for output in engine.generate(prompt_string, sampling_params, request_id):
            token_ids = list(output.outputs[0].token_ids)
        
        gen_time = time.time() - start
        
        # Count audio tokens
        audio_token_count = sum(1 for t in token_ids if isinstance(t, int) and t >= AUDIO_TOKEN_BASE)
        
        # Decode to audio
        audio = decode_tokens_to_audio(token_ids, snac)
        
        if audio is None:
            return {
                'success': False,
                'text': text[:30],
                'gen_time': gen_time,
                'tokens': len(token_ids),
                'audio_tokens': audio_token_count,
            }
        
        audio_duration = len(audio) / 24000
        rtf = gen_time / audio_duration if audio_duration > 0 else float('inf')
        
        return {
            'success': True,
            'text': text[:30],
            'gen_time': gen_time,
            'tokens': len(token_ids),
            'audio_tokens': audio_token_count,
            'audio_duration': audio_duration,
            'rtf': rtf,
            'audio': audio,
        }
    
    # =========================================================================
    # TEST 1: SEQUENTIAL (baseline)
    # =========================================================================
    print("\n" + "=" * 60)
    print("[4] Test SEQUENTIAL (baseline)")
    print("=" * 60)
    
    test_texts = [
        "Hello! This is the first test.",
        "Second test to measure performance.",
        "Third test for consistent results.",
    ]
    
    sequential_results = []
    total_seq_time = 0
    
    for i, text in enumerate(test_texts, 1):
        print(f"\n    Test {i}: \"{text}\"")
        result = await generate_speech(text, request_id=f"seq_{i}")
        
        if result['success']:
            print(f"    -> Time: {result['gen_time']:.2f}s | Audio: {result['audio_duration']:.2f}s | RTF: {result['rtf']:.3f}")
            sequential_results.append(result)
            total_seq_time += result['gen_time']
        else:
            print(f"    -> ERROR: {result['tokens']} tokens ({result['audio_tokens']} audio), no audio")
    
    # =========================================================================
    # TEST 2: PARALLEL (Continuous Batching)
    # =========================================================================
    print("\n" + "=" * 60)
    print("[5] Test PARALLEL (Continuous Batching)")
    print("=" * 60)
    
    parallel_texts = [
        "Hello, how are you today?",
        "The weather is beautiful outside.",
        "I love programming with Python.",
        "Machine learning is fascinating.",
    ]
    
    batch_results_summary = []
    
    for num_concurrent in [2, 4]:
        print(f"\n    === {num_concurrent} CONCURRENT REQUESTS ===")
        
        texts = parallel_texts[:num_concurrent]
        
        start_batch = time.time()
        tasks = [
            generate_speech(text, request_id=f"par_{num_concurrent}_{i}")
            for i, text in enumerate(texts)
        ]
        results = await asyncio.gather(*tasks)
        batch_time = time.time() - start_batch
        
        # Calculate metrics
        successful = [r for r in results if r['success']]
        total_audio = sum(r['audio_duration'] for r in successful)
        
        print(f"\n    Results:")
        for r in results:
            if r['success']:
                print(f"    - \"{r['text']}...\" -> {r['gen_time']:.2f}s | {r['audio_duration']:.2f}s | RTF: {r['rtf']:.3f}")
            else:
                print(f"    - \"{r['text']}...\" -> FAILED ({r['audio_tokens']} audio tokens)")
        
        if total_audio > 0:
            batch_rtf = batch_time / total_audio
            throughput = len(successful) / batch_time
            
            print(f"\n    Aggregate Metrics:")
            print(f"    - Total wall-clock time: {batch_time:.2f}s")
            print(f"    - Success rate: {len(successful)}/{len(texts)}")
            print(f"    - Total audio generated: {total_audio:.2f}s")
            print(f"    - Batch RTF: {batch_rtf:.3f}")
            print(f"    - Throughput: {throughput:.2f} req/s")
            print(f"    - Effective speed: {1/batch_rtf:.1f}x real-time")
            
            batch_results_summary.append({
                'concurrent': num_concurrent,
                'batch_time': batch_time,
                'total_audio': total_audio,
                'batch_rtf': batch_rtf,
                'throughput': throughput,
            })
    
    # =========================================================================
    # FINAL SUMMARY
    # =========================================================================
    print("\n" + "=" * 60)
    print("[6] === FINAL SUMMARY ===")
    print("=" * 60)
    
    if sequential_results:
        seq_rtfs = [r['rtf'] for r in sequential_results]
        avg_seq_rtf = sum(seq_rtfs) / len(seq_rtfs)
        total_seq_audio = sum(r['audio_duration'] for r in sequential_results)
        
        print(f"\n    SEQUENTIAL (baseline):")
        print(f"    - Average RTF: {avg_seq_rtf:.3f}")
        print(f"    - Speed: {1/avg_seq_rtf:.1f}x real-time")
        
        if batch_results_summary:
            print(f"\n    CONTINUOUS BATCHING:")
            for bs in batch_results_summary:
                speedup = avg_seq_rtf / bs['batch_rtf'] if bs['batch_rtf'] > 0 else 0
                print(f"    - {bs['concurrent']} concurrent: RTF={bs['batch_rtf']:.3f}, {bs['throughput']:.2f} req/s, {speedup:.1f}x speedup")
        
        # Capacity estimate
        users_seq = 1 / avg_seq_rtf if avg_seq_rtf > 0 else 0
        print(f"\n    RTX 4090 CAPACITY ESTIMATE:")
        print(f"    - Sequential: ~{users_seq:.0f} real-time users")
        if batch_results_summary:
            best_batch = min(batch_results_summary, key=lambda x: x['batch_rtf'])
            users_batch = best_batch['concurrent'] / best_batch['batch_rtf'] if best_batch['batch_rtf'] > 0 else 0
            print(f"    - With batching: ~{users_batch:.0f} real-time users")
    
    print("=" * 60)
    
    # Save audio
    if sequential_results and sequential_results[-1]['success']:
        audio = sequential_results[-1]['audio']
        output_path = "/root/test_batching_v3_output.wav"
        audio_int16 = (audio * 32767).astype(np.int16)
        with wave.open(output_path, "wb") as wf:
            wf.setnchannels(1)
            wf.setsampwidth(2)
            wf.setframerate(24000)
            wf.writeframes(audio_int16.tobytes())
        print(f"\n    Audio saved to: {output_path}")


if __name__ == "__main__":
    asyncio.run(run_tests())