import asyncio
import json

import pytest
import websockets

from edge_proxy.nav_backend.base import Pose2D
from edge_proxy.nav_backend.rosbridge_backend import RosbridgeBackend


@pytest.mark.asyncio
async def test_rosbridge_backend_navigate_flow():
    updates = []
    published = []

    async def rosbridge_handler(ws):
        async for raw in ws:
            msg = json.loads(raw)
            if msg.get("op") != "publish":
                continue

            published.append(msg)
            if msg.get("topic") != "gui/execute_plan":
                continue

            await ws.send(
                json.dumps(
                    {
                        "op": "publish",
                        "topic": "gui/get_plan_feedback",
                        "msg": {"data": 1},
                    }
                )
            )
            await ws.send(
                json.dumps(
                    {
                        "op": "publish",
                        "topic": "gui/get_waypoints_new",
                        "msg": {
                            "waypoints": [
                                {"status": 2},
                                {"status": 1},
                            ],
                            "next_waypoint": 1,
                        },
                    }
                )
            )
            await ws.send(
                json.dumps(
                    {
                        "op": "publish",
                        "topic": "gui/get_plan_feedback",
                        "msg": {"data": 0},
                    }
                )
            )

    ws_server = await websockets.serve(rosbridge_handler, "127.0.0.1", 0)
    try:
        port = ws_server.sockets[0].getsockname()[1]
        backend = RosbridgeBackend(
            rosbridge_url=f"ws://127.0.0.1:{port}",
            insecure_tls=False,
        )
        await backend.start()
        try:
            async def _on_update(update):
                updates.append((update.status, update.progress))

            await backend.navigate_to_pose(
                request_id="nav_001",
                destination="office",
                pose=Pose2D(x=1.0, y=2.0, theta=0.5),
                speed="normal",
                on_update=_on_update,
            )
        finally:
            await backend.stop()

        statuses = [s for s, _ in updates]
        assert "accepted" in statuses
        assert "navigating" in statuses
        assert "arrived" in statuses

        add_waypoint_msgs = [m for m in published if m.get("topic") == "gui/add_waypoint_new"]
        execute_msgs = [m for m in published if m.get("topic") == "gui/execute_plan"]
        assert add_waypoint_msgs, "No gui/add_waypoint_new publish observed"
        assert execute_msgs, "No gui/execute_plan publish observed"
    finally:
        ws_server.close()
        await ws_server.wait_closed()


@pytest.mark.asyncio
async def test_rosbridge_backend_telemetry_updates():
    sent_telemetry = False

    async def rosbridge_handler(ws):
        nonlocal sent_telemetry
        async for raw in ws:
            msg = json.loads(raw)
            if msg.get("op") == "subscribe" and not sent_telemetry:
                sent_telemetry = True
                await ws.send(
                    json.dumps(
                        {
                            "op": "publish",
                            "topic": "gui/get_robot_pose",
                            "msg": {"x": 3.2, "y": -1.5, "z": 0.7},
                        }
                    )
                )
                await ws.send(
                    json.dumps(
                        {
                            "op": "publish",
                            "topic": "/kstack/state/battery",
                            "msg": {"percentage": 0.62, "power_supply_status": 1},
                        }
                    )
                )

    ws_server = await websockets.serve(rosbridge_handler, "127.0.0.1", 0)
    try:
        port = ws_server.sockets[0].getsockname()[1]
        backend = RosbridgeBackend(
            rosbridge_url=f"ws://127.0.0.1:{port}",
            insecure_tls=False,
        )
        await backend.start()
        try:
            await asyncio.sleep(0.1)
            pose = backend.get_pose()
            battery = backend.get_battery()
            assert pose.x == pytest.approx(3.2)
            assert pose.y == pytest.approx(-1.5)
            assert pose.theta == pytest.approx(0.7)
            assert battery.level == 62
            assert battery.charging is True
        finally:
            await backend.stop()
    finally:
        ws_server.close()
        await ws_server.wait_closed()