#!/usr/bin/env python3
"""
Test script for MuseTalk WebRTC demo mode
Runs headless Chromium and captures metrics
"""

import asyncio
import json
from playwright.async_api import async_playwright

async def run_demo_test(url: str, timeout: int = 120000):
    """Run demo test and return metrics"""

    async with async_playwright() as p:
        # Launch headless browser
        browser = await p.chromium.launch(headless=True)
        context = await browser.new_context(
            permissions=['microphone']
        )
        page = await context.new_page()

        # Capture console logs
        logs = []
        page.on('console', lambda msg: logs.append(f"[{msg.type}] {msg.text}"))

        print(f"[Test] Opening {url}")

        try:
            await page.goto(url, timeout=30000)
            print("[Test] Page loaded, waiting for demo to complete...")

            # Wait for metrics (up to 2 minutes)
            try:
                await page.wait_for_selector('#demo-metrics-output', timeout=timeout)
                print("[Test] Metrics element found!")

                metrics = await page.evaluate('() => window.__DEMO_METRICS__')

                if metrics:
                    def fmt(v):
                        return f"{v:>6}" if v is not None else "  N/A "

                    # MODEL INFORMATION
                    print("\n📦 MODELS USED:")
                    print("┌────────────────────────────────────────────────────────────┐")
                    print("│  Component          Model                        Size     │")
                    print("├────────────────────────────────────────────────────────────┤")
                    print("│  STT                Whisper distil-large-v3-ptbr ~1.5 GB  │")
                    print("│  LLM                google/gemma-3-1b-it         ~2.0 GB  │")
                    print("│  TTS                espeak-ng / ElevenLabs       system   │")
                    print("│  MuseTalk UNet      UNet2DConditionModel         3.4 GB   │")
                    print("│  MuseTalk VAE       AutoencoderKL (sd-vae)       335 MB   │")
                    print("│  Face Detection     DWPose                       407 MB   │")
                    print("│  Face Parser        BiSeNet                      53 MB    │")
                    print("└────────────────────────────────────────────────────────────┘")

                    # AVALIAÇÃO PRINCIPAL
                    ev = metrics.get('evaluation', {})
                    response_lat = ev.get('responseLatency')
                    ratio = ev.get('latencyRatio')
                    grade = ev.get('latencyGrade', 'N/A')

                    print("\n" + "╔" + "═"*60 + "╗")
                    print("║" + " "*10 + "LATÊNCIA DE RESPOSTA" + " "*30 + "║")
                    print("║" + " "*10 + "(áudio enviado → resposta começa)" + " "*17 + "║")
                    print("╠" + "═"*60 + "╣")
                    print(f"║" + " "*20 + f">>> {response_lat} ms <<<" + " "*22 + "║")
                    print("║" + " "*60 + "║")
                    print(f"║   Referência humana:  250 ms" + " "*31 + "║")
                    print(f"║   Razão:              {ratio}x mais lento" + " "*25 + "║")
                    print(f"║   Nota:               {grade}" + " "*37 + "║")
                    print("╚" + "═"*60 + "╝")

                    # Latências detalhadas - calcular tempos individuais do servidor
                    lat = metrics.get('latency', {})
                    sm = metrics.get('serverMetrics', {})
                    t = sm.get('timings', {}) if sm else {}

                    # Tempos individuais (não acumulados)
                    stt_time = t.get('stt_done', 0)
                    llm_time = t.get('llm_done', 0) - t.get('stt_done', 0) if t.get('llm_done') else 0

                    # TTS separado: espeak vs ElevenLabs
                    espeak_time = t.get('espeak_done', 0) - t.get('llm_done', 0) if t.get('espeak_done') else 0
                    tts_time = t.get('tts_done', 0) - t.get('llm_done', 0) if t.get('tts_done') else 0
                    elevenlabs_time = t.get('elevenlabs_done', 0) - t.get('llm_done', 0) if t.get('elevenlabs_done') else 0

                    musetalk_first = t.get('first_frame', 0) - t.get('tts_done', 0) if t.get('first_frame') else 0
                    musetalk_total = t.get('all_done', 0) - t.get('tts_done', 0) if t.get('all_done') else 0

                    print("\n📊 BREAKDOWN DA LATÊNCIA (tempos individuais):")
                    print(f"   1. STT (Whisper):       {stt_time:>6} ms")
                    print(f"   2. LLM (vLLM/local):    {llm_time:>6} ms")
                    if espeak_time:
                        print(f"   3. TTS espeak:          {espeak_time:>6} ms  ← usado para vídeo")
                        print(f"      TTS ElevenLabs:      {elevenlabs_time:>6} ms  ← paralelo (áudio final)")
                    else:
                        print(f"   3. TTS:                 {tts_time:>6} ms")
                    print(f"   4. MuseTalk 1ºFrame:    {musetalk_first:>6} ms")
                    print(f"   ─────────────────────────────────")
                    print(f"   ⭐ Latência Resposta:    {t.get('first_frame', 0):>6} ms")
                    print(f"   ")
                    print(f"   5. MuseTalk Total:      {musetalk_total:>6} ms")
                    print(f"   ─────────────────────────────────")
                    print(f"   Total Pipeline:         {t.get('all_done', 0):>6} ms")

                    # Qualidade
                    vq = metrics.get('videoQuality', {})
                    print(f"\n🎬 QUALIDADE DO VÍDEO:")
                    print(f"   Frames:            {vq.get('totalFrames', 0)}/{vq.get('expectedFrames', 0)}")
                    print(f"   Frames perdidos:   {vq.get('droppedFrames', 0)} ({vq.get('dropRate', 'N/A')})")
                    print(f"   Travamentos:       {vq.get('stutterCount', 0)}")
                    print(f"   Max gap:           {vq.get('maxFrameGap', 0)} ms")
                    print(f"   Avg gap:           {vq.get('avgFrameGap', 0)} ms")
                    smooth = "✓ SIM" if vq.get('isSmooth') else "✗ NÃO"
                    print(f"   Fluido:            {smooth}")

                    # Sincronização
                    sync = metrics.get('sync', {})
                    print(f"\n🔊 SINCRONIZAÇÃO ÁUDIO/VÍDEO:")
                    print(f"   Offset A/V:        {sync.get('audioVideoOffset', 'N/A')} ms")
                    synced = "✓ SIM" if sync.get('isSynced') else "✗ NÃO"
                    print(f"   Sincronizado:      {synced}")

                    # Server metrics
                    sm = metrics.get('serverMetrics', {})
                    if sm and sm.get('timings'):
                        print(f"\n🖥️  SERVIDOR (tempos acumulados):")
                        t = sm['timings']
                        stt = t.get('stt_done', 0)
                        llm = t.get('llm_done', 0)
                        espeak = t.get('espeak_done', 0)
                        tts = t.get('tts_done', 0)
                        elevenlabs = t.get('elevenlabs_done', 0)
                        first_frame = t.get('first_frame', 0)
                        all_done = t.get('all_done', 0)

                        print(f"   STT:               {stt:>6} ms")
                        print(f"   LLM:               {llm:>6} ms  (+{llm - stt} ms)")
                        if espeak:
                            print(f"   espeak:            {espeak:>6} ms  (+{espeak - llm} ms)")
                            print(f"   TTS ready:         {tts:>6} ms  (+{tts - espeak} ms)")
                        else:
                            print(f"   TTS:               {tts:>6} ms  (+{tts - llm} ms)")
                        print(f"   First Frame:       {first_frame:>6} ms")
                        if elevenlabs:
                            print(f"   ElevenLabs:        {elevenlabs:>6} ms  (paralelo)")
                        print(f"   Total:             {all_done:>6} ms")

                        # MuseTalk specific timing
                        if tts and first_frame and all_done:
                            musetalk_first = first_frame - tts
                            musetalk_total = all_done - tts
                            total_frames = vq.get('totalFrames', 0)

                            print(f"\n🎭 MUSETALK (UNet2D + VAE inference):")
                            print(f"   Models:            UNet2DConditionModel (3.4GB) + AutoencoderKL (335MB)")
                            print(f"   1º Frame:          {musetalk_first:>6} ms")
                            print(f"   Total:             {musetalk_total:>6} ms")
                            if total_frames and musetalk_total > 0:
                                fps = total_frames / (musetalk_total / 1000)
                                # Calculate realtime factor (assume 25fps output)
                                audio_duration_sec = total_frames / 25.0
                                realtime_factor = (musetalk_total / 1000) / audio_duration_sec if audio_duration_sec > 0 else 0
                                
                                print(f"   Frames gerados:    {total_frames:>6}")
                                print(f"   Audio duration:    {audio_duration_sec:>6.2f} s")
                                print(f"   Gen FPS:           {fps:>6.1f} fps")
                                print(f"   Realtime factor:   {realtime_factor:>6.2f}x")

                                # Avaliação MuseTalk
                                if realtime_factor <= 1.0:
                                    mt_grade = f"A+ ({1/realtime_factor:.1f}x faster than realtime)"
                                elif fps >= 25:
                                    mt_grade = "A (tempo real)"
                                elif fps >= 15:
                                    mt_grade = "B (aceitável)"
                                elif fps >= 10:
                                    mt_grade = "C (lento)"
                                else:
                                    mt_grade = "F (muito lento)"
                                print(f"   Avaliação:         {mt_grade}")

                    # Resultado final
                    print("\n" + "="*60)
                    fluid = metrics.get('conversationFluid', False)
                    if fluid:
                        print("✅ CONVERSA FLUIDA: SIM")
                    else:
                        print("❌ CONVERSA FLUIDA: NÃO")

                    if metrics.get('error'):
                        print(f"⚠️  ERRO: {metrics['error']}")

                    print("="*60)

                    # Print full JSON
                    print("\n📋 JSON Completo:")
                    print(json.dumps(metrics, indent=2))

                    return metrics

            except Exception as e:
                print(f"[Test] Timeout: {e}")
                metrics = await page.evaluate('() => window.__DEMO_METRICS__ || null')
                if metrics:
                    print("\n[Test] Métricas parciais:")
                    print(json.dumps(metrics, indent=2))
                    return metrics

                print("\n[Test] Console logs:")
                for log in logs[-30:]:
                    print(f"  {log}")

                return {"error": str(e), "logs": logs[-30:]}

        except Exception as e:
            print(f"[Test] Error: {e}")
            return {"error": str(e)}

        finally:
            await browser.close()

if __name__ == "__main__":
    import sys

    url = sys.argv[1] if len(sys.argv) > 1 else "http://localhost:3000/?demo=true"

    print(f"[Test] Running demo test on: {url}")
    result = asyncio.run(run_demo_test(url))

    if result and result.get('error'):
        sys.exit(1)