import asyncio
import logging
import random
import time
from collections import deque
from contextlib import asynccontextmanager
from typing import Optional, List, Tuple
from urllib.parse import quote, urlencode

from clarabot_crypto.keyring.account_base import AbstractSelfAccountKeyring, AbstractPeerAccountKeyring
from clarabot_crypto.keyring.room_base import AbstractRoomKeyring

from clarabot_nano_tests.testing.client.gql_ws import GraphQlWebSocketClient
from clarabot_nano_tests.testing.client.manager import GraphQLClientManager

logger = logging.getLogger(__name__)


class DocumentWsReader:  # read from wsd server to client
    def __init__(self):
        self.n = 0
        self.buf = []
        self.buf_size = 0
        self.buf_target = 0

    def feed(self, msgs: List[Tuple[int, int, int, bytes]]) -> List[bytes]:
        ws_msgs = []
        for msg_n, msg_offset, msg_total_len, msg_chunk in msgs:
            if msg_n != self.n:
                raise Exception(f"Ws message num inconsistent {msg_n} {self.n} {[x[0] for x in msgs]}")
            if msg_offset != self.buf_size:
                raise Exception("Ws message offset inconsistent")
            if msg_offset == 0:
                self.buf_target = msg_total_len
            elif msg_total_len != self.buf_target:
                raise Exception("Ws message len inconsistent")
            self.buf.append(msg_chunk)
            self.buf_size += len(msg_chunk)
            if self.buf_size == self.buf_target:
                ws_msgs.append(b"".join(self.buf))
                self.buf_size = 0
                del self.buf[:]
                self.n += 1
        return ws_msgs


class DocumentWsWriter:
    def __init__(self, max_delay=0.021):
        self.max_delay = max_delay
        self.n = 0
        self.buf = deque()

    def drain(self, part_size: int) -> List[List[Tuple[int, int, int, bytes]]]:
        req_msgs = []
        if self.buf and self.buf[0][0] < time.perf_counter():
            msgs = []
            msgs_size = 0
            while self.buf:
                _, ws_msg = self.buf.popleft()
                offset = 0
                size = len(ws_msg)
                while True:
                    chunk_end = min(offset + min(size - offset, part_size - msgs_size), size)
                    chunk = ws_msg[offset:chunk_end]
                    msgs.append((self.n, offset, size, chunk))
                    msgs_size += 16 + len(chunk)
                    if chunk_end == size:
                        break
                    elif msgs_size >= part_size:
                        req_msgs.append(list(msgs))
                        del msgs[:]
                        msgs_size = 0
                    offset = chunk_end
                self.n += 1
            if msgs:
                req_msgs.append(msgs)
        return req_msgs

    def feed(self, msg: bytes) -> None:
        self.buf.append((time.perf_counter() + self.max_delay, msg))


class AbstractDocumentWsSessionTester:
    def __init__(
        self,
        testing_client_manager: GraphQLClientManager,
        room_id: int,
        req_wsc: GraphQlWebSocketClient,
        req_kr: AbstractSelfAccountKeyring,
        nano_kr: AbstractPeerAccountKeyring,
        room_kr: Optional[AbstractRoomKeyring] = None,
    ):
        self.testing_client_manager = testing_client_manager
        self.room_id = room_id
        self.req_wsc = req_wsc
        self.req_kr = req_kr
        self.nano_kr = nano_kr
        self.room_kr = room_kr
        self.session = None
        self.running = False

    @classmethod
    def create_wopi_src_url(cls, file_id: str) -> str:
        return f"https://wopi.clarabot.com/wopi/files/{file_id}"

    @classmethod
    def create_wsd_ws_url(cls, file_id: str, access_token: str) -> str:
        wopi_src_url = cls.create_wopi_src_url(file_id)
        wopi_at_url = f"{wopi_src_url}?{urlencode({'access_token': access_token})}"
        wsd_ws_url = (
            f"ws://doc.clarabot.com/cool/{quote(wopi_at_url, safe='')}/ws?{urlencode({'WOPISrc': wopi_src_url})}"
        )
        return wsd_ws_url

    @asynccontextmanager
    async def run(self, wsd_ws_url: str):
        self.session = await self.testing_client_manager.document_ws_open(
            wsd_ws_url, self.room_id, "ws-mut", self.req_wsc, self.req_kr, self.nano_kr, self.room_kr
        )
        rt = asyncio.create_task(self._read())
        wt = asyncio.create_task(self._write())
        self.running = True
        try:
            yield self
        finally:
            self.running = False
            rt.cancel()
            wt.cancel()
            await asyncio.gather(rt, wt, return_exceptions=True)
            await self.testing_client_manager.document_ws_close(
                self.session, self.room_id, "ws-mut", self.req_wsc, self.req_kr, self.nano_kr, self.room_kr
            )
            self.session = None

    async def _read(self):
        try:
            reader = DocumentWsReader()

            def _read_response(i, response):
                assert isinstance(response.get("response", {}).get("messages"), list), str(response)
                for ws_msg in reader.feed(response["response"]["messages"]):
                    # print("_Reader read", len(ws_msg), ws_msg[:100])
                    self._handle_ws_msg(ws_msg)

            await self.testing_client_manager.document_ws_read(
                self.session,
                _read_response,
                self.room_id,
                self.req_wsc,
                self.req_kr,
                self.nano_kr,
                self.room_kr,
                timeout=3600,  # this needs to long-running
            )
        except Exception as e:
            logger.exception("_Reader")
            raise e
        finally:
            self.running = False

    def _handle_ws_msg(self, msg: bytes) -> None:
        raise NotImplementedError()

    async def _write(self):
        try:
            while True:
                msgs_batches = await self._get_write_msgs_batches()
                for msgs in msgs_batches:
                    # for msg in msgs:
                    #     print("_Writer send", len(msg[-1]), msg[-1][:100])
                    await self.testing_client_manager.document_ws_write(
                        self.session, msgs, self.room_id, self.req_wsc, self.req_kr, self.nano_kr, self.room_kr
                    )
                await asyncio.sleep(0.005)  # poor man's polling here... okay for dev only
        except Exception as e:
            logger.exception("_Writer")
            raise e
        finally:
            self.running = False

    async def _get_write_msgs_batches(self) -> List[List[Tuple[int, int, int, bytes]]]:
        raise NotImplementedError()


