import asyncio
import json
import websockets
import sys

async def run_integration_test():
    url = "ws://localhost:3000/ws"
    
    print("🤖 Test: Starting Integration Test Suite...")
    
    # Client 1: Creator
    try:
        creator_ws = await websockets.connect(url)
        print("✅ Creator connected to WebSocket")
    except Exception as e:
        print(f"❌ Failed to connect to server: {e}")
        sys.exit(1)
        
    creator_msg = {
        "type": "create_room",
        "username": "TestCreator",
        "fileInfo": {
            "name": "test_video.mkv",
            "size": 1048576, # 1MB
            "duration": 0
        }
    }
    await creator_ws.send(json.dumps(creator_msg))
    print("📤 Creator sent 'create_room'")
    
    room_code = None
    
    # Wait for room_created
    try:
        response_str = await asyncio.wait_for(creator_ws.recv(), timeout=2.0)
        response = json.loads(response_str)
        if response.get("type") == "room_created":
            room_code = response.get("code")
            print(f"✅ Creator received room_created with code: {room_code}")
        else:
            print(f"❌ Unexpected response for creator: {response}")
            sys.exit(1)
    except asyncio.TimeoutError:
        print("❌ Timeout waiting for 'room_created'")
        sys.exit(1)
        
    # Client 2: Joiner
    try:
        joiner_ws = await websockets.connect(url)
        print("✅ Joiner connected to WebSocket")
    except Exception as e:
        print(f"❌ Joiner failed to connect to server: {e}")
        sys.exit(1)
        
    joiner_msg = {
        "type": "join_room",
        "username": "TestJoiner",
        "code": room_code
    }
    
    await joiner_ws.send(json.dumps(joiner_msg))
    print(f"📤 Joiner sent 'join_room' for code: {room_code}")
    
    # Wait for file_check needed
    try:
        response_str = await asyncio.wait_for(joiner_ws.recv(), timeout=2.0)
        response = json.loads(response_str)
        if response.get("type") == "room_file_check":
            print(f"✅ Joiner received room_file_check from server. Payload: {response}")
            
            # Send file confirmation
            confirm_msg = {
                "type": "confirm_join",
                "fileInfo": {
                    "name": "test_video.mkv",
                    "size": 1048576,
                    "duration": 0
                }
            }
            await joiner_ws.send(json.dumps(confirm_msg))
            print("📤 Joiner sent 'confirm_join'")
            
            # Wait for successful room_joined
            response_str = await asyncio.wait_for(joiner_ws.recv(), timeout=2.0)
            response = json.loads(response_str)
            if response.get("type") == "room_joined":
                print(f"✅ Joiner successfully received room_joined: {list(response.keys())}")
                
                # Check users payload from the join response itself
                users = response.get("users", [])
                print(f"✅ Joiner received user list from room_joined payload: {users}")
                assert len(users) == 2, "Expected 2 users in the room!"
                    
                # Setup wait task for the Joiner to receive a Play sync event
                print("⏳ Creator is sending a Play sync event...")
                await creator_ws.send(json.dumps({
                    "type": "sync",
                    "action": "play",
                    "position": 5.0
                }))
                
                # Creator will get an echo, Joiner will get the broadcast
                creator_echo = json.loads(await asyncio.wait_for(creator_ws.recv(), timeout=2.0))
                joiner_sync = json.loads(await asyncio.wait_for(joiner_ws.recv(), timeout=2.0))
                
                if joiner_sync.get("type") == "sync":
                    print(f"✅ Joiner received sync event: {joiner_sync}")
                    assert joiner_sync.get("action") == "play", "Expected 'play' sync event"
                    assert joiner_sync.get("position") == 5.0, "Expected position 5.0"
                else:
                    print(f"❌ Joiner expected sync, got: {joiner_sync}")
                    sys.exit(1)
                    
                # Setup wait task for the Joiner to send a Chat message
                print("⏳ Joiner is sending a Chat message...")
                await joiner_ws.send(json.dumps({
                    "type": "chat",
                    "message": "Hello from integration test!"
                }))
                
                creator_chat = json.loads(await asyncio.wait_for(creator_ws.recv(), timeout=2.0))
                joiner_chat = json.loads(await asyncio.wait_for(joiner_ws.recv(), timeout=2.0))
                
                if creator_chat.get("type") == "chat":
                    print(f"✅ Creator received chat event: {creator_chat}")
                    assert creator_chat.get("username") == "TestJoiner", "Expected TestJoiner author"
                    assert creator_chat.get("message") == "Hello from integration test!", "Expected correct message"
                else:
                    print(f"❌ Creator expected chat, got: {creator_chat}")
                    sys.exit(1)
                    
            else:
                print(f"❌ Joiner expected room_joined, got: {response}")
        else:
            print(f"❌ Joiner expected room_file_check, got: {response}")
    except asyncio.TimeoutError:
        print("❌ Timeout waiting for server response to 'join_room'")
        sys.exit(1)
        
    await creator_ws.close()
    await joiner_ws.close()
    print("🎉 Integration test completed successfully!")

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