import asyncio
import itertools
import random
import time
from typing import Sequence, Dict, Optional, Tuple

import pytest
import safe_radix32
from clarabot_crypto.context.nano import NANO_CRYPTO_CONTEXT, NanoConfigDumpArgs
from clarabot_crypto.context.nano_own import NANO_OWN_CRYPTO_CONTEXT, NanoOwnConfigDumpArgs

from clarabot_nano.common.crypto.keyring import NanoSelfAccountKeyring, NanoRoomKeyring, NanoPeerAccountKeyring
from clarabot_nano.common.handler.base import AbstractNanoHandler
from clarabot_nano.common.handler.registry import NanoHandlerClassRegistry
from clarabot_nano.common.protocol.errors import (
    NANO_REQUEST_ERROR_PREPARSE,
    NANO_REQUEST_ERROR_PAYLOAD_PARSE,
    NANO_REQUEST_ERROR_DENIED,
    NANO_REQUEST_ERROR_PAYLOAD_ROOM_INVALID,
    NANO_REQUEST_ERROR,
    NANO_REQUEST_ERROR_INTERNAL_FAILURE,
    NANO_REQUEST_ERROR_TIMEOUT,
    NANO_REQUEST_ERROR_TEMPORARY_INTERNAL_FAILURE,
    NANO_REQUEST_ERROR_PAYLOAD_TYPE_INVALID,
    NANO_REQUEST_ERROR_PEER_NOT_TRUSTED,
    NANO_REQUEST_ERROR_MUTATION_ID,
    NANO_REQUEST_ERROR_UNKNOWN_TYPE,
)
from clarabot_nano.common.protocol.server import NanoServerResponse, NanoServerRequest
from clarabot_nano.common.request.constants import HANDLER_KEY_PREFIX_MAP
from clarabot_nano.common.spawn.constants import ShortNames
from clarabot_nano.common.zmq.requester import NoWorkerError, AioWorkerPoolRequester
from clarabot_nano.gateway.config import GatewayConfig
from clarabot_nano.gateway.job.request import state_request_job
from clarabot_nano.gateway.state.config import RoomGroupConfigState
from clarabot_nano.gateway.state.gateway import GatewayState
from clarabot_nano.utility.serialize import msgpack_dumps


class _ErrorTestingRoomGroupConfigState(RoomGroupConfigState):
    READY_WAIT_BY_FLAGS = 0.015

    def __init__(self, *args, **kwargs):
        RoomGroupConfigState.__init__(self, *args, **kwargs)
        self.flags_response = None

    async def get_flags(self, account_id: int) -> Optional[Tuple[int, int]]:
        await self.wait_ready()
        # mock blockchain processing
        return self.flags_response


