import asyncio
import json
import requests
import websockets


async def test_websocket(
    host: str,
    access_token: str | None,
    motion_group: str,
    cell="cell",
    secure=True,
):
    protocol = "wss" if secure else "ws"

    uri = f"{protocol}://{host}/api/v1/cells/{cell}/motion-groups/{motion_group}/state-stream?response_rate=32"
    additional_headers = (
        {"Authorization": f"Bearer {access_token}"} if access_token else {}
    )

    print(f"## Testing websocket connection to {uri}")
    async with websockets.connect(
        uri, additional_headers=additional_headers
    ) as websocket:
        msg_counter = 0
        while msg_counter < 5:
            message = await websocket.recv()
            data = json.loads(message)
            print("Joint data:", data["result"]["state"]["joint_position"]["joints"])
            msg_counter += 1
    return True


async def test_request(
    host: str,
    access_token: str | None,
    motion_group: str,
    cell="cell",
    secure=True,
):
    protocol = "https" if secure else "http"

    uri = f"{protocol}://{host}/api/v1/cells/{cell}/motion-groups/{motion_group}/state"
    additional_headers = (
        {"Authorization": f"Bearer {access_token}"} if access_token else {}
    )

    print(f"## Testing http connection to {uri}")
    response = requests.get(url=uri, headers=additional_headers, timeout=3)
    if response.status_code != 200:
        print(f"{response.status_code} {response.reason}")

    return response.status_code == 200


def main(
    host: str, cell: str, motion_group: str, secure: bool, access_token: str | None
):
    # TESTS

    # Test HTTP
    test_status = False
    try:
        test_status = asyncio.run(
            test_request(
                host=host,
                access_token=access_token,
                motion_group=motion_group,
                secure=secure,
                cell=cell,
            )
        )
    except Exception as ex:
        print(f"{ex}")
    status = "succeeded" if test_status else "failed"
    print(f"HTTP test {status}")

    # Test Websocket
    test_status = False
    try:
        test_status = asyncio.run(
            test_websocket(
                host=host,
                access_token=access_token,
                motion_group=motion_group,
                secure=secure,
                cell=cell,
            )
        )
    except Exception as ex:
        print(f"{ex}")
    status = "succeeded" if test_status else "failed"
    print(f"Websocket test {status}")


if __name__ == "__main__":
    main(
        host="abc.instance.wandelbots.io",
        access_token="123",
        secure=True,
        cell="cell",
        motion_group="0@ur10e",
    )