#!/usr/bin/env python3
"""
Cliente WebRTC automatizado para testar conexão com servidor
"""

import asyncio
import aiohttp
import sys
from aiortc import RTCPeerConnection, RTCConfiguration, RTCIceServer, RTCSessionDescription
from aiortc.contrib.media import MediaRecorder
import time

# Configuração
SERVER_URL = "http://62.107.25.198:47898"

# Cores para terminal
GREEN = '\033[92m'
RED = '\033[91m'
YELLOW = '\033[93m'
BLUE = '\033[94m'
RESET = '\033[0m'

class WebRTCTester:
    def __init__(self):
        self.pc = None
        self.session_id = None
        self.video_frames = 0
        self.audio_frames = 0
        self.ice_candidates = []
        self.connected = False

    async def test_connection(self):
        """Testa conexão WebRTC completa"""

        print(f"\n{'='*60}")
        print(f"{BLUE}WebRTC Connection Test{RESET}")
        print(f"{'='*60}\n")

        try:
            # 1. Configurar ICE servers (STUN + TURN)
            print(f"{YELLOW}[1/5]{RESET} Configurando ICE servers...")
            ice_servers = [
                RTCIceServer(urls=["stun:stun.l.google.com:19302"]),
                RTCIceServer(urls=["stun:stun1.l.google.com:19302"]),
                RTCIceServer(
                    urls=["turn:openrelay.metered.ca:80"],
                    username="openrelayproject",
                    credential="openrelayproject"
                ),
                RTCIceServer(
                    urls=["turn:openrelay.metered.ca:443"],
                    username="openrelayproject",
                    credential="openrelayproject"
                ),
            ]

            config = RTCConfiguration(iceServers=ice_servers)
            self.pc = RTCPeerConnection(configuration=config)
            print(f"{GREEN}✓{RESET} ICE servers configurados\n")

            # 2. Configurar handlers
            print(f"{YELLOW}[2/5]{RESET} Configurando event handlers...")

            @self.pc.on("icecandidate")
            async def on_ice_candidate(candidate):
                if candidate:
                    self.ice_candidates.append({
                        'type': candidate.candidate.type if hasattr(candidate.candidate, 'type') else 'unknown',
                        'protocol': candidate.candidate.protocol if hasattr(candidate.candidate, 'protocol') else 'unknown'
                    })

            @self.pc.on("connectionstatechange")
            async def on_connection_state():
                state = self.pc.connectionState
                if state == "connected":
                    print(f"{GREEN}✓{RESET} WebRTC Estado: {GREEN}CONNECTED{RESET}")
                    self.connected = True
                elif state == "failed":
                    print(f"{RED}✗{RESET} WebRTC Estado: {RED}FAILED{RESET}")
                elif state == "connecting":
                    print(f"{YELLOW}⟳{RESET} WebRTC Estado: CONNECTING...")
                else:
                    print(f"  WebRTC Estado: {state}")

            @self.pc.on("iceconnectionstatechange")
            async def on_ice_state():
                state = self.pc.iceConnectionState
                if state == "connected":
                    print(f"{GREEN}✓{RESET} ICE Estado: {GREEN}CONNECTED{RESET}")
                elif state == "failed":
                    print(f"{RED}✗{RESET} ICE Estado: {RED}FAILED{RESET}")
                elif state == "checking":
                    print(f"{YELLOW}⟳{RESET} ICE Estado: CHECKING...")
                else:
                    print(f"  ICE Estado: {state}")

            @self.pc.on("track")
            async def on_track(track):
                print(f"{GREEN}✓{RESET} Track recebido: {BLUE}{track.kind}{RESET}")

                if track.kind == "video":
                    while True:
                        try:
                            frame = await track.recv()
                            self.video_frames += 1
                            if self.video_frames % 25 == 0:  # A cada segundo (25fps)
                                print(f"  {BLUE}Video:{RESET} {self.video_frames} frames recebidos")
                        except Exception as e:
                            break

                elif track.kind == "audio":
                    while True:
                        try:
                            frame = await track.recv()
                            self.audio_frames += 1
                            if self.audio_frames % 50 == 0:  # A cada ~1 segundo
                                print(f"  {BLUE}Audio:{RESET} {self.audio_frames} frames recebidos")
                        except Exception as e:
                            break

            print(f"{GREEN}✓{RESET} Handlers configurados\n")

            # 3. Criar transceivers
            print(f"{YELLOW}[3/5]{RESET} Criando transceivers...")
            self.pc.addTransceiver("video", direction="recvonly")
            self.pc.addTransceiver("audio", direction="recvonly")
            print(f"{GREEN}✓{RESET} Transceivers criados\n")

            # 4. Criar offer e enviar para servidor
            print(f"{YELLOW}[4/5]{RESET} Criando offer SDP...")
            offer = await self.pc.createOffer()
            await self.pc.setLocalDescription(offer)

            # Aguardar ICE gathering
            print(f"{YELLOW}⟳{RESET} Aguardando ICE gathering...")
            await self.wait_for_ice_gathering()

            print(f"{GREEN}✓{RESET} Offer criado\n")

            # 5. Enviar offer para servidor
            print(f"{YELLOW}[5/5]{RESET} Enviando offer para servidor...")
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    f"{SERVER_URL}/offer",
                    json={
                        "sdp": self.pc.localDescription.sdp,
                        "type": self.pc.localDescription.type
                    }
                ) as resp:
                    if resp.status != 200:
                        raise Exception(f"Erro HTTP {resp.status}")

                    answer = await resp.json()
                    self.session_id = answer.get('session_id')

                    # Aplicar answer
                    await self.pc.setRemoteDescription(
                        RTCSessionDescription(
                            sdp=answer["sdp"],
                            type=answer["type"]
                        )
                    )

            print(f"{GREEN}✓{RESET} Answer recebido")
            print(f"{GREEN}✓{RESET} Session ID: {BLUE}{self.session_id}{RESET}\n")

            # 6. Aguardar conexão
            print(f"{YELLOW}⟳{RESET} Aguardando conexão WebRTC...\n")

            # Aguardar até 15 segundos
            for i in range(15):
                await asyncio.sleep(1)
                if self.connected:
                    break

            # 7. Mostrar resultados
            await self.show_results()

            # 8. Manter vivo por 10 segundos para receber frames
            if self.connected:
                print(f"\n{YELLOW}⟳{RESET} Monitorando recepção de frames por 10 segundos...\n")
                await asyncio.sleep(10)

            # 9. Resultados finais
            await self.show_final_results()

        except Exception as e:
            print(f"\n{RED}✗ ERRO:{RESET} {e}\n")
            import traceback
            traceback.print_exc()
            return False

        finally:
            if self.pc:
                await self.pc.close()

        return self.connected

    async def wait_for_ice_gathering(self):
        """Aguarda ICE gathering completar"""
        max_wait = 5
        for _ in range(max_wait * 10):
            if self.pc.iceGatheringState == "complete":
                return
            await asyncio.sleep(0.1)

    async def show_results(self):
        """Mostra resultados da conexão"""
        print(f"\n{'='*60}")
        print(f"{BLUE}Resultados da Conexão{RESET}")
        print(f"{'='*60}\n")

        # Estado da conexão
        conn_state = self.pc.connectionState
        if conn_state == "connected":
            print(f"WebRTC State:     {GREEN}✓ CONNECTED{RESET}")
        elif conn_state == "failed":
            print(f"WebRTC State:     {RED}✗ FAILED{RESET}")
        else:
            print(f"WebRTC State:     {YELLOW}{conn_state.upper()}{RESET}")

        # Estado ICE
        ice_state = self.pc.iceConnectionState
        if ice_state == "connected":
            print(f"ICE State:        {GREEN}✓ CONNECTED{RESET}")
        elif ice_state == "failed":
            print(f"ICE State:        {RED}✗ FAILED{RESET}")
        else:
            print(f"ICE State:        {YELLOW}{ice_state.upper()}{RESET}")

        # Candidatos ICE
        print(f"\nICE Candidates:   {len(self.ice_candidates)} gerados")

        candidate_types = {}
        for c in self.ice_candidates:
            ctype = c['type']
            candidate_types[ctype] = candidate_types.get(ctype, 0) + 1

        for ctype, count in candidate_types.items():
            icon = "✓" if ctype == "relay" else "•"
            color = GREEN if ctype == "relay" else BLUE
            print(f"  {color}{icon}{RESET} {ctype}: {count}")

        # Session ID
        if self.session_id:
            print(f"\nSession ID:       {BLUE}{self.session_id}{RESET}")

        print()

    async def show_final_results(self):
        """Mostra resultados finais"""
        print(f"\n{'='*60}")
        print(f"{BLUE}Resultados Finais{RESET}")
        print(f"{'='*60}\n")

        if self.connected:
            print(f"Status:           {GREEN}✓ SUCESSO{RESET}")
        else:
            print(f"Status:           {RED}✗ FALHA{RESET}")

        print(f"Video Frames:     {BLUE}{self.video_frames}{RESET}")
        print(f"Audio Frames:     {BLUE}{self.audio_frames}{RESET}")

        if self.video_frames > 0 and self.audio_frames > 0:
            print(f"\n{GREEN}✓ WebRTC funcionando perfeitamente!{RESET}")
        elif self.connected:
            print(f"\n{YELLOW}⚠ Conectado mas sem receber frames{RESET}")
        else:
            print(f"\n{RED}✗ Falha na conexão WebRTC{RESET}")

        print()


async def main():
    tester = WebRTCTester()
    success = await tester.test_connection()
    sys.exit(0 if success else 1)


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