import asyncio
import time
from pathlib import Path
from typing import Sequence, Union, List, Tuple

import aiohttp
import safe_radix32
from clarabot_crypto.context.account_key.room import ACCOUNT_ROOM_KEY_CRYPTO_CONTEXT, AccountRoomKeyConfigDumpArgs
from clarabot_crypto.context.room_key import ROOM_KEY_CRYPTO_CONTEXT, RoomKeyConfigDumpArgs
from clarabot_crypto.keyring.account_peer import PeerAccountKeyring
from clarabot_crypto.keyring.room import RoomKeyring
from clarabot_nano_common.utility.serialize import msgpack_loads

from clarabot_nano_tests.testing.privy.common import parse_setup_url


class DevPrivyManager:
    ENGINE_DEV_PORT = 8338

    def __init__(self, session: aiohttp.ClientSession):
        self.session = session

    def _get_url(self, path: str) -> str:
        return f"http://127.0.0.1:{self.ENGINE_DEV_PORT}/{path}"

    async def get_status_message(
        self, wait_status: Union[str, Sequence[str]] = None, timeout: int = 15, sleep: float = 0.15
    ) -> dict:
        if isinstance(wait_status, str):
            wait_status = (wait_status,)
        until = time.perf_counter() + timeout
        while True:
            try:
                async with self.session.get(self._get_url("status/message")) as r:
                    r.raise_for_status()
                    sm = await r.json()
                if isinstance(sm, dict):
                    if wait_status is None or sm.get("status") in wait_status:
                        return sm
            except Exception:
                if time.perf_counter() > until:
                    raise
            else:
                assert time.perf_counter() <= until, f"status_message wait failed for {wait_status} by {sm}"
            await asyncio.sleep(sleep)

    async def get_signing_pk(self) -> dict:
        async with self.session.get(self._get_url("keyring/p2")) as r:
            r.raise_for_status()
            return msgpack_loads(await r.read())["p2"]

    async def get_remote_config(self) -> dict:
        async with self.session.get(self._get_url("config/get")) as r:
            r.raise_for_status()
            return msgpack_loads(await r.read())

    async def update_remote_config_drive(self, drives: List[dict], filters: List[str]):
        async with self.session.post(
            self._get_url("config/update/drive"), json={"drives": drives, "filters": filters}
        ) as r:
            r.raise_for_status()

    async def setup_privy_room(
        self, testing_client_manager, wsc, wait_timeout: int = 15
    ) -> Tuple[int, PeerAccountKeyring, RoomKeyring]:
        sm = await self.get_status_message("waiting-config", timeout=wait_timeout)
        assert isinstance(sm.get("setup_url"), str)

        # 0. We need the keyring of the account (self-keyring).
        self_kr = testing_client_manager.session_storage(wsc.session)["acc_kr"]

        # 1. Get setup_url and parse it to have the proper values.
        #       setup_url : "https://{url_domain}/privy/setup#{fragment}"
        #       fragment  : {"sk": base64_encoded_setup_key, "p1": base64_encoded_privy_curve25519_pk}
        #       setup_args: {"sk": uint8array_setup_key, "p1": uint8array_privy_curve25519_pk}
        setup_url = sm["setup_url"]
        setup_args = parse_setup_url(setup_url)
        # 2. User confirms they want to setup the agent.
        # 3. Call "createNewId" WS mutation with type="room-id" to get room_id.
        room_id = await testing_client_manager.gql_create_new_id(wsc, "room-id")
        # 4. Generate a new keyring for the privy's room.
        room_kr = RoomKeyring.generate(room_id)
        # 5. Load the privy peer-keyring with the room_id and the "p1" key.
        peer_kr = PeerAccountKeyring(room_id, {PeerAccountKeyring.CURVE25519: setup_args["p1"]})
        # 6. Encrypt the room secret for the host/self account.
        self_room_secret = ACCOUNT_ROOM_KEY_CRYPTO_CONTEXT.dump(
            AccountRoomKeyConfigDumpArgs(self_kr, RoomKeyring.PINNED, room_kr)
        )
        # 7. Encrypt the room secret for the privy account.
        privy_room_secret = ROOM_KEY_CRYPTO_CONTEXT.dump(
            RoomKeyConfigDumpArgs(self_kr, peer_kr, room_kr, RoomKeyring.PINNED)
        )
        # 8. Call the setupPrivy WS mutation with the arguments:
        #       roomId=room_id
        #       privySetupKey=setup_args["sk"]
        #       selfRoomSecret=self_room_secret
        #       privyRoomSecret=privy_room_secret
        await testing_client_manager.gql_privy_setup(
            wsc, room_id, setup_args["sk"], self_room_secret, privy_room_secret
        )

        await self.get_status_message(("connected", "active"))
        await self.get_status_message("active")
        # NOTE: Get complete privy-kr.
        peer_kr = PeerAccountKeyring(
            room_id,
            {
                PeerAccountKeyring.CURVE25519: setup_args["p1"],
                PeerAccountKeyring.ED25519: await self.get_signing_pk(),
            },
        )

        return safe_radix32.decode(room_id), peer_kr, room_kr

    async def configure_drives(self, drive_paths: List[Path], filters: List[str]):
        drives = [{"name": drive_path.name, "path": str(drive_path)} for drive_path in drive_paths]
        await self.update_remote_config_drive(drives, filters)
        rc = await self.get_remote_config()
        assert rc.get("drives") == drives
        assert rc.get("filters") == filters
