import asyncio
import base64
import json
import logging
import os
import time
from collections import deque
from functools import partial
from typing import Dict, List, Tuple, Deque, Optional, Generator
from urllib.parse import unquote, quote, urlparse, parse_qsl, urlunparse, urlencode

import aiohttp
import wait_for2
from aiohttp import WSMsgType, FormData
from clarabot_nano_common.utility.serialize import ulong2bytes
from sortedcontainers import SortedList

from clarabot_nano.common.handler.stream import ServerResponsePartSerializationHelper
from clarabot_nano.server.document.config import DocumentServerConfig

logger = logging.getLogger(__name__)


class DocumentWsMessage:
    __slots__ = ("delay_until", "n", "data", "i")

    def __init__(self, delay_until: float, n: int, data: bytes):
        self.delay_until = delay_until
        self.n = n
        self.data = data
        self.i = 0

    def chunk(self, size_limit: int) -> Tuple[bool, Tuple[int, int, int, bytes]]:
        from_ = self.i
        until = from_ + size_limit
        chunk = self.data[from_:until]
        data_len = len(self.data)
        self.i = until
        return until >= data_len, (self.n, from_, data_len, chunk)


class DocumentWsCommand:
    __slots__ = ("size", "buf", "buf_size", "wait_handle")

    def __init__(self, size: int):
        self.size = size
        self.buf = SortedList(key=lambda x: x[0])
        self.buf_size = 0
        self.wait_handle = None

    def close(self):
        if self.wait_handle is not None:
            self.wait_handle.cancel()

    @property
    def complete(self) -> bool:
        return self.size == self.buf_size

    def get_data(self) -> bytes:
        assert self.complete
        self.close()
        return b"".join(x[1] for x in self.buf)

    def add_chunk(self, offset: int, size: int, data: bytes) -> int:
        if size != self.size:
            raise Exception("Inconsistent document command size")
        if offset < 0 or offset + len(data) > size:
            raise Exception("Invalid document command offset and data length relative to size")
        item = (offset, data)
        i = self.buf.bisect_left(item)
        if i < len(self.buf):
            if item[0] + len(item[1]) > self.buf[i][0]:
                raise Exception("Conflicting document command chunk overlaps following part")
        if i > 0:
            i -= 1
            if item[0] < self.buf[i][0] + len(self.buf[i][1]):
                raise Exception("Conflicting document command chunk overlaps previous part")
        self.buf.add(item)
        self.buf_size += len(item[1])
        return len(item[1])


