#!/usr/bin/env python3
"""
Simple WebSocket test for hwGpt chatbot.
This tests basic WebSocket connectivity without complex logic.
"""

import asyncio
import websockets
import json
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


async def test_simple_websocket():
    """Test basic WebSocket connection."""
    uri = "ws://localhost:8000/hw_gpt/chat/1/35"
    
    try:
        logger.info(f"Connecting to: {uri}")
        
        async with websockets.connect(uri) as websocket:
            logger.info("✅ WebSocket connected successfully!")
            
            # Wait for connection message
            try:
                response = await asyncio.wait_for(websocket.recv(), timeout=5.0)
                data = json.loads(response)
                logger.info(f"Received: {data}")
                
                if data.get("type") == "connection_established":
                    logger.info("✅ Connection established message received!")
                else:
                    logger.warning(f"Unexpected message type: {data.get('type')}")
                    
            except asyncio.TimeoutError:
                logger.warning("⏰ No response received within 5 seconds")
            
            # Test sending a simple message
            test_message = {
                "type": "chat",
                "message": "Hello! This is a test message."
            }
            
            logger.info(f"📤 Sending test message: {test_message}")
            await websocket.send(json.dumps(test_message))
            
            # Wait for response
            try:
                response = await asyncio.wait_for(websocket.recv(), timeout=10.0)
                data = json.loads(response)
                logger.info(f"📥 Received response: {data}")
                
            except asyncio.TimeoutError:
                logger.warning("⏰ No response to test message within 10 seconds")
            
            logger.info("✅ WebSocket test completed successfully!")
            
    except Exception as e:
        logger.error(f"❌ WebSocket test failed: {str(e)}")
        import traceback
        traceback.print_exc()


async def test_websocket_with_branch():
    """Test WebSocket connection with branch ID."""
    uri = "ws://localhost:8000/hw_gpt/chat/1/35/53"
    
    try:
        logger.info(f"Connecting to: {uri}")
        
        async with websockets.connect(uri) as websocket:
            logger.info("✅ WebSocket with branch connected successfully!")
            
            # Wait for connection message
            try:
                response = await asyncio.wait_for(websocket.recv(), timeout=5.0)
                data = json.loads(response)
                logger.info(f"Received: {data}")
                
            except asyncio.TimeoutError:
                logger.warning("⏰ No response received within 5 seconds")
            
            logger.info("✅ WebSocket with branch test completed!")
            
    except Exception as e:
        logger.error(f"❌ WebSocket with branch test failed: {str(e)}")


if __name__ == "__main__":
    print("🚀 Simple WebSocket Test for hwGpt")
    print("=" * 50)
    print("Make sure your FastAPI server is running on localhost:8000")
    print()
    
    # Test basic connection
    print("1️⃣ Testing basic WebSocket connection...")
    asyncio.run(test_simple_websocket())
    
    print("\n2️⃣ Testing WebSocket with branch ID...")
    asyncio.run(test_websocket_with_branch())
    
    print("\n🎉 All tests completed!")