#!/usr/bin/env python3
"""WebSocket handshake tester.

Usage:
  python3 scripts/ws_handshake_test.py --url wss://openapivts.koreainvestment.com:29443/websocket/realtime \
      --origin https://securities.koreainvestment.com --subprotocols realtime

Performs an HTTP GET diagnostic (https) and then attempts a websocket.connect()
while printing responses and exceptions.
"""
import argparse
import asyncio
import textwrap

try:
    import httpx
except Exception:
    httpx = None

try:
    import websockets
    from websockets.exceptions import InvalidStatusCode, ConnectionClosed
except Exception:
    websockets = None
    InvalidStatusCode = Exception
    ConnectionClosed = Exception


async def main():
    p = argparse.ArgumentParser()
    p.add_argument("--url", required=True)
    p.add_argument("--origin", required=False)
    p.add_argument("--subprotocols", required=False, help="comma separated")
    args = p.parse_args()

    url = args.url
    origin = args.origin
    subs = None
    if args.subprotocols:
        subs = [s.strip() for s in args.subprotocols.split(",") if s.strip()]

    print("-- HTTP GET diagnostic --")
    if httpx is None:
        print("httpx not installed. Skipping HTTP diagnostic.")
    else:
        http_url = url.replace("wss://", "https://").replace("ws://", "http://")
        headers = {"User-Agent": "ws-handshake-test/1.0"}
        if origin:
            headers["Origin"] = origin
        try:
            async with httpx.AsyncClient(verify=True, timeout=10.0) as client:
                resp = await client.get(http_url, headers=headers)
                print(f"HTTP {resp.status_code}")
                print("Headers:")
                for k, v in list(resp.headers.items())[:10]:
                    print(f"  {k}: {v}")
                print("Body snippet:")
                print(resp.text[:800])
        except Exception as e:
            print("HTTP diagnostic failed:", repr(e))

    print("\n-- WebSocket handshake attempt --")
    if websockets is None:
        print("websockets package not installed. Install with: pip install websockets")
        return

    headers = [("User-Agent", "ws-handshake-test/1.0")]
    if origin:
        headers.append(("Origin", origin))

    kwargs = dict(extra_headers=dict(headers))
    if subs:
        kwargs["subprotocols"] = subs

    try:
        async with websockets.connect(url, **kwargs) as ws:
            print("Connected. Subprotocol:", ws.subprotocol)
            await ws.send("ping")
            try:
                msg = await asyncio.wait_for(ws.recv(), timeout=3.0)
                print("Received:", msg)
            except asyncio.TimeoutError:
                print("No response received (timeout)")
    except InvalidStatusCode as isc:
        status = getattr(isc, "status_code", None)
        print(f"InvalidStatusCode: HTTP {status}")
    except ConnectionClosed as cc:
        print("ConnectionClosed:", cc)
    except Exception as e:
        print("WebSocket connect failed:", repr(e))


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