class ChunkingDocumentWsSessionTester(AbstractDocumentWsSessionTester):
    def __init__(
        self,
        part_size: int,
        testing_client_manager: GraphQLClientManager,
        room_id: int,
        req_wsc: GraphQlWebSocketClient,
        req_kr: AbstractSelfAccountKeyring,
        nano_kr: AbstractPeerAccountKeyring,
        room_kr: Optional[AbstractRoomKeyring] = None,
    ):
        AbstractDocumentWsSessionTester.__init__(
            self, testing_client_manager, room_id, req_wsc, req_kr, nano_kr, room_kr
        )
        self.part_size = part_size
        self.send_buffer = None
        self.read_queue = None
        self.do_shuffle = True

    @asynccontextmanager
    async def run(self, wsd_ws_url: str):
        self.send_buffer = DocumentWsWriter()
        self.read_queue = asyncio.Queue()
        try:
            async with AbstractDocumentWsSessionTester.run(self, wsd_ws_url) as x:
                yield x
        except Exception as e:
            raise e
        finally:
            self.read_queue = None
            self.send_buffer = None

    def send(self, msg: bytes):
        assert self.running, "not running"
        self.send_buffer.feed(msg)

    def _handle_ws_msg(self, msg: bytes) -> None:
        self.read_queue.put_nowait(msg)

    async def _get_write_msgs_batches(self) -> List[List[Tuple[int, int, int, bytes]]]:
        msgs_batches = self.send_buffer.drain(self.part_size)
        if self.do_shuffle:
            random.shuffle(msgs_batches)
        return msgs_batches


class EchoDocumentWsSessionTester(ChunkingDocumentWsSessionTester):
    def __init__(
        self,
        part_size: int,
        testing_client_manager: GraphQLClientManager,
        room_id: int,
        req_wsc: GraphQlWebSocketClient,
        req_kr: AbstractSelfAccountKeyring,
        nano_kr: AbstractPeerAccountKeyring,
        room_kr: Optional[AbstractRoomKeyring] = None,
    ):
        ChunkingDocumentWsSessionTester.__init__(
            self, part_size, testing_client_manager, room_id, req_wsc, req_kr, nano_kr, room_kr
        )
        self.sent_queue = asyncio.Queue()

    def send(self, msg: bytes, drain=True):
        ChunkingDocumentWsSessionTester.send(self, msg)
        self.sent_queue.put_nowait(msg)
        if drain:
            self.drain_queues()

    def drain_queues(self):
        while not self.read_queue.empty() and not self.sent_queue.empty():
            r = self.read_queue.get_nowait()
            s = self.sent_queue.get_nowait()
            assert r == s, f"echo mismatch\n{r}\n{s}"

    async def wait_drain_queues(self, wait=3.0):
        while True:
            assert self.running, "not running"
            self.drain_queues()
            if wait is None:
                break
            elif self.read_queue.empty() and self.sent_queue.empty():
                break
            elif wait > 0:
                wait -= 0.1
                await asyncio.sleep(0.1)
            else:
                assert False, "missing queue msg"