# NOTE: editor-api proxy handlers must validate the consistency of the current authorization context and the
# access_token that's passed to the wsd service. Failing this would allow the pass-the-token vulnerability to
# be exploited, and lead to potential privilege escalation in a room.
# The websocket supports access_token reset, this command must be checked in-flight to be consistent with the
# context as well.
# To ensure that the above are enough for a specific version of wsd to be secured check all the following:
#   1. Ensure that the Authorization class at common/Authorization handles its fields (_type and _data) privately.
#      They are only set at creation and can only be changed with the resetAccessToken() method.
#   2. Ensure that the Authorization object gets created along with ClientSession creation in wsd/ClientSession only:
#       - _auth(Authorization::create(uriPublic)) only called there by ClientSession ctor
#       - no rouge instances are made with "new Authorization(...)"
#   3. Ensure that the Authorization::resetAccessToken method is only ever called from the ClientSession input
#      processing for the request "resetaccesstoken":
#       _auth.resetAccessToken(tokens[1]);
#   4. Ensure that the message is tokenized according to our regexp in ClientSession::_handleInput.
#   5. Ensure that WopiStorage from wsd/WopiStorage implements the WOPI calls with an Authorization object from
#      using Authorization::authorizeURI and Authorization::authorizeRequest only:
#       auth.authorizeURI(uriObject);
#       auth.authorizeRequest(request);
# TBD:
#   - we could omit the access_token from the user side and always only pass to the wsd the current user token
#     from the context
#       -> Makes pass the token impossible, since there is never a token
#       -> Could make token timeout/auth behaviour issues because the wsd-user-editor may expect to handle it
#       -> This could make the user side path (from url) more useful and not a token for security reasons
#
# HTTP request wrapping
# The client has a few places that interacts with the server in relevant ways to the proxying:
#   - src/map/Clipboard uses XMLHttpRequest/fetch to get/post clipboard data from server (maybe for big selections?)
#   - src/map/handler/Map.FileInserter uses XMLHttpRequest to upload files (embedding images for example)
#   - src/layer/tile/CanvasTileLayer governs print, slideshow and export operation starts
#     These can be delegated to be handled by the integrating application by setting DownloadAsPostMessage.
#     Otherwise, different facilities take over and load the relevant URL in DOM or in new window/tab.
# The first two needs client side patching of the XMLHttpRequest object to send the messages through the WS, while
# the 3rd shall delegate to the integrating app and get the resources through the WS itself.
# We need to introduce custom messages, that can be separated from the WSD communication. These should be processed
# parallel among the normal requests. This makes it possible to move these transfers onto the secondary ws. Furthermore
# we need a utility that transforms the FormData (POST) into a binary serialization.
# The custom HTTP proxy message shall be "nanohttpproxy <num>\n<bin>". This conforms to the wsd protocol, so even if
# such messages were to be received, they will be parsed and discarded correctly.
#   For the request, the binary part shall be
#       {"method": "GET", "url": "http...", ...}
#   or
#       {"method": "POST", "url": "http...", "fields": {...}, ...}
#       with N optional fields of (custom segmentation needed for binary data)
#           \n{"name": "data-name", "len": M, ...}\n<data>
#   For the response, the binary part shall be
#       {"status": 200, "content_type"?: "text/plain", ...}\n<data>
#   or
#       {"status"?: 400, "error"?: "x", ...}\n
class DocumentWsSession:
    TOUCH_TIMEOUT = 10.0
    TOUCH_UPDATE_THRESHOLD = 1.0
    COMMAND_MAX_WAIT = 8.0
    COMMAND_BUFFER_SIZE_LIMIT = 4000000
    STREAM_MESSAGE_MAX_DELAY = 0.002
    STREAM_RESPONSE_MAX_DELAY = 2.5
    STREAM_BUFFER_SIZE_LIMIT = 2000000

    @classmethod
    def wsd_command_tokenizer(cls, data: bytes) -> Generator[bytes, None, None]:
        token_start = 0
        data_len = len(data)
        while True:
            while True:  # go to token start
                if token_start >= data_len or data[token_start] == 10:  # 10 is \n
                    return
                if data[token_start] == 32:  # 32 is "space"
                    token_start += 1
                else:
                    break
            token_end = token_start + 1
            while token_end < data_len:  # go to token end
                if data[token_end] == 10 or data[token_end] == 32:
                    break
                token_end += 1
            yield data[token_start:token_end]
            token_start = token_end

    __slots__ = (
        "_loop",
        "_s",
        "extracted_config",
        "session_key",
        "account_id",
        "room_id",
        "can_write",
        "secret_key",
        "_wopi_src",
        "work_task",
        "_touch_t",
        "_touch_handle",
        "_stream_buffer_n",
        "_stream_buffer_size",
        "_stream_buffer",
        "_stream_cond",
        "_stream_ser",
        "_command_n",
        "_command_buffer",
        "_command_buffer_size",
        "_command_lock",
        "_command_queue",
    )

    def __init__(
        self,
        s: aiohttp.ClientSession,
        extracted_config: DocumentServerConfig,
        session_key: str,
        account_id: int,
        room_id: int,
        can_write: bool,
        connected: asyncio.Event,
        ws_url: str,
    ):
        self._loop = asyncio.get_running_loop()
        self._s = s
        self.extracted_config = extracted_config
        self.session_key = session_key
        self.account_id = account_id
        self.room_id = room_id
        self.can_write = can_write
        self.secret_key = os.urandom(8)  # used for unique transfer-key of stream for example
        self._wopi_src = unquote(dict(parse_qsl(urlparse(ws_url).query))["WOPISrc"])  # stored for translation checks
        self.work_task = self._loop.create_task(self._work(connected, ws_url))
        # abandon/auto-cleanup
        self._touch_t = time.perf_counter() + self.TOUCH_UPDATE_THRESHOLD
        self._touch_handle = self._loop.call_later(self.TOUCH_TIMEOUT, self._handle_touch_timeout)
        # proxy from document
        self._stream_buffer_n = 0
        self._stream_buffer_size = 0
        self._stream_buffer: Deque[DocumentWsMessage] = deque()
        self._stream_cond = asyncio.Condition()
        self._stream_ser = ServerResponsePartSerializationHelper("document session", self.TOUCH_TIMEOUT)
        # proxy to document
        self._command_n = 0
        self._command_buffer: Dict[int, DocumentWsCommand] = dict()
        self._command_buffer_size = 0
        self._command_lock = asyncio.Lock()
        self._command_queue = asyncio.Queue()

    async def get_wsd_major_version(self) -> int:
        async with self._s.get(
            f"http://127.0.0.1:{self.extracted_config.cool_server_port}/hosting/capabilities"
        ) as resp:
            resp.raise_for_status()
            capabilities = await resp.json()
            version = capabilities["productVersion"]
            if isinstance(version, str):
                try:
                    return int(version.partition(".")[0])
                except Exception:
                    pass
        return 0

    async def _work(self, connected: asyncio.Event, ws_url: str):
        wsd_version = await self.get_wsd_major_version()
        async with self._s.ws_connect(ws_url) as ws:
            connected.set()  # signal the manager that the session has started
            del connected  # free
            del ws_url  # free

            http_request_tasks = set()

            async def _http_request(num: bytes, msg: bytes):
                try:
                    op_end = msg.find(b"\n")
                    if op_end == -1:
                        raise Exception("Invalid op-sep")
                    head_end = msg.find(b"\n", op_end + 1)
                    if head_end == -1:
                        head_end = len(msg)
                    head = json.loads(msg[op_end + 1 : head_end])
                    if not isinstance(head, dict):
                        raise Exception("Invalid head type")
                    method = head.get("method")
                    if method not in ("GET", "POST"):
                        raise Exception("Invalid method")
                    url = head.get("url")
                    if not isinstance(url, str):
                        raise Exception("Invalid url")
                    url = self._translate_http_request_url(url)
                    filename = None
                    if method == "GET":
                        if "fields" in head or head_end != len(msg):
                            raise Exception("Fields are invalid for GET")
                        async with self._s.get(url) as resp:
                            cd = resp.content_disposition
                            if cd and cd.filename:
                                filename = cd.filename
                            status = resp.status
                            content_type = resp.content_type
                            body = await resp.read()
                    else:
                        fields = head.get("fields")
                        if not isinstance(fields, dict):
                            raise Exception("Invalid fields type")
                        data = FormData()
                        for k, v in fields.items():
                            data.add_field(k, v)
                        prev_end = head_end
                        while True:
                            if prev_end >= len(msg):
                                break
                            if msg[prev_end] != 10:  # not newline
                                raise Exception("Malformed bin-fields separator")
                            bin_head_end = msg.find(b"\n", prev_end + 1)
                            if bin_head_end == -1:
                                raise Exception("Malformed bin-fields inner separator")
                            bin_head = json.loads(msg[prev_end + 1 : bin_head_end])
                            if not isinstance(bin_head, dict):
                                raise Exception("Invalid bin-head type")
                            bin_name = bin_head.get("name")
                            bin_len = bin_head.get("len")
                            if not isinstance(bin_name, str) or not isinstance(bin_len, int):
                                raise Exception("Invalid bin-head fields")
                            bin_end = bin_head_end + 1 + bin_len
                            if bin_end > len(msg):
                                raise Exception("Invalid bin-len value")
                            data.add_field(bin_name, msg[bin_head_end + 1 : bin_end])
                            prev_end = bin_end
                        async with self._s.post(url, data=data) as resp:
                            cd = resp.content_disposition
                            if cd and cd.filename:
                                filename = cd.filename
                            status = resp.status
                            content_type = resp.content_type
                            body = await resp.read()
                    resp_head = json.dumps(
                        {"status": status, "content_type": content_type, "filename": filename}
                    ).encode("utf-8")
                    msg = b"".join((b"nanohttpproxy ", num, b"\n", resp_head, b"\n", body))
                    await self._push_message_to_client(msg)
                except Exception:
                    logger.exception("Document WS session HTTP proxy error for %s", msg[:250])
                    msg = b"".join((b"nanohttpproxy ", num, b'\n{"status": 500, "error": "internal"}\n'))
                    await self._push_message_to_client(msg)
                    raise

            def _cleanup_http_request(t: asyncio.Task):
                http_request_tasks.discard(t)
                if not t.cancelled() and not t.exception():
                    t.result()

            async def _reader():
                try:
                    while True:
                        ws_msg = await ws.receive()
                        if ws_msg.type == WSMsgType.TEXT:
                            msg = ws_msg.data.encode("utf-8")
                            prefix = b"nanowstext " if wsd_version >= 24 else None
                        elif ws_msg.type == WSMsgType.BINARY:
                            msg = ws_msg.data
                            prefix = None
                        else:
                            if ws_msg.type == WSMsgType.ERROR:
                                logger.error("Document WS session work reader error %s", ws.exception())
                            elif ws_msg.type == WSMsgType.CLOSE:
                                logger.info("Document WS session work reader close")
                            else:
                                logger.error("Document WS session work reader exit %s", ws_msg.type)
                            break
                        # To debug wsd -> client messages:
                        # print("_WSD2C", msg[: (x if (x := msg.find(10)) >= 0 else len(msg))])
                        msg_tokenizer = iter(self.wsd_command_tokenizer(msg))
                        try:
                            cmd = next(msg_tokenizer)
                        except StopIteration:
                            logger.error("Document WS session push failed to tokenize: %s", msg[:100])
                            break
                        if cmd == b"saveas:":
                            for token in msg_tokenizer:
                                if token.startswith(b"url="):
                                    # Strangely the save-as response does not quote the url, so we do it adaptively.
                                    quoted = token.startswith(b"url=http%") or token.startswith(b"url=https%")
                                    url = token[4:].decode("utf-8")
                                    if quoted:
                                        url = unquote(url)
                                    url = self.extracted_config.translate_from_wopi_host(url)
                                    if quoted:
                                        url = quote(url, safe="")
                                    msg = msg.replace(token, b"url=" + url.encode("utf-8"), 1)
                                    break
                            else:
                                logger.error("Document WS session push %s has not url", cmd)
                                break
                        # NOTE: Other commands that are related to URL translation in some aspect:
                        #   - downloadas: Response has an id that the client constructs a request by.
                        if prefix:  # We tell the client that this message must be decoded into text with a prefix.
                            msg = prefix + msg
                        await self._push_message_to_client(msg)
                except Exception:
                    logger.exception("Document WS session work reader failed")
                    raise

            async def _writer():
                try:
                    while True:
                        msg = await self._command_queue.get()
                        msg_tokenizer = iter(self.wsd_command_tokenizer(msg))
                        try:
                            cmd = next(msg_tokenizer)
                        except StopIteration:
                            logger.error("Document WS session command failed to tokenize: %s", msg[:100])
                            break
                        # To debug client -> wsd messages:
                        # print("_C2WSD", msg[: (x if (x := msg.find(10)) >= 0 else len(msg))])
                        if cmd == b"nanohttpproxy":
                            tokens = list(msg_tokenizer)
                            if len(tokens) != 1:
                                logger.error("Document WS session http-req failed to tokenize: %s", msg[:100])
                                break
                            t = self._loop.create_task(_http_request(tokens[0], msg))
                            t.add_done_callback(_cleanup_http_request)
                            http_request_tasks.add(t)
                            continue
                        elif cmd == b"resetaccesstoken":
                            # Security check of the "resetaccesstoken" command.
                            # TODO: We do not implement this because it should not be used?
                            #   If it is used by PutRelativeFile response processing, then the "load" command needs
                            #   to be observed as well and allow change of file_id and access_token in-sync?
                            #     - This is not used by PutRelativeFile, a whole new socket is started instead.
                            #   Why would these not be a single command anyway? What would their order be?
                            # access_token = msg_tokens[1].decode("utf-8")
                            logger.error("Got access_token reset command for document ws, but it is not implemented")
                            break
                        elif cmd == b"load":
                            for token in msg_tokenizer:
                                if token.startswith(b"url="):
                                    msg = msg.replace(
                                        token,
                                        b"url="
                                        + quote(
                                            self.extracted_config.translate_to_wopi_host(
                                                unquote(token[4:].decode("utf-8"))
                                            ),
                                            safe="",
                                        ).encode("utf-8"),
                                        1,
                                    )
                                    break
                            else:
                                logger.error("Document WS session command %s has not url", cmd)
                                break
                        # NOTE: Other commands that are related to URL translation in some aspect:
                        #   - "saveas": The URL in the command is like "wopi:filename.ext", no translation needed.
                        await ws.send_bytes(msg)
                except Exception:
                    logger.exception("Document WS session work writer failed")
                    raise

            reader_t = self._loop.create_task(_reader())
            writer_t = self._loop.create_task(_writer())
            try:
                await asyncio.wait((reader_t, writer_t), return_when=asyncio.FIRST_COMPLETED)
            finally:
                reader_t.cancel()
                writer_t.cancel()
                for t in http_request_tasks:
                    t.cancel()
                await asyncio.gather(reader_t, writer_t, *http_request_tasks, return_exceptions=True)

    async def _push_message_to_client(self, msg: bytes):
        async with self._stream_cond:
            while 0 < self._stream_buffer_size > self.STREAM_BUFFER_SIZE_LIMIT:
                # No timeout on wait, there is no other way to continue besides drain by client.
                await self._stream_cond.wait()
            delay_until = time.perf_counter() + self.STREAM_MESSAGE_MAX_DELAY
            self._stream_buffer.append(DocumentWsMessage(delay_until, self._stream_buffer_n, msg))
            self._stream_buffer_size += len(msg)
            self._stream_buffer_n += 1
            self._stream_cond.notify_all()

    def _translate_http_request_url(self, url: str) -> str:
        parsed = urlparse(url)
        path_parts = parsed.path.split("/")
        if not (3 <= len(path_parts) <= 5) or path_parts[0] != "" or path_parts[1] != "cool":
            raise Exception("Invalid http proxy url path")
        if len(path_parts) == 3:
            if path_parts[2] != "clipboard":
                raise Exception("Invalid http proxy clipboard path")
        else:
            if path_parts[3] == "insertfile":
                if len(path_parts) != 4:
                    raise Exception("Invalid http proxy insertfile path")
            elif path_parts[3] == "download":
                if len(path_parts) != 5:
                    raise Exception("Invalid http proxy download path")
            else:
                raise Exception(f"Invalid http proxy url path-3 {path_parts[3][:30]}")
            # validate WOPISrc in path param
            path_wopi_src = self.extracted_config.translate_to_wopi_host(unquote(path_parts[2]))
            if path_wopi_src != self._wopi_src:
                raise Exception("Inconsistent http proxy path wopi-src")
            path_parts[2] = quote(self._wopi_src, safe="")
        # validate WOPISrc in query param
        query = dict(parse_qsl(parsed.query))
        if "WOPISrc" not in query:
            raise Exception("Missing wopi-src http proxy query")
        query_wopi_src = self.extracted_config.translate_to_wopi_host(unquote(query["WOPISrc"]))
        if query_wopi_src != self._wopi_src:
            raise Exception("Inconsistent http proxy query wopi-src")
        query["WOPISrc"] = self._wopi_src
        return urlunparse(
            (
                "http",
                f"127.0.0.1:{self.extracted_config.cool_server_port}",
                "/".join(path_parts),
                "",
                urlencode(query),
                "",
            )
        )

    def touch(self):
        t = time.perf_counter()
        if self._touch_t < t:
            self._touch_t = t + self.TOUCH_UPDATE_THRESHOLD
            self._touch_handle.cancel()
            self._touch_handle = self._loop.call_later(self.TOUCH_TIMEOUT, self._handle_touch_timeout)

    def _handle_touch_timeout(self):
        logger.warning("Closing abandoned document session %s", self.session_key)
        self.work_task.cancel()

    def _handle_command_wait_timeout(self):
        logger.warning("Closing staggered document session %s", self.session_key)
        self.work_task.cancel()

    def stop(self) -> asyncio.Task:
        logger.warning("Stopping document session %s", self.session_key)
        self.work_task.cancel()
        self._touch_handle.cancel()
        return self.work_task

    # proxy api
    async def read(self, part: int, size_limit: int) -> List[Tuple[int, int, int, bytes]]:  # wsd server to client
        async with self._stream_ser.serialize_processing(part):
            msgs = []
            delay_until = time.perf_counter() + self.STREAM_RESPONSE_MAX_DELAY
            async with self._stream_cond:
                while True:
                    # firstly drain the buffer in-sync to return already delayed items quickly
                    while self._stream_buffer:
                        stream_msg = self._stream_buffer[0]
                        delay_until = min(stream_msg.delay_until, delay_until)
                        msg_done, stream_tuple = stream_msg.chunk(size_limit)
                        size_limit -= len(stream_tuple[-1])
                        size_limit -= 16  # approximate item serialization overhead (msgpack)
                        msgs.append(stream_tuple)
                        if msg_done:
                            self._stream_buffer_size -= len(self._stream_buffer.popleft().data)
                            self._stream_cond.notify_all()
                        if size_limit <= 0:
                            return msgs
                    delay_wait = delay_until - time.perf_counter()
                    if delay_wait <= 0.0:
                        return msgs
                    try:
                        await wait_for2.wait_for(self._stream_cond.wait(), delay_wait)
                    except asyncio.TimeoutError:
                        pass  # re-acquires the cond-lock, we can quickly check the deque and not return asap

    async def write(self, msgs: List[Tuple[int, int, int, bytes]]) -> None:  # client to wsd server
        try:
            # NOTE: The sending messages are not guaranteed to arrive in-order, so we'll have a time-window to fix them.
            async with self._command_lock:
                for msg in msgs:
                    if not isinstance(msg, (list, tuple)) or len(msg) != 4:
                        raise Exception("Document session desync by command message")
                    num, offset, size, data = msg
                    if not isinstance(num, int) or num < self._command_n:
                        raise Exception("Document session desync by command number")
                    if not isinstance(offset, int):
                        raise Exception("Document session desync by command offset")
                    if not isinstance(size, int):
                        raise Exception("Document session desync by command size")
                    if not isinstance(data, bytes):
                        raise Exception("Document session desync by command data")
                    cmd = self._command_buffer.get(num)
                    if cmd is None:
                        cmd = DocumentWsCommand(size)
                        cmd.add_chunk(offset, size, data)
                        if num == self._command_n and cmd.complete:
                            # If the command is the next one in line and already complete, we can optimize a little by
                            # avoiding add/remove it from the buffer collection.
                            await self._command_queue.put(cmd.get_data())
                            self._command_n += 1
                        else:
                            # Ensure client does not send commands far into the future as abuse or error. Also detects
                            # when a command stream may have "stuck" by loosing an intermediary message.
                            cmd.wait_handle = self._loop.call_later(
                                self.COMMAND_MAX_WAIT, self._handle_command_wait_timeout
                            )
                            self._command_buffer[num] = cmd
                            self._command_buffer_size += cmd.buf_size
                            continue  # no next completion, do not fall through to cyclic complete check
                    else:
                        self._command_buffer_size += cmd.add_chunk(offset, size, data)
                        if num == self._command_n and cmd.complete:
                            await self._command_queue.put(cmd.get_data())
                            del self._command_buffer[num]
                            self._command_buffer_size -= cmd.buf_size
                            self._command_n += 1
                        else:
                            continue  # no next completion, do not fall through to cyclic complete check
                    # reaching this point means the next command was completed, so we may be able to push more
                    while True:
                        cmd = self._command_buffer.get(self._command_n)
                        if cmd is None or not cmd.complete:
                            break
                        await self._command_queue.put(cmd.get_data())
                        del self._command_buffer[self._command_n]
                        self._command_buffer_size -= cmd.buf_size
                        self._command_n += 1
            if self._command_buffer_size > self.COMMAND_BUFFER_SIZE_LIMIT:
                raise Exception("Document session buffer limit by command")
        except Exception:
            logger.exception("Write failed to document session %s", self.session_key)
            self.stop()
            raise


