"""
Teste de Limite MÁXIMO - Até onde vai antes de travar?
"""
import os
os.environ["VLLM_ATTENTION_BACKEND"] = "FLASH_ATTN"

import torch
import time
import asyncio
from transformers import AutoTokenizer
from vllm import AsyncLLMEngine, AsyncEngineArgs, SamplingParams

START_TOKEN = 128259
END_TOKENS = [128009, 128260, 128261, 128257]
STOP_TOKEN = 128258
AUDIO_TOKEN_BASE = 128266

async def main():
    print("=" * 70)
    print("TESTE DE LIMITE MÁXIMO - ATÉ ONDE VAI?")
    print("=" * 70)

    print("\n[1] Carregando modelo com max_num_seqs=256...")
    tokenizer = AutoTokenizer.from_pretrained("canopylabs/orpheus-3b-0.1-ft")

    # Máximo possível
    engine_args = AsyncEngineArgs(
        model="canopylabs/orpheus-3b-0.1-ft",
        dtype="bfloat16",
        max_model_len=4096,
        gpu_memory_utilization=0.95,
        max_num_seqs=256,  # Tentando o máximo
        enable_chunked_prefill=True,
        enable_prefix_caching=True,
        enforce_eager=False,
    )
    engine = AsyncLLMEngine.from_engine_args(engine_args)

    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, voice="tara"):
        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)
        return tokenizer.decode(all_input_ids[0])

    async def measure_request(text, request_id):
        prompt_string = format_prompt(text)
        start = time.time()
        ttff = None
        audio_token_count = 0

        async for output in engine.generate(prompt_string, sampling_params, request_id):
            current_audio = sum(1 for t in output.outputs[0].token_ids if t >= AUDIO_TOKEN_BASE)
            if ttff is None and current_audio >= 7:
                ttff = time.time() - start
            audio_token_count = current_audio

        total_time = time.time() - start
        audio_duration = (audio_token_count // 7) * 0.023

        return {
            'ttff': ttff or total_time,
            'total_time': total_time,
            'audio_duration': audio_duration,
        }

    # Gerar 256 frases de teste
    base_texts = [
        "Hello, how are you?", "Good morning!", "Nice to meet you.",
        "How is the weather?", "I love learning.", "This is great.",
        "Thank you so much.", "Have a nice day.", "See you later.",
        "What time is it?", "Where are you from?", "I am happy.",
        "Let's practice.", "Very interesting.", "Good job today.",
        "Keep it up.", "Well done!", "Excellent work.",
    ]
    test_texts = (base_texts * 15)[:256]  # 256 frases

    print("\n[2] Warmup...")
    await measure_request("Warmup.", "warmup")

    print("\n[3] Testando limites...")
    print("=" * 70)

    results_summary = []

    for num_users in [32, 48, 64, 96, 128, 160, 200, 256]:
        print(f"\n>>> TESTANDO {num_users} USUÁRIOS <<<")

        texts = test_texts[:num_users]

        try:
            start_batch = time.time()
            tasks = [measure_request(text, f"u{num_users}_{i}") for i, text in enumerate(texts)]
            results = await asyncio.gather(*tasks, return_exceptions=True)
            batch_time = time.time() - start_batch

            errors = [r for r in results if isinstance(r, Exception)]
            successful = [r for r in results if not isinstance(r, Exception)]

            if errors:
                print(f"    ERROS: {len(errors)}/{num_users}")
                print(f"    Tipo: {type(errors[0]).__name__}")
                if len(successful) == 0:
                    print(f"    >>> FALHA TOTAL - LIMITE ATINGIDO <<<")
                    break

            if successful:
                ttffs = [r['ttff'] for r in successful]
                total_audio = sum(r['audio_duration'] for r in successful)

                rtf = batch_time / total_audio if total_audio > 0 else 999
                realtime = total_audio / batch_time if batch_time > 0 else 0

                print(f"    OK: {len(successful)}/{num_users} | Time: {batch_time:.1f}s | TTFF max: {max(ttffs)*1000:.0f}ms | RTF: {rtf:.3f}")

                results_summary.append({
                    'users': num_users,
                    'success': len(successful),
                    'errors': len(errors),
                    'batch_time': batch_time,
                    'ttff_max': max(ttffs),
                    'rtf': rtf,
                    'realtime': realtime
                })

                # Se TTFF > 2 segundos, parar
                if max(ttffs) > 2.0:
                    print(f"    >>> TTFF muito alto ({max(ttffs)*1000:.0f}ms) - PARANDO <<<")
                    break

        except Exception as e:
            print(f"    CRASH: {type(e).__name__}: {str(e)[:80]}")
            break

    print("\n" + "=" * 70)
    print("RESUMO FINAL")
    print("=" * 70)

    print("\n| Users | OK  | Erros | Time   | TTFF Max  | RTF   |")
    print("|-------|-----|-------|--------|-----------|-------|")
    for r in results_summary:
        print(f"| {r['users']:5} | {r['success']:3} | {r['errors']:5} | {r['batch_time']:6.1f}s | {r['ttff_max']*1000:7.0f}ms | {r['rtf']:.3f} |")

    if results_summary:
        last_good = [r for r in results_summary if r['errors'] == 0 and r['ttff_max'] < 1.0]
        if last_good:
            best = max(last_good, key=lambda x: x['users'])
            print(f"\n>>> MÁXIMO ESTÁVEL: {best['users']} usuários (TTFF < 1s, 0 erros) <<<")

    print("=" * 70)

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