class _ErrorTestingGatewayState(GatewayState):
    def __init__(self, self_kr: NanoSelfAccountKeyring, deny_anon: bool):
        # init
        GatewayState.__init__(self, None, self.mock_gateway_config(deny_anonymous=deny_anon), self_kr, None)
        self.mock_connection_ack()
        self.test_resps = []
        # setup rooms for testing
        self.room_ids = [9000, 9001, 9002, 9003, 9004, 9005, 9006]
        for room_id in self.room_ids:
            self.room_config_store[room_id] = _ErrorTestingRoomGroupConfigState(room_id)
        self.room_config_store[9001].state = RoomGroupConfigState.STATE_UNVERIFIED
        self.room_config_store[9002].state = RoomGroupConfigState.STATE_READY
        self.room_config_store[9003].state = RoomGroupConfigState.STATE_REVOKED
        self.room_config_store[9004].state = RoomGroupConfigState.STATE_ERROR
        self.room_config_store[9005].state = RoomGroupConfigState.STATE_READY  # config loaded but has no keyring data
        self.room_config_store[9006].state = RoomGroupConfigState.STATE_READY  # config loaded but has no apis
        self.room_krs = {
            i: NanoRoomKeyring.generate(safe_radix32.encode_fw(i)) for i in range(self.room_ids[0], self.room_ids[-2])
        }
        self.room_krs[9006] = NanoRoomKeyring.generate(safe_radix32.encode_fw(9006))
        # convenience
        self.room_id_loading = self.room_ids[0]
        self.room_id_unverified = self.room_ids[1]
        self.room_id_ready = self.room_ids[2]
        self.room_id_revoked = self.room_ids[3]
        self.room_id_error = self.room_ids[4]
        self.room_id_ready_no_kr = self.room_ids[5]
        self.room_id_ready_no_api = self.room_ids[6]
        # setup accounts for testing
        self.anon_krs = [NanoSelfAccountKeyring.generate_anonymous() for _ in range(3)]
        self.acc_ids = list(range(15000, 15010))
        self.acc_krs: Dict[int, NanoSelfAccountKeyring] = {
            i: NanoSelfAccountKeyring.generate(safe_radix32.encode_fw(i)) for i in self.acc_ids
        }
        self.peer_krs: Dict[int, NanoPeerAccountKeyring] = {
            i: NanoPeerAccountKeyring(kr.id.decode(), dict(kr._public_keys)) for i, kr in self.acc_krs.items()
        }
        self.peer_krs[15001].trusted = NanoPeerAccountKeyring.TRUSTED_FULLY
        self.peer_krs[15002].trusted = NanoPeerAccountKeyring.TRUSTED_STALE
        self.peer_krs[15003].trusted = NanoPeerAccountKeyring.TRUSTED_REQUIRES_UPGRADE
        self.peer_krs[15004].trusted = NanoPeerAccountKeyring.TRUSTED_NOT
        # convenience
        self.peer_id_unknown = self.acc_ids[0]
        self.peer_id_fully = self.acc_ids[1]
        self.peer_id_stale = self.acc_ids[2]
        self.peer_id_requires_upgrade = self.acc_ids[3]
        self.peer_id_not = self.acc_ids[4]
        # update extracted config for test mocking of room data
        self.extracted_config = self.mock_gateway_config(
            drives=[
                {
                    "path": "xy",
                    "rooms": [safe_radix32.encode_fw(i) for i in range(self.room_ids[0], self.room_ids[-1])],
                }
            ],
            deny_anonymous=deny_anon,
        )

    def mock_gateway_config(
        self,
        name="test",
        remote_admin_policy="restrict",
        admin_password=None,
        room_blocks={},
        deny_anonymous=True,
        drives=(),
    ):
        return GatewayConfig(
            {
                "local": {"remote_admin_policy": remote_admin_policy, "data_dir": "n/a"},
                "remote": {
                    "name": name,
                    "admin_password": admin_password,
                    "room_blocks": dict(room_blocks),
                    "deny_anonymous": deny_anonymous,
                    "drives": list(drives),
                    "secrets": [b"n/a"],
                },
            }
        )

    def mock_connection_ack(self, part_size=1000000):
        self._init_from_connection_ack(safe_radix32.decode(self.keyring_cache.self_kr.id.decode()), part_size)

    async def send_response(self, op_id, response):
        """
        Mocks the real implementation with a local collection.
        """
        self.test_resps.append((op_id, response))

    async def get_own_room_keyring(self, room_id: int) -> NanoRoomKeyring:
        kr = self.room_krs.get(room_id)
        if not kr:
            raise Exception("No room test kr")
        return kr

    async def get_peer_account_keyring(self, account_id: int) -> NanoPeerAccountKeyring:
        kr = self.peer_krs.get(account_id)
        if not kr:
            raise Exception("No peer test kr")
        return kr


class _ErrorTestingAioWorkerPoolRequester(AioWorkerPoolRequester):
    def __init__(self):
        # NOTE: purposefully not calling super inits
        self.test_error = None
        self.test_reqs = []

    async def request(self, multipart_message: Sequence[bytes], timeout=12.0, server_prefix: str = None):
        """
        Mocks the real implementation with synthetic results.
        """
        if self.test_error:
            try:
                raise self.test_error
            finally:
                self.test_error = None
        if not self.test_reqs:
            self.test_reqs.append(Exception("Request not expected for server"))
        return self.test_reqs.pop(-1)


class _ErrorTestingNanoHandlerClassRegistry(NanoHandlerClassRegistry):
    pass


class _ErrorTestingHandler(AbstractNanoHandler):
    HANDLER_KEY = "err-hndl"