class DocumentWsManager:
    __slots__ = ("_proxy_session", "_extracted_config", "_document_sessions")

    def __init__(self, s: aiohttp.ClientSession, extracted_config: DocumentServerConfig):
        self._proxy_session = s
        self._extracted_config = extracted_config
        self._document_sessions: Dict[str, DocumentWsSession] = dict()

    def _lose_document_session(self, k: str, f) -> None:
        self._document_sessions.pop(k)

    def get_document_session(
        self, session_key: str, account_id: int, room_id: int, can_write: Optional[bool], touch: bool = True
    ) -> Optional[DocumentWsSession]:
        ds = self._document_sessions.get(session_key)
        if (
            ds
            and account_id == ds.account_id
            and room_id == ds.room_id
            and (can_write is None or can_write is ds.can_write)
        ):
            if touch:
                ds.touch()
            return ds

    async def create_document_session(self, account_id: int, room_id: int, can_write: bool, ws_url: str) -> str:
        while True:
            k = base64.b64encode(os.urandom(10) + ulong2bytes(int(time.time() * 1000))).decode("ascii")
            if k not in self._document_sessions:
                break
        connected = asyncio.Event()
        ds = DocumentWsSession(
            self._proxy_session, self._extracted_config, k, account_id, room_id, can_write, connected, ws_url
        )
        ds.work_task.add_done_callback(partial(self._lose_document_session, k))
        self._document_sessions[k] = ds
        try:
            await wait_for2.wait_for(connected.wait(), 8.0)
        except (Exception, asyncio.CancelledError):
            ds.stop()
            raise
        else:
            return k

    async def close_document_sessions(self):
        while self._document_sessions:
            ds_list = list(self._document_sessions.values())
            await asyncio.gather(*[ds.stop() for ds in ds_list], return_exceptions=True)
