import grpc
import time
import sys
import wave
from concurrent.futures import ThreadPoolExecutor, as_completed

sys.path.insert(0, 'server/grpc')
import avatar_pb2
import avatar_pb2_grpc

# Carregar audio de teste
audio_path = 'data/audio/orpheus_demo_3_phrases.wav'
with wave.open(audio_path, 'rb') as wf:
    sample_rate = wf.getframerate()
    audio_data = wf.readframes(wf.getnframes())

audio_duration = len(audio_data) / (sample_rate * 2)
print(f'Audio: {audio_duration:.2f}s')
print()

def run_session(session_id, stub):
    def audio_chunks():
        chunk_size = 4800  # 200ms chunks
        for i in range(0, len(audio_data), chunk_size):
            yield avatar_pb2.AudioChunk(
                session_id=f'user_{session_id}',
                audio_data=audio_data[i:i+chunk_size],
                sample_rate=sample_rate,
                is_final=(i + chunk_size >= len(audio_data))
            )
    
    start = time.time()
    first_frame_time = None
    frame_count = 0
    
    try:
        for response in stub.StreamingGenerate(audio_chunks()):
            if first_frame_time is None:
                first_frame_time = time.time() - start
            frame_count += 1
    except Exception as e:
        return {'error': str(e), 'session_id': session_id}
    
    total_time = time.time() - start
    fps = frame_count / total_time if total_time > 0 else 0
    
    return {
        'session_id': session_id,
        'ttff': first_frame_time * 1000 if first_frame_time else 0,
        'frames': frame_count,
        'fps': fps,
        'total_time': total_time
    }

def test_concurrent_users(num_users):
    channel = grpc.insecure_channel('localhost:8080')
    stub = avatar_pb2_grpc.AvatarServiceStub(channel)
    
    results = []
    with ThreadPoolExecutor(max_workers=num_users) as executor:
        futures = [executor.submit(run_session, i, stub) for i in range(num_users)]
        for future in as_completed(futures):
            results.append(future.result())
    
    return results

# Testar de 1 a 8 usuarios
print('=' * 70)
print('TESTE DE CAPACIDADE MULTI-USUARIO (torch.compile ENABLED)')
print('=' * 70)
print(f'Users | Avg TTFF | Avg FPS | Min FPS | Real-time?')
print('-' * 70)

realtime_fps = 25

for num_users in [1, 2, 3, 4, 5, 6, 8]:
    results = test_concurrent_users(num_users)
    valid = [r for r in results if 'error' not in r]
    
    if valid:
        avg_ttff = sum(r['ttff'] for r in valid) / len(valid)
        avg_fps = sum(r['fps'] for r in valid) / len(valid)
        min_fps = min(r['fps'] for r in valid)
        is_realtime = 'YES' if min_fps >= realtime_fps else 'NO'
        
        print(f'{num_users:>5} | {avg_ttff:>8.0f}ms | {avg_fps:>7.1f} | {min_fps:>7.1f} | {is_realtime}')
    else:
        print(f'{num_users:>5} | ERROR')
    
    time.sleep(2)

print('=' * 70)
print(f'Real-time threshold: {realtime_fps} FPS (video is 25 FPS)')
print('=' * 70)