class _ErrorTestingGatewayProcess(object):
    def __init__(self, deny_anon: bool):
        # NOTE: purposefully not calling super inits
        self.gws = _ErrorTestingGatewayState(
            NanoSelfAccountKeyring.generate(safe_radix32.encode_fw(random.randint(2 ** 42, 2 ** 52))), deny_anon
        )
        self._worker_requester = _ErrorTestingAioWorkerPoolRequester()
        self._handler_registry = _ErrorTestingNanoHandlerClassRegistry([_ErrorTestingHandler])

    async def dbg_req_proc_err(self, server_request, error, server_resps=()):
        await self.dbg_req_proc(server_request, {NanoServerResponse.ERROR: error}, server_resps)

    async def dbg_req_proc(self, server_request, response, server_resps=()):
        assert not self.get_responses()
        self.set_server_responses(server_resps)
        op_id = random.randint(0, 2 ** 48)
        await state_request_job(self.gws, self._worker_requester, self._handler_registry, op_id, server_request)
        rs = self.get_responses()
        assert len(rs) == 1, f"{rs}\n{response}"
        rs[0][1].pop("l", None)
        assert rs[0] == (op_id, response), f"{rs[0]}\n{response}"

    def set_server_responses(self, reqs):
        self._worker_requester.test_reqs = list(reqs)

    def set_server_error(self, error: Exception):
        self._worker_requester.test_error = error

    def get_responses(self):
        tmp = self.gws.test_resps
        self.gws.test_resps = []
        return tmp


