#!/usr/bin/env python3
"""
Probe the TIDY-CRE WebSocket endpoint and report status.
Usage: probe_ws.py [host] [port] [timeout]
"""

import sys
import asyncio
import json

try:
    import websockets
except ImportError:
    print("[ERROR] websockets package not installed. Install with: pip install websockets")
    sys.exit(1)


async def probe_websocket(host: str, port: int, timeout: float):
    """Connect to WebSocket and return connection info."""
    uri = f"ws://{host}:{port}"
    print(f"Probing WebSocket at {uri}...")
    
    try:
        # Use asyncio.wait_for for timeout compatibility
        async with websockets.connect(uri) as websocket:
            print(f"[OK] Connected to {uri}")
            
            # Try to get some info from the connection
            print(f"  - Local address: {websocket.local_address}")
            print(f"  - Remote address: {websocket.remote_address}")
            print(f"  - Protocol: {websocket.subprotocol or 'default'}")
            
            return True
    except asyncio.TimeoutError:
        print(f"[FAIL] Connection timed out after {timeout}s")
        return False
    except ConnectionRefusedError:
        print(f"[FAIL] Connection refused - service not running")
        return False
    except Exception as e:
        print(f"[FAIL] Error: {e}")
        return False


async def run_with_timeout(host: str, port: int, timeout: float):
    """Run probe with timeout wrapper."""
    try:
        return await asyncio.wait_for(probe_websocket(host, port, timeout), timeout=timeout)
    except asyncio.TimeoutError:
        print(f"[FAIL] Overall operation timed out after {timeout}s")
        return False

def main():
    host = sys.argv[1] if len(sys.argv) > 1 else "127.0.0.1"
    port = int(sys.argv[2]) if len(sys.argv) > 2 else 18789
    timeout = float(sys.argv[3]) if len(sys.argv) > 3 else 5.0
    
    result = asyncio.run(run_with_timeout(host, port, timeout))
    sys.exit(0 if result else 1)


if __name__ == "__main__":
    main()