import asyncio
import logging
import math
from collections import defaultdict
from contextlib import asynccontextmanager
from functools import partial
from typing import Optional, Dict, Sequence, Any, AsyncContextManager, Set, List, Tuple

import aiohttp
import safe_radix32
from aiohttp import WSMsgType

from clarabot_nano import __version__
from clarabot_nano.common.crypto.keyring import NanoSelfAccountKeyring, NanoPeerAccountKeyring, NanoRoomKeyring
from clarabot_nano.common.handler.base import NanoHandleHelper
from clarabot_nano.common.zmq.requester import AioClientRequester
from clarabot_nano.gateway.config import GatewayConfig
from clarabot_nano.gateway.state.config import RoomGroupConfigState
from clarabot_nano.gateway.state.keyring import KeyringCache
from clarabot_nano.gateway.state.mutate import GatewayMutationTracking
from clarabot_nano.utility.serialize import msgpack_dumps, msgpack_loads

logger = logging.getLogger(__name__)


class NanoVersionRejected(Exception):
    pass


def _discard_task_from_map(m, k, t):
    m.pop(k, None)
    # t.exception()  # avoid logging?


class GatewayState(NanoHandleHelper):
    """
    This object holds pretty much all the state that the current connection is managing.
    """

    __WRAP_REQ_ID = 2 ** 31 - 1

    def __init__(
        self,
        ws: aiohttp.ClientWebSocketResponse,
        extracted_config: GatewayConfig,
        self_kr: NanoSelfAccountKeyring,
        core_requester: AioClientRequester,
    ):
        self.ws = ws
        self.account_id = None
        self.part_size: int = None
        self.part_approx: Tuple[int, int] = None
        self.s_hmac_parts: Tuple[int, int] = None
        self.core_requester = core_requester
        # config stores and states
        self.extracted_config = extracted_config
        self.mutate_tracker = GatewayMutationTracking()
        self.room_config_store: Dict[int, RoomGroupConfigState] = dict()
        self.keyring_cache = KeyringCache(self_kr)  # simple cache, access to it is async-managed by this object
        self._kr_peer_load_tasks: Dict[int, asyncio.Future] = dict()
        # job queues
        self._session_job_queue = asyncio.Queue()
        self._config_job_queue = asyncio.Queue()
        self._keyring_job_queue = asyncio.Queue()
        # other internal states
        self.__request_id = 1
        self.__tasks = set()
        self.__tracked_tasks = dict()  # map to support cancellation
        self.__server_futures: Dict[int, asyncio.Future] = dict()
        self.__server_responses: Dict[int, list] = defaultdict(list)

    async def initialise(self, csrf_token, continue_token) -> int:
        request_id = await self._send_websocket(
            "connection_init", {"version": __version__, "headers": {"pcsrf": csrf_token, "cont": continue_token}}
        )
        # wait for ack
        msg = await self.ws.receive()
        if msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSED, WSMsgType.CLOSING):  # FIXME: aiohttp specific codes here?
            if msg.data == 4100:
                raise NanoVersionRejected()
            raise aiohttp.ClientConnectionError(f"Close: {msg.data}")
        if msg.type != WSMsgType.BINARY:
            raise aiohttp.ClientError(f"invalid connection ack message: {msg.type} {msg.data}")
        try:
            response = msgpack_loads(msg.data)
        except Exception:
            raise aiohttp.ClientError(f"invalid connection ack message: {msg.type} {msg.data}")
        if (
            not isinstance(response, dict)
            or response.get("type") != "connection_ack"
            or response.get("id") != request_id
            or not isinstance(payload := response.get("payload"), dict)
        ):
            raise aiohttp.ClientError(f"invalid connection ack response: {response}")
        # Check sanity and consistency of the server reported account-id.
        if not isinstance(account_id := payload.get("account_id"), int):
            raise aiohttp.ClientError(f"invalid account-id in connection ack response: {response}")
        account_id_rad32 = safe_radix32.encode_fw(account_id)
        kr_account_id_rad32 = self.keyring_cache.self_kr.id.decode()
        if account_id_rad32 != kr_account_id_rad32:
            raise aiohttp.ClientError(
                f"inconsistent account-id from keyring and server: {account_id_rad32} {kr_account_id_rad32}"
            )
        # Check validity of the configs that are in the domain of the server.
        part_size = payload.get("part_size")  # Servers may adjust the desired part sizes for the cache.
        if not isinstance(part_size, int) or part_size < 50_000 or part_size > 5_000_000:
            raise aiohttp.ClientError(f"invalid part-size in connection ack response: {response}")
        self._init_from_connection_ack(account_id, part_size)
        update_state = payload.get("update_state")  # Servers will report update recommendation.
        if not isinstance(update_state, int):
            update_state = -1  # Unknown recommendation.
        return update_state

    def _init_from_connection_ack(self, account_id: int, part_size: int):
        # TODO: check and use update_state from server
        self.account_id = account_id
        # part size config is constant for the connection
        self.part_size = part_size
        part_approx_delta = part_size >> 2
        self.part_approx = [part_size - part_approx_delta, part_approx_delta]
        s_hmac_parts = math.ceil(24_000_000 / part_size)
        self.s_hmac_parts = [s_hmac_parts, s_hmac_parts << 1]

    async def close(self):
        # clean up depending tasks if any
        for t in self.__tasks:
            t.cancel()
        for t in self.__tracked_tasks.values():
            t.cancel()
        ts = list(self.__tracked_tasks.values())
        ts.extend(self.__tasks)
        await asyncio.gather(*ts, return_exceptions=True)
        # closing ws will be done by the aio session

    def ensure_bound_task(self, coroutine_or_future) -> asyncio.Future:
        future = asyncio.ensure_future(coroutine_or_future)
        future.add_done_callback(lambda t: self.__tasks.discard(t))
        self.__tasks.add(future)
        return future

    def bind_client_task(self, task: asyncio.Task, operation_id: int) -> asyncio.Task:
        """
        Ensure task cleanup before this state is closed.
        If operation_id is provided the task may be cancelled by it.
        """
        if operation_id in self.__tracked_tasks:
            raise Exception("Conflicting tracked task.")
        task.add_done_callback(partial(_discard_task_from_map, self.__tracked_tasks, operation_id))
        self.__tracked_tasks[operation_id] = task
        return task

    def cancel_task(self, operation_id: int):
        if t := self.__tracked_tasks.get(operation_id):
            t.cancel()

    async def send_response(self, op_id, payload):
        """
        Response to a Nano client-request.
        """
        await self.ws.send_bytes(msgpack_dumps({"id": op_id, "type": "response", "payload": payload}))

    def __next_request_id(self) -> int:
        while True:
            r_id = self.__request_id
            if self.__request_id >= self.__WRAP_REQ_ID:
                self.__request_id = 1
            else:
                self.__request_id += 1
            if r_id not in self.__server_futures:
                return r_id

    async def _send_websocket(self, typ: str, payload: Optional[dict]) -> int:
        request_id = self.__next_request_id()
        await self.ws.send_bytes(msgpack_dumps({"id": request_id, "type": typ, "payload": payload}))
        return request_id

    @asynccontextmanager
    async def server_request(self, payload: dict) -> AsyncContextManager[Sequence[Any]]:
        request_id = self.__next_request_id()
        self.__server_futures[request_id] = future = asyncio.get_event_loop().create_future()
        try:
            await self.ws.send_bytes(msgpack_dumps({"id": request_id, "type": "start", "payload": payload}))
            yield future
        finally:
            # noinspection PyAsyncCall
            self.__server_futures.pop(request_id, None)
            self.__server_responses.pop(request_id, None)
            # TODO: could support stop on server if future is not done

    def collect_response_data(self, request_id: int, payload: dict):
        self.__server_responses[request_id].append(payload)

    def complete_response_data(self, request_id: int):
        if f := self.__server_futures.get(request_id):
            try:
                f.set_result(self.__server_responses.get(request_id, ()))
            except asyncio.InvalidStateError:
                logger.warning("Ignoring redundant complete marker of request")

    def handle_new_extracted_config(self, extracted_config: GatewayConfig, old_extracted_config: GatewayConfig):
        """
        This helps to update the running state from config changes without requiring a reconnect to the server.
        """
        self.extracted_config = extracted_config
        if self.extracted_config.rooms != old_extracted_config.rooms:
            self._session_job_queue.put_nowait({"type": "serve-change", "serves": self.extracted_config.get_serves()})
        # Notify the room-config manager task of the verified block changes.
        old_room_blocks = old_extracted_config.room_blocks
        room_blocks = extracted_config.room_blocks
        room_block_keys = set(old_extracted_config.room_blocks.keys())
        room_block_keys.update(extracted_config.room_blocks.keys())
        for room_id_rad32 in room_block_keys:
            if old_room_blocks.get(room_id_rad32) != room_blocks.get(room_id_rad32):
                self._config_job_queue.put_nowait(({"type": "verified", "room_id_rad32": room_id_rad32}))

    # keyring access api
    async def get_peer_account_keyring(self, account_id: int) -> NanoPeerAccountKeyring:
        kr = self.keyring_cache.get_peer_kr(account_id)
        if kr:
            return kr
        return await asyncio.shield(self._get_peer_account_keyring_task(account_id))

    def _get_peer_account_keyring_task(self, account_id: int) -> asyncio.Task:
        task = self._kr_peer_load_tasks.get(account_id)
        if not task:
            self._kr_peer_load_tasks[account_id] = task = asyncio.create_task(
                self._load_peer_account_keyring(account_id)
            )
            task.add_done_callback(partial(_discard_task_from_map, self._kr_peer_load_tasks, account_id))
        return task

    async def _load_peer_account_keyring(self, account_id: int):
        f = asyncio.get_running_loop().create_future()
        self._keyring_job_queue.put_nowait({"type": "load-peer", "id": account_id, "future": f})
        return await f

    async def get_own_room_keyring(self, room_id: int) -> NanoRoomKeyring:
        kr = self.keyring_cache.get_own_room_kr(room_id)
        if kr:
            return kr
        return await asyncio.shield(self._get_room_keyring_task(room_id))

    def _get_room_keyring_task(self, room_id: int) -> asyncio.Task:
        task = self._kr_peer_load_tasks.get(room_id)
        if not task:
            self._kr_peer_load_tasks[room_id] = task = asyncio.create_task(self._load_room_keyring(room_id))
            task.add_done_callback(partial(_discard_task_from_map, self._kr_peer_load_tasks, room_id))
        return task

    async def _load_room_keyring(self, room_id: int):
        f = asyncio.get_running_loop().create_future()
        self._keyring_job_queue.put_nowait({"type": "load-room", "id": room_id, "future": f})
        return await f

    # room config utility api
    def set_served_rooms(self, served: Set[int]):
        """
        Updates the store of the group-configs of served rooms. Also produces events to load the new rooms.
        """
        old = []
        for room_id in self.room_config_store.keys():
            try:
                served.remove(room_id)
            except KeyError:
                old.append(room_id)
        for room_id in old:
            self.room_config_store.pop(room_id).close()  # cleanup waiters
        for room_id in served:
            self.room_config_store[room_id] = RoomGroupConfigState(room_id)
            self._config_job_queue.put_nowait({"type": "load-chain", "room_id": room_id})

    def get_verified_room_blocks(self, room_id: int) -> List[str]:
        verified_room_blocks = self.extracted_config.room_blocks.get(safe_radix32.encode_fw(room_id))
        if not verified_room_blocks:
            return []
        return sorted(verified_room_blocks)

    # event handlers
    def handle_session_change(self, status: str):
        self._session_job_queue.put_nowait({"type": "session-change", "status": status})

    def handle_config_change(self, room_id: int, block: bytes, signature: bytes):
        self._config_job_queue.put_nowait(
            {"type": "add-block", "room_id": room_id, "block": block, "signature": signature}
        )

    def handle_serves_change(self, room_ids: Set[int]):
        self._config_job_queue.put_nowait({"type": "serves-change", "room_ids": set(room_ids)})

    def handle_trust_change(self, account_id: int, trust: bytes):
        self._keyring_job_queue.put_nowait({"type": "trust-peer", "id": account_id, "trust": trust})