@pytest.mark.parametrize("deny_anon", [True, False])
@pytest.mark.asyncio
async def test_gateway_request_errors(monkeypatch, deny_anon):
    p = _ErrorTestingGatewayProcess(deny_anon)
    monkeypatch.setitem(HANDLER_KEY_PREFIX_MAP, _ErrorTestingHandler.HANDLER_KEY, ShortNames.GATEWAY)
    monkeypatch.setitem(HANDLER_KEY_PREFIX_MAP, "test-gw", ShortNames.GATEWAY)

    def _deny_anon_error(sample_, error_):
        if deny_anon and sample_.get(NanoServerRequest.ACCOUNT_ID) == 0:
            return NANO_REQUEST_ERROR_DENIED
        return error_

    # preparse errors
    await p.dbg_req_proc_err(b"xxxx", NANO_REQUEST_ERROR_PREPARSE)
    for sample in [
        {},
        {NanoServerRequest.PAYLOAD: b"xxxx"},
        {NanoServerRequest.PAYLOAD: b"xxxx", NanoServerRequest.CACHE_KEY: 5},
        {NanoServerRequest.PAYLOAD: b"xxxx", NanoServerRequest.TRANSFER_KEY: 5},
        {NanoServerRequest.PAYLOAD: b"xxxx", NanoServerRequest.PART: "xa"},
        {
            NanoServerRequest.PAYLOAD: b"xxxx",
            NanoServerRequest.CACHE_KEY: b"a",
            NanoServerRequest.TRANSFER_KEY: b"b",
        },
        {NanoServerRequest.PAYLOAD: b"xxxx", NanoServerRequest.CACHE_KEY: b"a"},
        {NanoServerRequest.PAYLOAD: b"xxxx", NanoServerRequest.TRANSFER_KEY: b"a"},
        {NanoServerRequest.PAYLOAD: b"xxxx", NanoServerRequest.CACHE_KEY: b"a", NanoServerRequest.PART: 0},
        {NanoServerRequest.PAYLOAD: b"xxxx", NanoServerRequest.ACCOUNT_ID: "asdf"},
        {NanoServerRequest.PAYLOAD: b"xxxx", NanoServerRequest.ACCOUNT_ID: 0},
        {NanoServerRequest.PAYLOAD: b"xxxx", NanoServerRequest.ACCOUNT_ID: 0, NanoServerRequest.ROOM_ID: "asdf"},
    ]:
        await p.dbg_req_proc_err(msgpack_dumps(sample), NANO_REQUEST_ERROR_PREPARSE)

    # early denied anon
    await p.dbg_req_proc_err(
        msgpack_dumps(
            {
                NanoServerRequest.PAYLOAD: b"xxxx",
                NanoServerRequest.ACCOUNT_ID: 0,
                NanoServerRequest.ROOM_ID: p.gws.room_id_ready,
            }
        ),
        NANO_REQUEST_ERROR_DENIED if deny_anon else NANO_REQUEST_ERROR_PREPARSE,
    )

    # early denied for not served room
    await p.dbg_req_proc_err(
        msgpack_dumps(
            {NanoServerRequest.PAYLOAD: b"xxxx", NanoServerRequest.ACCOUNT_ID: 0, NanoServerRequest.ROOM_ID: 13}
        ),
        NANO_REQUEST_ERROR_DENIED,
    )
    # peer not trusted
    await p.dbg_req_proc_err(
        msgpack_dumps(
            {
                NanoServerRequest.PAYLOAD: b"xxxx",
                NanoServerRequest.ACCOUNT_ID: p.gws.peer_id_not,
                NanoServerRequest.ROOM_ID: p.gws.room_id_ready,
            }
        ),
        NANO_REQUEST_ERROR_PEER_NOT_TRUSTED,
    )

    # decrypt errors
    for sample in [
        {
            NanoServerRequest.ACCOUNT_ID: 0,
            NanoServerRequest.EPHEMERAL_PUBLIC_KEY: p.gws.anon_krs[0].get_public_key(NanoPeerAccountKeyring.CURVE25519),
            NanoServerRequest.ROOM_ID: p.gws.room_id_ready,
        },
        {NanoServerRequest.ACCOUNT_ID: p.gws.account_id},
        {NanoServerRequest.ACCOUNT_ID: p.gws.account_id, NanoServerRequest.ROOM_ID: p.gws.room_id_ready},
        {NanoServerRequest.ACCOUNT_ID: p.gws.peer_id_fully, NanoServerRequest.ROOM_ID: p.gws.room_id_ready},
    ]:
        tmp = dict(sample)
        tmp[NanoServerRequest.PAYLOAD] = b"\x00xxxx"
        await p.dbg_req_proc_err(msgpack_dumps(tmp), _deny_anon_error(sample, NANO_REQUEST_ERROR_PAYLOAD_PARSE))
        tmp = dict(sample)
        tmp[NanoServerRequest.PAYLOAD] = b"\x01xxxx"
        await p.dbg_req_proc_err(msgpack_dumps(tmp), _deny_anon_error(sample, NANO_REQUEST_ERROR_PAYLOAD_PARSE))

    _COPY_SENTINEL = object()

    def _create_authorized_samples_for_room(
        payload: bytes, direct_payload: Optional[bytes] = _COPY_SENTINEL, room_id: int = p.gws.room_id_ready
    ):
        samples = [
            {
                NanoServerRequest.PAYLOAD: NANO_CRYPTO_CONTEXT.dump(
                    NanoConfigDumpArgs(
                        payload,
                        p.gws.anon_krs[0],
                        p.gws.keyring_cache.self_kr,
                        p.gws.room_krs[room_id],
                        time.time(),
                    )
                ),
                NanoServerRequest.ACCOUNT_ID: 0,
                NanoServerRequest.ROOM_ID: room_id,
                NanoServerRequest.EPHEMERAL_PUBLIC_KEY: p.gws.anon_krs[0].get_public_key(
                    NanoPeerAccountKeyring.CURVE25519
                ),
            },
            {
                NanoServerRequest.PAYLOAD: NANO_CRYPTO_CONTEXT.dump(
                    NanoConfigDumpArgs(
                        payload,
                        p.gws.acc_krs[p.gws.peer_id_fully],
                        p.gws.keyring_cache.self_kr,
                        p.gws.room_krs[room_id],
                        time.time(),
                    )
                ),
                NanoServerRequest.ACCOUNT_ID: p.gws.peer_id_fully,
                NanoServerRequest.ROOM_ID: room_id,
            },
            {
                NanoServerRequest.PAYLOAD: NANO_OWN_CRYPTO_CONTEXT.dump(
                    NanoOwnConfigDumpArgs(
                        payload,
                        p.gws.keyring_cache.self_kr,
                        time.time(),
                    )
                ),
                NanoServerRequest.ACCOUNT_ID: p.gws.account_id,
                NanoServerRequest.ROOM_ID: room_id,
            },
        ]
        if direct_payload is _COPY_SENTINEL:
            direct_payload = payload
        if direct_payload is not None:
            samples.append(
                {
                    NanoServerRequest.PAYLOAD: NANO_OWN_CRYPTO_CONTEXT.dump(
                        NanoOwnConfigDumpArgs(
                            direct_payload,
                            p.gws.keyring_cache.self_kr,
                            time.time(),
                        )
                    ),
                    NanoServerRequest.ACCOUNT_ID: p.gws.account_id,
                }
            )
        return samples

    # payload parse errors
    for sample in _create_authorized_samples_for_room(b"xxxx"):
        await p.dbg_req_proc_err(msgpack_dumps(sample), _deny_anon_error(sample, NANO_REQUEST_ERROR_PAYLOAD_PARSE))

    # handler type serialization errors
    for sample in [
        *_create_authorized_samples_for_room(msgpack_dumps({})),
        *_create_authorized_samples_for_room(msgpack_dumps({"type": None})),
        *_create_authorized_samples_for_room(msgpack_dumps({"type": b"asdf"})),
    ]:
        await p.dbg_req_proc_err(
            msgpack_dumps(sample), _deny_anon_error(sample, NANO_REQUEST_ERROR_PAYLOAD_TYPE_INVALID)
        )

    # room-id consistency and serialization errors
    for sample in [
        *_create_authorized_samples_for_room(
            msgpack_dumps({"type": "asdf"}), msgpack_dumps({"type": "asdf", "room_id": p.gws.room_id_ready})
        ),
        *_create_authorized_samples_for_room(msgpack_dumps({"type": "asdf", "room_id": b"XXX"})),
        *_create_authorized_samples_for_room(msgpack_dumps({"type": "asdf", "room_id": "XXX"})),
        *_create_authorized_samples_for_room(msgpack_dumps({"type": "asdf", "room_id": p.gws.room_id_unverified})),
    ]:
        await p.dbg_req_proc_err(
            msgpack_dumps(sample), _deny_anon_error(sample, NANO_REQUEST_ERROR_PAYLOAD_ROOM_INVALID)
        )

    # accessing room whose config is not ready
    for sample in list(
        itertools.chain(
            *[
                _create_authorized_samples_for_room(msgpack_dumps({"type": "asdf", "room_id": room_id}), None, room_id)
                for room_id in [p.gws.room_id_loading, p.gws.room_id_unverified]
            ]
        )
    ):
        await p.dbg_req_proc_err(
            msgpack_dumps(sample), _deny_anon_error(sample, NANO_REQUEST_ERROR_TEMPORARY_INTERNAL_FAILURE)
        )

    # accessing room whose config is not ready or have no apis
    for sample in list(
        itertools.chain(
            *[
                _create_authorized_samples_for_room(msgpack_dumps({"type": "asdf", "room_id": room_id}), None, room_id)
                for room_id in [p.gws.room_id_revoked, p.gws.room_id_error, p.gws.room_id_ready_no_api]
            ]
        )
    ):
        await p.dbg_req_proc_err(msgpack_dumps(sample), NANO_REQUEST_ERROR_DENIED)

    # denying by room membership or policy
    # Forcing "None" response will indicate no access for any request.
    p.gws.room_config_store[p.gws.room_id_ready].flags_response = None
    for sample in [
        *_create_authorized_samples_for_room(msgpack_dumps({"type": "asdf", "room_id": p.gws.room_id_ready}), None)
    ]:
        await p.dbg_req_proc_err(msgpack_dumps(sample), NANO_REQUEST_ERROR_DENIED)

    # denying room by local policy override
    # Forcing "(0, 1)" response will indicate all requests are only granted anonymous-access.
    p.gws.room_config_store[p.gws.room_id_ready].flags_response = (0, 1)  # not accurate for all samples
    for sample in [
        *_create_authorized_samples_for_room(msgpack_dumps({"type": "asdf", "room_id": p.gws.room_id_ready}), None)
    ]:
        if deny_anon:
            error = NANO_REQUEST_ERROR_DENIED
        elif sample[NanoServerRequest.ACCOUNT_ID] == 0:
            error = NANO_REQUEST_ERROR_PAYLOAD_TYPE_INVALID
        else:
            error = NANO_REQUEST_ERROR_MUTATION_ID
        await p.dbg_req_proc_err(msgpack_dumps(sample), error)

    # mutation-id failures
    # Forcing "(0, 1)" response will indicate all requests are granted member access.
    p.gws.room_config_store[p.gws.room_id_ready].flags_response = (1, 1)  # not accurate for all samples
    for sample in [
        *_create_authorized_samples_for_room(
            msgpack_dumps({"type": "asdf", "room_id": p.gws.room_id_ready}), msgpack_dumps({"type": "asdf"})
        ),
        *_create_authorized_samples_for_room(
            msgpack_dumps({"type": "asdf", "room_id": p.gws.room_id_ready, "mut_key": "mk"}),
            msgpack_dumps({"type": "asdf", "mut_key": "mk"}),
        ),
        *_create_authorized_samples_for_room(
            msgpack_dumps({"type": "asdf", "room_id": p.gws.room_id_ready, "mut_key": "mk", "mut_seq": 1}),
            msgpack_dumps({"type": "asdf", "mut_key": "mk", "mut_seq": 1}),
        ),
    ]:
        await p.dbg_req_proc_err(
            msgpack_dumps(sample),
            _deny_anon_error(
                sample,
                NANO_REQUEST_ERROR_PAYLOAD_TYPE_INVALID
                if sample[NanoServerRequest.ACCOUNT_ID] == 0
                else NANO_REQUEST_ERROR_MUTATION_ID,
            ),
        )

    # unexpected error by failing to load room keyring
    for sample in [
        {
            NanoServerRequest.PAYLOAD: b"xxxx",
            NanoServerRequest.ACCOUNT_ID: p.gws.account_id,
            NanoServerRequest.ROOM_ID: p.gws.room_id_ready_no_kr,
        },
    ]:
        await p.dbg_req_proc_err(msgpack_dumps(sample), NANO_REQUEST_ERROR)

    own_mut_key = p.gws.mutate_tracker.create(p.gws.account_id)
    own_mut_seq = 1

    def _create_direct_server_req(type_: str = "asdf"):
        nonlocal own_mut_seq
        sample = {
            NanoServerRequest.PAYLOAD: NANO_OWN_CRYPTO_CONTEXT.dump(
                NanoOwnConfigDumpArgs(
                    msgpack_dumps({"type": type_, "mut_key": own_mut_key, "mut_seq": own_mut_seq}),
                    p.gws.keyring_cache.self_kr,
                    time.time(),
                )
            ),
            NanoServerRequest.ACCOUNT_ID: p.gws.account_id,
        }
        own_mut_seq += 1
        return msgpack_dumps(sample)

    # server processing errors
    await p.dbg_req_proc_err(_create_direct_server_req(), NANO_REQUEST_ERROR_INTERNAL_FAILURE, server_resps=([],))
    await p.dbg_req_proc_err(
        _create_direct_server_req(), NANO_REQUEST_ERROR_INTERNAL_FAILURE, server_resps=([b"x"] * 3,)
    )
    await p.dbg_req_proc_err(
        _create_direct_server_req(), NANO_REQUEST_ERROR_INTERNAL_FAILURE, server_resps=([b"xxxx"],)
    )
    await p.dbg_req_proc_err(
        _create_direct_server_req(), NANO_REQUEST_ERROR_INTERNAL_FAILURE, server_resps=([msgpack_dumps([42])],)
    )
    await p.dbg_req_proc_err(
        _create_direct_server_req(), NANO_REQUEST_ERROR_INTERNAL_FAILURE, server_resps=([msgpack_dumps({})],)
    )
    p.set_server_error(asyncio.TimeoutError())
    await p.dbg_req_proc_err(_create_direct_server_req(), NANO_REQUEST_ERROR_TIMEOUT)
    p.set_server_error(NoWorkerError())
    await p.dbg_req_proc_err(_create_direct_server_req(), NANO_REQUEST_ERROR_TEMPORARY_INTERNAL_FAILURE)

    # gateway handlers (invocation errors)
    await p.dbg_req_proc_err(_create_direct_server_req("test-gw"), NANO_REQUEST_ERROR_UNKNOWN_TYPE)
    await p.dbg_req_proc_err(_create_direct_server_req(_ErrorTestingHandler.HANDLER_KEY), NANO_REQUEST_ERROR)

    # mutation fail
    own_mut_seq = 2
    await p.dbg_req_proc_err(_create_direct_server_req(), NANO_REQUEST_ERROR_MUTATION_ID)
