"""
Debug script for testing WebSocket streaming functionality
"""
import asyncio
import websockets
import json
import time

async def debug_websocket():
    """Test the WebSocket connection and message streaming"""
    
    # Test configuration
    test_config = {
        "subiect": "Test streaming functionality",
        "numar_runde": 2,
        "pauza_intre_runde": 1,
        "agenti": [
            {
                "ip": "127.0.0.1",
                "port": "11434",
                "rol": "generator",
                "model": "llama2",
                "max_tokens": 100,
                "timeout": 30
            },
            {
                "ip": "127.0.0.1",
                "port": "11434",
                "rol": "critic",
                "model": "llama2", 
                "max_tokens": 100,
                "timeout": 30
            }
        ]
    }
    
    client_id = f"debug-{int(time.time())}"
    ws_url = f"ws://localhost:8002/ws/{client_id}"
    
    try:
        print(f"Connecting to WebSocket: {ws_url}")
        
        async with websockets.connect(ws_url) as websocket:
            print("✅ WebSocket connection established")
            
            # Send start command
            start_message = {
                "action": "start",
                "config": test_config
            }
            
            print("📤 Sending start command...")
            await websocket.send(json.dumps(start_message))
            print("✅ Start command sent")
            
            # Listen for messages
            message_count = 0
            token_count = 0
            
            async for message in websocket:
                try:
                    data = json.loads(message)
                    message_type = data.get("type", "unknown")
                    message_data = data.get("data", {})
                    
                    message_count += 1
                    
                    print(f"\n📬 Message #{message_count}: {message_type}")
                    
                    if message_type == "started":
                        print(f"   ✅ Conversation started: {message_data.get('conversation_id')}")
                    
                    elif message_type == "round_start":
                        print(f"   🏁 Round {message_data.get('round')}/{message_data.get('total_rounds')} started")
                    
                    elif message_type == "token":
                        token_count += 1
                        token = message_data.get("token", "")
                        agent_role = message_data.get("agent_role", "")
                        print(f"   🔤 Token #{token_count} from {agent_role}: '{token}'")
                    
                    elif message_type == "message":
                        content = message_data.get("content", "")[:50] + "..."
                        agent_role = message_data.get("agent_role", "")
                        tokens = message_data.get("tokens", "N/A")
                        print(f"   💬 Message from {agent_role} ({tokens} tokens): {content}")
                    
                    elif message_type == "round_end":
                        print(f"   🏁 Round ended. Total ideas: {message_data.get('total_ideas')}")
                    
                    elif message_type == "conversation_end":
                        print(f"   🎉 Conversation completed!")
                        print(f"      Rounds: {message_data.get('total_rounds')}")
                        print(f"      Messages: {message_data.get('total_messages')}")
                        print(f"      Ideas: {message_data.get('total_ideas')}")
                        break
                    
                    elif message_type == "error":
                        print(f"   ❌ Error: {message_data.get('message')}")
                        if "Connection" in str(message_data.get('message', '')):
                            print("   💡 This might be expected if Ollama is not running")
                        break
                    
                    elif message_type == "saved":
                        print(f"   💾 Conversation saved: {message_data.get('conversation_id')}")
                        break
                        
                    else:
                        print(f"   📋 Data: {message_data}")
                    
                    # Stop after 50 messages to avoid infinite loop
                    if message_count > 50:
                        print("⚠️  Stopping due to message limit")
                        break
                        
                except json.JSONDecodeError as e:
                    print(f"❌ JSON decode error: {e}")
                    print(f"Raw message: {message}")
                except Exception as e:
                    print(f"❌ Error processing message: {e}")
            
            print(f"\n📊 Summary:")
            print(f"   Total messages: {message_count}")
            print(f"   Token messages: {token_count}")
            
    except websockets.exceptions.ConnectionClosed as e:
        print(f"❌ WebSocket connection closed: {e}")
    except Exception as e:
        print(f"❌ WebSocket error: {e}")

if __name__ == "__main__":
    print("🚀 Starting WebSocket debug test...")
    asyncio.run(debug_websocket())
    print("🏁 Debug test completed")