"""Exercise the gateway over real TCP and WebSocket transports."""

from __future__ import annotations

import asyncio
import json
import os
from typing import Any

import httpx
import websockets


async def receive_type(device: Any, message_type: str) -> dict[str, Any]:
    for _ in range(10):
        message = json.loads(await asyncio.wait_for(device.recv(), 3))
        if message.get("type") == message_type:
            return message
    raise AssertionError(f"did not receive {message_type}")


async def run() -> None:
    base_url = os.environ.get("GATEWAY_HTTP_URL", "http://127.0.0.1:18080")
    ws_url = os.environ.get("GATEWAY_WS_URL", "ws://127.0.0.1:18080/device/ws")
    device_token = os.environ["GATEWAY_TEST_DEVICE_TOKEN"]
    pi_token = os.environ["GATEWAY_TEST_PI_TOKEN"]
    async with httpx.AsyncClient(base_url=base_url, timeout=5) as client:
        health = await client.get("/health")
        health.raise_for_status()
        assert health.json()["status"] == "ok"
        transport = "HTTPS" if base_url.startswith("https://") else "HTTP"
        print(f"{transport} PASS: {base_url}/health")
        admin = await client.get("/admin")
        admin.raise_for_status()
        assert "ESP32 Device Lab" in admin.text
        status = await client.get("/api/status")
        status.raise_for_status()
        assert status.json()["server"]["type"] == "server_snapshot"
        print("Admin PASS: HTML and API status")

        async with websockets.connect(ws_url) as device:
            await device.send(
                json.dumps(
                    {
                        "type": "hello",
                        "device_id": "smoke-device",
                        "device_token": device_token,
                        "board": "m5stack-stopwatch",
                        "firmware": "0.1.0",
                    }
                )
            )
            assert await receive_type(device, "hello_ack") == {"type": "hello_ack"}
            print(f"WSS PASS: hello/hello_ack via {ws_url}")
            snapshot = await receive_type(device, "server_snapshot")
            assert snapshot["type"] == "server_snapshot"
            assert 0 <= snapshot["cpu_percent"] <= 100
            assert snapshot["network_rx_bps"] >= 0
            print("Gateway -> Device PASS: server_snapshot received")
            assert await receive_type(device, "target_state") == {
                "type": "target_state",
                "target": "server",
            }

            listed = (await client.get("/devices")).json()
            smoke_device = next(item for item in listed if item["device_id"] == "smoke-device")
            assert smoke_device["online"] is True

            await device.send(json.dumps({"type": "heartbeat", "uptime": 42}))
            await device.send(
                json.dumps(
                    {
                        "type": "button_event",
                        "button": "A",
                        "action": "click",
                        "timestamp": 1000,
                    }
                )
            )

            for _ in range(50):
                listed = (await client.get("/devices")).json()
                smoke_device = next(
                    item for item in listed if item["device_id"] == "smoke-device"
                )
                if smoke_device["last_event"] == {
                    "type": "button_event",
                    "button": "A",
                    "action": "click",
                    "timestamp": 1000,
                }:
                    break
                await asyncio.sleep(0.1)
            else:
                raise AssertionError("gateway did not record device button_event")
            print("Device -> Cloud PASS: heartbeat and button_event recorded")

            sent = await client.post(
                "/devices/smoke-device/message",
                json={"title": "Smoke", "body": "transport test"},
            )
            sent.raise_for_status()
            assert await receive_type(device, "display_message") == {
                "type": "display_message",
                "title": "Smoke",
                "body": "transport test",
            }
            print("Cloud -> Device PASS: display_message received")

            pi_event = {
                "type": "pi_session_state",
                "session_id": "smoke-pi-session",
                "name": "esp32-dev",
                "cwd": "/home/lighthouse/esp32-dev",
            }
            headers = {"x-device-lab-pi-token": pi_token}
            for event in ("session_start", "agent_start"):
                response = await client.post(
                    "/integrations/pi/events", json={**pi_event, "event": event}, headers=headers
                )
                response.raise_for_status()
            selected = await client.post(
                "/devices/smoke-device/target", json={"target": "pi:smoke-pi-session"}
            )
            selected.raise_for_status()
            running = await receive_type(device, "agent_state")
            assert running["state"] == "RUNNING"
            completed = await client.post(
                "/integrations/pi/events",
                json={**pi_event, "event": "agent_end"},
                headers=headers,
            )
            completed.raise_for_status()
            done = await receive_type(device, "agent_state")
            assert done["state"] == "DONE"
            switched = await client.post(
                "/devices/smoke-device/target", json={"target": "server"}
            )
            switched.raise_for_status()
            assert (await receive_type(device, "target_state"))["target"] == "server"
            print("Pi -> Gateway -> Device PASS: RUNNING and DONE target states")

    print("gateway TCP/WebSocket smoke test passed")


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