import asyncio
import base64
import html
import json
import logging
import os
import weakref
from contextlib import asynccontextmanager
from functools import partial
from pathlib import Path
from typing import Optional, List, Tuple
from urllib.parse import urlencode

import pytest
import safe_radix32
from aiohttp import web, WSCloseCode, WSMessage, WSMsgType
from clarabot_crypto.keyring.account_base import AbstractSelfAccountKeyring, AbstractPeerAccountKeyring
from clarabot_crypto.keyring.room_base import AbstractRoomKeyring
from clarabot_nano_common.utility.path.common import normalize_absolute_path
from clarabot_nano_common.utility.serialize import msgpack_dumps

from clarabot_nano.utility import parse_environment_flag
from clarabot_nano_tests.server.document.resources import DOCUMENT_TEST_FILES, SERVER_APP_TEST_HTML
from clarabot_nano_tests.testing.client.gql_ws import GraphQlWebSocketClient
from clarabot_nano_tests.testing.client.manager import GraphQLClientManager
from clarabot_nano_tests.testing.config import room_blocks_from_room_data
from clarabot_nano_tests.testing.core.common import clear_nano_slots
from clarabot_nano_tests.testing.core.instance import ConfigTestingNanoInstanceHandler
from clarabot_nano_tests.testing.core.manager import DevCoreProcessManager
from clarabot_nano_tests.testing.document import ChunkingDocumentWsSessionTester
from clarabot_nano_tests.testing.file import temp_path
from clarabot_nano_tests.testing.gateway import do_gateway_login
from clarabot_nano_tests.testing.server.manager import DevServerManager

logger = logging.getLogger(__name__)


# class DirectDocumentWsSessionTester(AbstractDocumentWsSessionTester):


