#!/usr/bin/env python3
"""
Comprehensive WebSocket performance test with detailed metrics display.
This script tests the TTS pipeline and displays:
- Soprano synthesis time
- RVC processing time
- Overall pipeline realtime factor
"""

import asyncio
import websockets
import json
import time
import subprocess
import re

async def test_synthesis():
    uri = "ws://localhost:8765/ws/stream"
    
    # Test sentences of varying lengths
    test_cases = [
        ("Short", "Hello, this is a test."),
        ("Medium", "The quick brown fox jumps over the lazy dog."),
        ("Long", "Artificial intelligence is revolutionizing the way we interact with technology.")
    ]
    
    print("\n" + "="*80)
    print(" "*20 + "TTS PIPELINE PERFORMANCE TEST")
    print("="*80 + "\n")
    
    async with websockets.connect(uri) as websocket:
        
        for idx, (label, sentence) in enumerate(test_cases, 1):
            print(f"\n{'─'*80}")
            print(f"Test {idx} ({label}): {sentence}")
            print('─'*80)
            
            # Track client-side timing
            client_start = time.time()
            
            # Send the sentence token by token
            for char in sentence:
                await websocket.send(json.dumps({"token": char}))
            
            # Flush to trigger synthesis
            await websocket.send(json.dumps({"flush": True}))
            
            # Receive all audio chunks
            chunk_count = 0
            total_bytes = 0
            first_chunk_time = None
            
            try:
                while True:
                    audio_data = await asyncio.wait_for(websocket.recv(), timeout=10.0)
                    chunk_count += 1
                    total_bytes += len(audio_data)
                    
                    if first_chunk_time is None:
                        first_chunk_time = time.time() - client_start
                        
            except asyncio.TimeoutError:
                client_end = time.time()
                client_total = client_end - client_start
                
                # Calculate audio duration from received bytes
                # Format: PCM float32, 48kHz, mono
                samples = total_bytes / 4
                audio_duration = samples / 48000
                client_rtf = audio_duration / client_total if client_total > 0 else 0
                
                # Get server-side metrics from logs
                result = subprocess.run(
                    ["docker", "logs", "miku-rvc-api", "--tail", "50"],
                    capture_output=True,
                    text=True
                )
                
                # Parse the last job complete log
                pattern = r"Job complete: (\d+) blocks in ([\d.]+)s.*?Audio duration: ([\d.]+)s.*?Realtime factor: ([\d.]+)x.*?Avg Soprano: ([\d.]+)ms.*?Avg RVC: ([\d.]+)ms"
                matches = re.findall(pattern, result.stderr + result.stdout, re.DOTALL)
                
                if matches:
                    server_blocks, server_time, server_audio, server_rtf, avg_soprano, avg_rvc = matches[-1]
                    
                    print(f"\n📊 PERFORMANCE METRICS:")
                    print(f"\n  🎤 SOPRANO (Text-to-Speech):")
                    print(f"     └─ Generation time: {avg_soprano} ms")
                    
                    print(f"\n  🎙️  RVC (Voice Conversion):")
                    print(f"     ├─ Blocks processed: {server_blocks}")
                    print(f"     └─ Avg per block: {avg_rvc} ms")
                    
                    print(f"\n  ⚡ PIPELINE:")
                    print(f"     ├─ Audio generated: {server_audio}s")
                    print(f"     ├─ Processing time: {server_time}s")
                    print(f"     └─ Realtime factor: {server_rtf}x")
                    
                    print(f"\n  🌐 CLIENT (End-to-end):")
                    print(f"     ├─ Chunks received: {chunk_count}")
                    print(f"     ├─ Total bytes: {total_bytes:,} ({total_bytes/1024:.1f} KB)")
                    print(f"     ├─ Time to first chunk: {first_chunk_time:.2f}s (latency)")
                    print(f"     ├─ Total time: {client_total:.2f}s")
                    print(f"     └─ Realtime factor: {client_rtf:.2f}x")
                    
                    # Performance assessment
                    server_rtf_val = float(server_rtf)
                    print(f"\n  💡 ASSESSMENT:")
                    if server_rtf_val >= 1.0:
                        print(f"     ✅ EXCELLENT - Faster than realtime ({server_rtf}x)")
                    elif server_rtf_val >= 0.7:
                        print(f"     ✅ GOOD - Near realtime ({server_rtf}x)")
                    elif server_rtf_val >= 0.5:
                        print(f"     ⚠️  ACCEPTABLE - Moderate speed ({server_rtf}x)")
                    else:
                        print(f"     ❌ SLOW - Needs optimization ({server_rtf}x)")
                else:
                    print(f"\n  ⚠️  Could not parse server metrics from logs")
                    print(f"\n  📊 CLIENT METRICS:")
                    print(f"     ├─ Chunks: {chunk_count}")
                    print(f"     ├─ Audio: {audio_duration:.2f}s")
                    print(f"     ├─ Total time: {client_total:.2f}s")
                    print(f"     └─ RTF: {client_rtf:.2f}x")
    
    print(f"\n{'='*80}")
    print(" "*30 + "TEST COMPLETE")
    print("="*80 + "\n")

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