!/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()