class DevDocumentClientServer:
    def __init__(
        self,
        port: int,
        wopi_port: int,
        cool_port: int,
        part_size: int,
        testing_client_manager: GraphQLClientManager,
        room_id: int,
        req_wsc: GraphQlWebSocketClient,
        req_kr: AbstractSelfAccountKeyring,
        nano_kr: AbstractPeerAccountKeyring,
        room_kr: Optional[AbstractRoomKeyring] = None,
    ):
        self.port = port
        self.wopi_port = wopi_port
        self.cool_port = cool_port
        self.testing_client_manager = testing_client_manager
        self.part_size = part_size
        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

    async def _app_on_shutdown(self, app):
        for ws in set(app["websockets"]):
            await ws.close(code=WSCloseCode.GOING_AWAY, message="Server shutdown")

    @asynccontextmanager
    async def run(self, wsd_version_hash: str, file_preps: List[Tuple[str, str, str]]):
        self.direct_url_srcs = dict()
        for name, file_id, access_token in file_preps:
            wopi_src_url = ChunkingDocumentWsSessionTester.create_wopi_src_url(file_id)
            direct_url = (
                f"http://127.0.0.1:{self.port}/browser/{wsd_version_hash}/cool.html"
                f"?{urlencode({'WOPISrc': wopi_src_url})}"
            )
            self.direct_url_srcs[name] = f"{direct_url}#at!{access_token}"

        browser_path = Path(os.path.expanduser(f"~/.clarabot-doc/browser/{wsd_version_hash}"))
        assert browser_path.is_dir(), str(browser_path)
        app = web.Application()
        app.add_routes(
            [
                # helpers
                web.get("/app.html", self._handle_app),
                web.static(f"/browser/{wsd_version_hash}", browser_path),
                # For testing with "direct/native" client connection.
                web.get("/cool/{url}/ws", self._handle_cool_ws),
                # For testing a client connection with nano protocol.
                web.get("/nano/cool/ws", self._handle_nano_ws),
            ]
        )
        app["websockets"] = weakref.WeakSet()
        app.on_shutdown.append(self._app_on_shutdown)
        runner = web.AppRunner(app)
        await runner.setup()
        site = web.TCPSite(runner, "127.0.0.1", self.port)
        await site.start()
        try:
            yield self
        finally:
            await runner.cleanup()

    async def _handle_app(self, request: web.Request):
        native_links = []
        direct_links = []
        frame_links = []
        for name, url in self.direct_url_srcs.items():
            direct_url = url.replace(
                "https%3A%2F%2Fwopi.clarabot.com%2F", f"http%3A%2F%2F127.0.0.1%3A{self.wopi_port}%2F", 1
            )
            native_url = direct_url.replace(f"http://127.0.0.1:{self.port}", f"http://127.0.0.1:{self.cool_port}", 1)
            native_links.append(f'<a href="{html.escape(native_url)}">{html.escape(name)}<a/><br/>')
            direct_links.append(f'<a href="{html.escape(direct_url)}">{html.escape(name)}<a/><br/>')
            frame_links.append(
                f'<a href="#" onclick="load_frame(this, \'{html.escape(url)}\')">{html.escape(name)}<a/><br/>'
            )
        text = SERVER_APP_TEST_HTML.read_text()
        text = text.replace("%STAGE_1%", "\n    ".join(native_links), 1)
        text = text.replace("%STAGE_2%", "\n    ".join(direct_links), 1)
        text = text.replace("%STAGE_3%", "\n    ".join(frame_links), 1)
        return web.Response(text=text, content_type="text/html")

    async def _handle_cool_ws(self, request: web.Request):
        ws = web.WebSocketResponse()
        try:
            await ws.prepare(request)
            request.app["websockets"].add(ws)
            await self._run_cool_ws(request, ws)
        except Exception:
            logger.exception("ws-error")
            raise
        finally:
            request.app["websockets"].discard(ws)
        return ws

    async def _run_cool_ws(self, request: web.Request, ws: web.WebSocketResponse):
        wsd_ws_url = str(request.url)
        assert wsd_ws_url.startswith("http"), str(wsd_ws_url)
        wsd_ws_url = "ws" + wsd_ws_url[4:]
        async with ChunkingDocumentWsSessionTester(
            self.part_size,
            self.testing_client_manager,
            self.room_id,
            self.req_wsc,
            self.req_kr,
            self.nano_kr,
            self.room_kr,
        ).run(wsd_ws_url) as wst0:
            wst0: ChunkingDocumentWsSessionTester

            async def _read():  # from server to client
                try:
                    while True:
                        msg = await wst0.read_queue.get()
                        await ws.send_bytes(msg)
                except Exception as e:
                    logger.exception("_ClientRead")
                    raise e

            async def _write():  # from client to server
                try:
                    async for msg in ws:
                        msg: WSMessage
                        if msg.type == WSMsgType.TEXT:
                            wst0.send(msg.data.encode("utf-8"))
                        elif msg.type == WSMsgType.BINARY:
                            wst0.send(msg.data)
                        elif msg.type == WSMsgType.ERROR:
                            logger.error("client ws connection error %s %s", msg, ws.exception())
                            break
                        else:
                            logger.error("client ws connection unexpected message %s", msg)
                            break
                except Exception as e:
                    logger.exception("_ClientWrite")
                    raise e

            rt = asyncio.create_task(_read())
            wt = asyncio.create_task(_write())
            try:
                await asyncio.wait([rt, wt], return_when=asyncio.FIRST_COMPLETED)
            finally:
                rt.cancel()
                wt.cancel()
                await asyncio.gather(rt, wt, return_exceptions=True)

    async def _handle_nano_ws(self, request: web.Request):
        ws = web.WebSocketResponse()
        try:
            await ws.prepare(request)
            request.app["websockets"].add(ws)
            await self._run_nano_ws(request, ws)
        except Exception:
            logger.exception("ws-error")
            raise
        finally:
            request.app["websockets"].discard(ws)
        return ws

    async def _run_nano_ws(self, request: web.Request, ws: web.WebSocketResponse):
        response_queue = asyncio.Queue()
        readers = dict()

        def _read_response(session, i, response):
            assert isinstance(response.get("response", {}).get("messages"), list), str(response)
            msgs = [[a, b, c, d.hex()] for a, b, c, d in response["response"]["messages"]]
            if msgs:
                response_queue.put_nowait(json.dumps({"session": session, "msgs": msgs}))

        async def _read():  # from client to server
            try:
                async for msg in ws:
                    msg: WSMessage
                    if msg.type in (WSMsgType.TEXT, WSMsgType.BINARY):
                        data = msg.json()
                    elif msg.type == WSMsgType.ERROR:
                        logger.error("client ws connection error %s %s", msg, ws.exception())
                        break
                    else:
                        logger.error("client ws connection unexpected message %s", msg)
                        break
                    cmd = data["cmd"]
                    if cmd == "open":
                        wsd_ws_url = data["url"]
                        try:
                            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,
                            )
                            readers[session] = asyncio.get_running_loop().create_task(
                                self.testing_client_manager.document_ws_read(
                                    session,
                                    partial(_read_response, session),
                                    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:
                            response_queue.put_nowait(json.dumps({"key": data["key"], "error": str(e)}))
                        else:
                            response_queue.put_nowait(json.dumps({"key": data["key"], "session": session}))
                    elif cmd == "send":
                        try:
                            await self.testing_client_manager.document_ws_write(
                                data["session"],
                                [
                                    (a, b, c, bytes.fromhex(d[2:] if d.startswith("0x") else d))
                                    for a, b, c, d in data["msgs"]
                                ],
                                self.room_id,
                                self.req_wsc,
                                self.req_kr,
                                self.nano_kr,
                                self.room_kr,
                            )
                        except Exception as e:
                            logger.exception("asdddd %s", data)
                            response_queue.put_nowait(json.dumps({"session": data.get("session"), "error": str(e)}))
                    elif cmd == "close":
                        try:
                            await self.testing_client_manager.document_ws_close(
                                data["session"],
                                self.room_id,
                                "ws-mut",
                                self.req_wsc,
                                self.req_kr,
                                self.nano_kr,
                                self.room_kr,
                            )
                        except Exception as e:
                            response_queue.put_nowait(json.dumps({"session": data.get("session"), "error": str(e)}))
                        else:
                            rt_ = readers.pop(data["session"])
                            await asyncio.gather(rt_, return_exceptions=True)
                            response_queue.put_nowait(json.dumps({"session": data.get("session"), "close": True}))
                    else:
                        raise Exception(f"Unexpected cmd {data}")

            except Exception as e:
                logger.exception("_ProtocolRead")
                raise e

        async def _write():
            try:
                while True:
                    msg = await response_queue.get()
                    await ws.send_str(msg)
            except Exception as e:
                logger.exception("_ProtocolWrite")
                raise e

        rt = asyncio.create_task(_read())
        wt = asyncio.create_task(_write())
        try:
            await asyncio.wait([rt, wt], return_when=asyncio.FIRST_COMPLETED)
        finally:
            rt.cancel()
            wt.cancel()
            await asyncio.gather(rt, wt, return_exceptions=True)


# TODO: Could automatize the functionality checks using selenium of documents loading?
@pytest.mark.skipif(True, reason="This can only be used for interactive integration testing and development for now.")
@pytest.mark.skipif(parse_environment_flag("CBNANO_TEST_WITHOUT_SERVER"), reason="server-dependant tests disabled")
@pytest.mark.asyncio
async def test_server_document_client(
    tmpdir,
    testing_server_manager: DevServerManager,
    reset_user_app_dirs,
    dev_nano_root_dir,
    testing_client_manager: GraphQLClientManager,
    testing_server_user0,
):
    u0_id = testing_server_user0["account_id"]
    u0_id_sr32 = safe_radix32.encode_fw(u0_id)
    room0 = testing_server_user0["rooms"][0]
    room0_id = room0["id"]
    room0_sr32 = safe_radix32.encode_fw(room0_id)
    room0_kr = room0["_room_kr"]
    await testing_server_manager.call_nano_clear_sessions([room0_sr32])
    await clear_nano_slots(testing_server_manager, [u0_id])

    r0 = Path(normalize_absolute_path(temp_path(tmpdir)))
    r0.mkdir()
    for fp in DOCUMENT_TEST_FILES:
        (r0 / fp.name).write_bytes(fp.read_bytes())

    u0_kr = testing_server_user0["_acc_kr"]
    u0wsc = await testing_client_manager.create_logged_in_wsc(testing_server_user0)
    assert not u0wsc.ws.closed

    wopi_port = 12151
    cool_port = 9980  # aka WSD

    def _mut_loc_conf(c_):
        c_["cool_wopi_port"] = wopi_port
        c_["cool_server_port"] = cool_port

    async with ConfigTestingNanoInstanceHandler(
        reset_user_app_dirs,
        dev_nano_root_dir,
        testing_server_manager,
        local_config_mut=_mut_loc_conf,
    ).run(with_document_server=True) as (pm, debug_info, session, dbg_config):
        pm: DevCoreProcessManager

        await asyncio.sleep(0.25)

        dbg_state = (await pm.get_wait_dbg_config())["state"]
        gateway_port = dbg_state["gateway_port"]
        gateway_key = dbg_state["gateway_key"]
        # successful login
        session_id = await do_gateway_login(
            testing_server_manager, pm, debug_info, gateway_port, gateway_key, testing_server_user0
        )
        await testing_server_manager.wait_nano_bind(u0_id, {session_id: set()})
        await pm.wait_dbg_states(
            [
                ("gateway_port", gateway_port),
                ("gateway_key", gateway_key),
                ("gateway_status", "active"),
            ],
            timeout=15.0,
        )
        assert session_id == pm.get_session_id_from_continue_token()
        await testing_server_manager.wait_nano_bind(u0_id, {session_id: set()})
        nano_id = await testing_server_manager.get_nano_id_by_nano_session_id(u0_id, session_id)
        assert nano_id is not None

        # setup admin pw
        a_pwd = base64.b64encode(os.urandom(9)).decode()
        await testing_client_manager.nano_request_response(
            msgpack_dumps(
                await testing_client_manager.prepare_request_with_mutation_id(
                    {"type": "A-CH-PWD", "new_password": a_pwd},
                    "dbg0",
                    None,
                    nano_id,
                    u0wsc,
                    u0_kr,
                    u0_kr,
                )
            ),
            None,
            nano_id,
            u0wsc,
            u0_kr,
            u0_kr,
        )
        a_sess = await testing_client_manager.new_admin_session_request(
            a_pwd, "dbg0", None, nano_id, u0wsc, u0_kr, u0_kr
        )
        await testing_client_manager.single_response_admin_request(
            {"type": "A-D-CREATE", "path": str(r0)}, "dbg0", a_sess, None, nano_id, u0wsc, u0_kr, u0_kr
        )
        await testing_client_manager.single_response_admin_request(
            {
                "type": "A-D-ATTACH",
                "path": str(r0),
                "target_room_id": room0_sr32,
                "room_blocks": room_blocks_from_room_data(room0),
            },
            "dbg0",
            a_sess,
            None,
            nano_id,
            u0wsc,
            u0_kr,
            u0_kr,
        )
        await asyncio.sleep(0.33)  # a little time is needed so serve advertisements become available for routing

        file_preps = []
        for fp in DOCUMENT_TEST_FILES:
            wsd_ver, f_id, f_at = await testing_client_manager.document_ws_prepare(
                fp.name, room0_id, u0wsc, u0_kr, u0_kr
            )
            file_preps.append((fp.name, f_id, f_at))

        part_size = 500_000  # TODO: check from context
        async with DevDocumentClientServer(
            9981, wopi_port, cool_port, part_size, testing_client_manager, room0_id, u0wsc, u0_kr, u0_kr
        ).run(wsd_ver, file_preps) as dds:
            print(f"Visit http://127.0.0.1:{dds.port}/app.html")
            while True:
                await asyncio.sleep(5.0)
