#!/usr/bin/env python3
"""WebTransport fallback test — verifies hybrid mode falls back to WS when WT is unavailable."""

import asyncio
import json
import sys
import time

try:
    import websockets
except ImportError:
    import subprocess
    subprocess.check_call([sys.executable, "-m", "pip", "install", "websockets", "-q"])
    import websockets

WS_URL = "ws://localhost:28888/ws/chat"
API_BASE = "http://localhost:28888"
RESULTS = []

def log(tag, msg):
    ts = time.strftime("%H:%M:%S")
    line = f"[{ts}] [{tag}] {msg}"
    print(line)
    RESULTS.append(line)

async def recv_json(ws, timeout=5):
    raw = await asyncio.wait_for(ws.recv(), timeout=timeout)
    return json.loads(raw)

async def send_json(ws, data):
    await ws.send(json.dumps(data))

async def test_wt_fallback():
    log("TEST", "=== WebTransport Fallback Test Start ===")

    log("WT", "--- Test 1: WT not configured, WS fallback should work ---")
    log("WT", "connecting via WS (fallback path)...")
    ws = await websockets.connect(WS_URL, additional_headers={"x-auth-user-id": "2001"})
    resp = await recv_json(ws)
    log("WT", f"fallback WS connected: {resp}")
    assert resp["type"] == "connected", f"expected connected, got {resp['type']}"

    log("WT", "--- Test 2: Full chat flow over fallback WS ---")
    await send_json(ws, {"type": "create_group", "name": "wt-fallback-room"})
    resp = await recv_json(ws)
    log("WT", f"create_group via fallback: {resp}")
    assert resp["type"] == "group_created"
    group_id = resp["sessionId"]

    await send_json(ws, {"type": "join", "sessionId": group_id})
    resp1 = await recv_json(ws)
    log("WT", f"join broadcast: {resp1}")
    resp2 = await recv_json(ws)
    log("WT", f"join confirm: {resp2}")

    await send_json(ws, {"type": "message", "sessionId": group_id, "content": "hello from WT fallback!"})
    resp = await recv_json(ws)
    log("WT", f"message via fallback: {resp}")
    assert resp["type"] == "message"
    assert resp["content"] == "hello from WT fallback!"

    log("WT", "--- Test 3: DM over fallback ---")
    await send_json(ws, {"type": "create_dm", "targetId": 2002})
    resp = await recv_json(ws)
    log("WT", f"DM created via fallback: {resp}")
    assert resp["type"] == "dm_created"
    dm_id = resp["sessionId"]

    await send_json(ws, {"type": "message", "sessionId": dm_id, "content": "DM via fallback"})
    resp = await recv_json(ws)
    log("WT", f"DM message via fallback: {resp}")
    assert resp["type"] == "message"

    log("WT", "--- Test 4: History over fallback ---")
    await send_json(ws, {"type": "history", "sessionId": group_id})
    resp = await recv_json(ws)
    log("WT", f"history via fallback: type={resp['type']}")
    assert resp["type"] == "history"

    log("WT", "--- Test 5: Multi-user over fallback ---")
    ws2 = await websockets.connect(WS_URL, additional_headers={"x-auth-user-id": "2002"})
    resp = await recv_json(ws2)
    assert resp["type"] == "connected"
    log("WT", "user2 connected via fallback WS")

    await send_json(ws, {"type": "message", "sessionId": dm_id, "content": "cross-user DM"})
    resp = await recv_json(ws)
    log("WT", f"sender got broadcast: {resp}")
    assert resp["type"] == "message"

    try:
        resp2 = await recv_json(ws2, timeout=2)
        log("WT", f"user2 got message: {resp2}")
    except asyncio.TimeoutError:
        log("WT", "user2 did not receive (not joined to session, expected)")

    await ws.close()
    await ws2.close()

    log("WT", "--- Test 6: Verify WT port is not serving (no TLS configured) ---")
    try:
        wt_ws = await asyncio.wait_for(
            websockets.connect("ws://localhost:28443/wt/chat"),
            timeout=2
        )
        await wt_ws.close()
        log("WT", "WT port unexpectedly open (might be OK if hybrid exposes it)")
    except (ConnectionRefusedError, asyncio.TimeoutError, OSError):
        log("WT", "WT port not available (expected — no TLS cert configured)")

    log("TEST", "=== WebTransport Fallback Test PASSED ===")

async def main():
    try:
        await test_wt_fallback()
        return 0
    except Exception as e:
        log("FAIL", f"Test failed: {e}")
        import traceback
        log("FAIL", traceback.format_exc())
        return 1

if __name__ == "__main__":
    rc = asyncio.run(main())
    with open("logs/wt_test.log", "w") as f:
        f.write("\n".join(RESULTS) + "\n")
    sys.exit(rc)
