#!/usr/bin/env python3
"""
PROPER GPU Contention Test v2
Tests the ACTUAL async pipeline vs isolated components to find the real bottleneck
"""
import time
import requests
import subprocess
from statistics import mean

def test_soprano_isolated(num_runs=3):
    """Test Soprano isolated endpoint"""
    print("=" * 70)
    print("TEST 1: Soprano ISOLATED (baseline)")
    print("=" * 70)
    
    text = "This is a comprehensive test of raw Soprano TTS streaming performance without any RVC processing to isolate the true bottleneck and measure generation speed."
    times = []
    audio_durations = []
    
    for i in range(num_runs):
        print(f"\nRun {i+1}/{num_runs}...")
        start = time.time()
        
        # Download the audio file
        response = requests.get(
            f"http://localhost:8765/api/test_soprano_stream",
            params={"text": text},
            stream=True
        )
        
        output_file = f"/tmp/soprano_isolated_{i}.wav"
        with open(output_file, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        
        elapsed = time.time() - start
        times.append(elapsed)
        
        # Get audio duration using ffprobe
        result = subprocess.run(
            ['ffprobe', '-v', 'error', '-show_entries', 'format=duration', 
             '-of', 'default=noprint_wrappers=1:nokey=1', output_file],
            capture_output=True, text=True
        )
        duration = float(result.stdout.strip())
        audio_durations.append(duration)
        
        realtime = duration / elapsed
        print(f"  Elapsed: {elapsed:.2f}s, Audio: {duration:.2f}s, Realtime: {realtime:.2f}x")
        time.sleep(1)
    
    avg_elapsed = mean(times)
    avg_duration = mean(audio_durations)
    avg_realtime = avg_duration / avg_elapsed
    
    print(f"\n✓ Soprano Isolated Average:")
    print(f"  Elapsed: {avg_elapsed:.2f}s")
    print(f"  Audio: {avg_duration:.2f}s")
    print(f"  Realtime: {avg_realtime:.2f}x")
    
    return avg_elapsed, avg_duration, avg_realtime

def test_async_pipeline(num_runs=3):
    """Test the ACTUAL async pipeline (Soprano + RVC integrated)"""
    print("\n" + "=" * 70)
    print("TEST 2: ASYNC PIPELINE (Soprano + RVC integrated)")
    print("=" * 70)
    
    text = "This is a comprehensive test of raw Soprano TTS streaming performance without any RVC processing to isolate the true bottleneck and measure generation speed."
    times = []
    audio_durations = []
    
    # First, submit job and connect to stream
    for i in range(num_runs):
        print(f"\nRun {i+1}/{num_runs}...")
        
        # Submit the job
        job_response = requests.post(
            "http://localhost:8765/api/speak",
            json={"text": text}
        )
        job_data = job_response.json()
        print(f"  Job queued: {job_data.get('job_id', 'unknown')[:8]}")
        
        # Connect to streaming endpoint and consume audio
        start = time.time()
        response = requests.get(
            "http://localhost:8765/api/stream/continuous",
            stream=True
        )
        
        output_file = f"/tmp/async_pipeline_{i}.wav"
        bytes_received = 0
        with open(output_file, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                if chunk:
                    f.write(chunk)
                    bytes_received += len(chunk)
                    # Stop after receiving enough data (roughly one job's worth)
                    # Assuming 48kHz 16-bit mono: ~96KB per second of audio
                    # For ~6s audio, that's ~576KB
                    if bytes_received > 600000:  # 600KB threshold
                        break
        
        elapsed = time.time() - start
        times.append(elapsed)
        
        # Close the connection
        response.close()
        
        # Get audio duration
        result = subprocess.run(
            ['ffprobe', '-v', 'error', '-show_entries', 'format=duration',
             '-of', 'default=noprint_wrappers=1:nokey=1', output_file],
            capture_output=True, text=True
        )
        duration = float(result.stdout.strip()) if result.stdout.strip() else 0
        audio_durations.append(duration)
        
        realtime = duration / elapsed if elapsed > 0 else 0
        print(f"  Elapsed: {elapsed:.2f}s, Audio: {duration:.2f}s, Realtime: {realtime:.2f}x")
        print(f"  Bytes received: {bytes_received}")
        time.sleep(2)
    
    avg_elapsed = mean(times)
    avg_duration = mean(audio_durations)
    avg_realtime = avg_duration / avg_elapsed if avg_elapsed > 0 else 0
    
    print(f"\n✓ Async Pipeline Average:")
    print(f"  Elapsed: {avg_elapsed:.2f}s")
    print(f"  Audio: {avg_duration:.2f}s")
    print(f"  Realtime: {avg_realtime:.2f}x")
    
    return avg_elapsed, avg_duration, avg_realtime

def main():
    print("\n" + "=" * 70)
    print("PROPER GPU CONTENTION TEST v2")
    print("=" * 70)
    print("\nThis test compares:")
    print("1. Soprano isolated (no RVC)")
    print("2. Full async pipeline (Soprano + RVC integrated)")
    print("\nIf GPU contention is the issue:")
    print("- Soprano should be slower in the pipeline due to RVC competing")
    print("- Performance loss should correlate with GPU memory pressure")
    print("\nStarting in 3 seconds...")
    time.sleep(3)
    
    # Run tests
    soprano_elapsed, soprano_duration, soprano_realtime = test_soprano_isolated(num_runs=3)
    pipeline_elapsed, pipeline_duration, pipeline_realtime = test_async_pipeline(num_runs=3)
    
    # Analysis
    print("\n" + "=" * 70)
    print("ANALYSIS")
    print("=" * 70)
    
    # Calculate how much slower the pipeline is
    slowdown_factor = soprano_realtime / pipeline_realtime if pipeline_realtime > 0 else float('inf')
    slowdown_percent = ((soprano_realtime - pipeline_realtime) / soprano_realtime) * 100 if soprano_realtime > 0 else 0
    
    print(f"\nSoprano Isolated:  {soprano_realtime:.2f}x realtime")
    print(f"Async Pipeline:    {pipeline_realtime:.2f}x realtime")
    print(f"Performance Loss:  {slowdown_percent:.1f}%")
    print(f"Slowdown Factor:   {slowdown_factor:.2f}x")
    
    print("\n" + "=" * 70)
    print("CONCLUSION")
    print("=" * 70)
    
    if slowdown_percent > 40:
        print("⚠️  MAJOR BOTTLENECK DETECTED!")
        print(f"\n   The pipeline runs at {slowdown_percent:.1f}% slower than Soprano alone.")
        print("\n   Possible causes:")
        print("   1. GPU contention (RVC competing with Soprano)")
        print("   2. Memory bandwidth limitations")  
        print("   3. Pipeline overhead (queue, threading)")
        print("   4. CPU bottleneck (PCIe lanes, context switching)")
        
        # Check logs for timing breakdown
        print("\n   Recommendation: Check server logs for detailed timing breakdown")
        print("   Look for '[Soprano]' and '[RVC]' timing to see where time is spent")
        
    elif slowdown_percent > 20:
        print("⚠️  MODERATE PERFORMANCE LOSS")
        print(f"\n   The pipeline runs at {slowdown_percent:.1f}% slower than Soprano alone.")
        print("   This is expected overhead from RVC processing")
        
    elif slowdown_percent > 10:
        print("✓ MINOR OVERHEAD (ACCEPTABLE)")
        print(f"\n   The pipeline runs at {slowdown_percent:.1f}% slower than Soprano alone.")
        print("   This is normal for integrated pipeline processing")
        
    else:
        print("✓ EXCELLENT PERFORMANCE!")
        print(f"\n   Pipeline performance: {pipeline_realtime:.2f}x realtime")
        print("   Minimal overhead from integration")
    
    # Target assessment
    print(f"\n   Target: >1.0x realtime")
    if pipeline_realtime >= 1.0:
        print(f"   Status: ✓ ACHIEVED ({pipeline_realtime:.2f}x)")
    else:
        needed_improvement = (1.0 / pipeline_realtime - 1.0) * 100
        print(f"   Status: ✗ MISSED (need {needed_improvement:.1f}% improvement)")
    
    print("\n" + "=" * 70)

if __name__ == "__main__":
    main()
