from __future__ import annotations

import asyncio
import base64
import json
import time
from contextlib import contextmanager, asynccontextmanager
from pathlib import Path
from typing import ContextManager, Dict, List

import aiohttp
import psutil
from clarabot_nano_common.utility.serialize import msgpack_loads, msgpack_dumps

from clarabot_nano.common.core.message import MSG_TYPE_REQUEST_CONFIG_UPDATE, MSG_TYPE_REQUEST_STATE_UPDATE
from clarabot_nano.common.spawn.constants import ShortNames
from clarabot_nano_tests.testing.process import pypopen


class DevCoreProcessManager(object):
    @classmethod
    @contextmanager
    def startup(cls, config_dir, data_dir, module, args=None, env=None) -> ContextManager[DevCoreProcessManager]:
        popen_args = [module]
        if args:
            popen_args.extend(args)
        with pypopen(popen_args, env=env) as p:
            yield cls(p, config_dir, data_dir)

    def __init__(self, p: psutil.Popen, config_dir: Path, data_dir: Path):
        self.p = p
        self.config_dir = config_dir
        self.data_dir = data_dir

    def children(self):
        try:
            # Some plaftorms spawn intermediate processes, so recursive is needed.
            return psutil.Process(self.p.pid).children(recursive=True)
        except psutil.NoSuchProcess:
            return []

    def read_child_debug_info(self, pid):
        info_path = self.config_dir / "dbg-srv" / f"{pid}.info"
        if info_path.is_file():
            try:
                with info_path.open() as f:
                    return json.load(f)
            except Exception:
                pass  # write-read race condition shall not be hard error

    def children_debug_info(self):
        debug_info = {}
        for c in self.children():
            child_info = self.read_child_debug_info(c.pid)
            if child_info and (child_info["pid"] != c.pid or child_info["create_time"] != c.create_time()):
                child_info = None
            elif child_info:
                child_info["_p"] = c
            debug_info[c.pid] = child_info
        return debug_info

    async def wait_children(self, children_pattern: Dict[str, int], timeout=12.0) -> Dict[str, Dict[int, dict]]:
        """
        :param children_pattern: mapping of process type and number of instances
        :param timeout: time to wait
        :return:
        """
        timeout += time.perf_counter()
        children_spawns = {}
        while time.perf_counter() < timeout:
            children_spawns.clear()
            for c in self.children_debug_info().values():
                if not c:
                    continue
                if c["name"] not in children_spawns:
                    children_spawns[c["name"]] = {}
                children_spawns[c["name"]][c["pid"]] = c
            current_pattern = dict((n, len(x)) for n, x in children_spawns.items())
            for name, num in children_pattern.items():
                if name not in current_pattern and num == 0:
                    current_pattern[name] = 0
            if children_pattern == current_pattern:
                return children_spawns
            await asyncio.sleep(0.25)
        assert False, f"{children_pattern}\n{dict((k, len(v)) for k, v in children_spawns.items())}\n{children_spawns}"

    async def wait_children_config(
        self,
        aio_session: aiohttp.ClientSession,
        children_spawns: Dict[str, Dict[int, dict]],
        fixed_config=None,
        timeout=8.0,
    ) -> dict:
        """
        Wait and assert that the config written by the debug-core is the same as the one received from the debug client.
        """
        t = time.perf_counter() + timeout
        child_config = None
        debug_info = self.get_debug_info(children_spawns)
        while True:
            if fixed_config:
                config = fixed_config
            else:
                config = self.get_dbg_config()
            if not debug_info or "port" not in debug_info:
                break
            try:
                child_config = await self.get_children_config(
                    aio_session, debug_info["port"], timeout=min(1.5, max(0.0, t - time.perf_counter()))
                )
                assert child_config.pop("_local_revision") > 0
                if config == child_config:
                    break
                await asyncio.sleep(0.25)
            except asyncio.TimeoutError:
                pass
            if time.perf_counter() > t:
                assert False, f"configs are not in sync\n{fixed_config}\n{config}\n{child_config}"
        return config

    def get_dbg_config(self):
        dbg_config = self.config_dir / "dbg-config.mpk"
        if dbg_config.is_file():
            try:
                return msgpack_loads(dbg_config.read_bytes())
            except Exception:
                # There is no synchronization between reading and writing this file, so race conditions may happen
                # when decoding. These shall be ignored and the next attempt will likely succeed.
                pass

    async def get_wait_dbg_config(self, timeout=5.0):
        t = time.perf_counter() + timeout
        while True:
            config = self.get_dbg_config()
            if config is not None:
                return config
            if time.perf_counter() > t:
                assert False, f"config is not readable"
            await asyncio.sleep(0.05)

    async def get_children_config(self, aio_session, debug_port, timeout=5.0):
        async with aio_session.get(f"http://127.0.0.1:{debug_port}/get-config", timeout=timeout) as response:
            assert response.status == 200
            return msgpack_loads(await response.read())

    def get_debug_info(self, children_spawns):
        debug_spawn = children_spawns.get(ShortNames.DEBUG)
        if not debug_spawn:
            return
        return next(iter(debug_spawn.items()))[1]

    def get_session_id_from_continue_token(self):
        c_token = self.data_dir / "gateway" / ".ctoken"
        if c_token.is_file():
            try:
                token = c_token.read_text("utf-8").partition(".")[0]
                token = base64.urlsafe_b64decode(token + "===")
                return msgpack_loads(token)["s"]
            except Exception:
                # There is no synchronization between reading and writing this file, so race conditions may happen
                # when decoding. These shall be ignored and the next attempt will likely succeed.
                pass

    async def wait_dbg_states(self, states, timeout=5.0):
        timeout += time.perf_counter()
        while time.perf_counter() < timeout:
            dbg_config = self.get_dbg_config()
            if dbg_config and (state := dbg_config["state"]):
                for k, v in states:
                    if callable(v):
                        if v(state.get(k)):
                            continue
                    elif state.get(k) == v:
                        continue
                    break
                else:
                    return state
            await asyncio.sleep(0.05)
        assert False, f"{dbg_config}"

    @asynccontextmanager
    async def dbg_request(
        self, aio_session: aiohttp.ClientSession, port, path, assert_status=200, post=False, timeout=5.0, **req_kwargs
    ) -> ContextManager[aiohttp.ClientResponse]:
        url = f"http://127.0.0.1:{port}/{path}"
        if post:
            async with aio_session.post(url, timeout=timeout, **req_kwargs) as response:
                if assert_status is not None:
                    assert response.status == assert_status
                yield response
        else:
            async with aio_session.get(url, timeout=timeout, **req_kwargs) as response:
                if assert_status is not None:
                    assert response.status == assert_status
                yield response

    async def assert_request_config_update(
        self, aio_session: aiohttp.ClientSession, port: int, request: dict, ok=True, message=None
    ):
        await self.assert_core_request(aio_session, port, MSG_TYPE_REQUEST_CONFIG_UPDATE, request, ok, message)

    async def assert_request_state_update(
        self, aio_session: aiohttp.ClientSession, port: int, request: dict, ok=True, message=None
    ):
        await self.assert_core_request(aio_session, port, MSG_TYPE_REQUEST_STATE_UPDATE, request, ok, message)

    async def assert_core_request(
        self, aio_session: aiohttp.ClientSession, port: int, command: bytes, request: dict, ok=True, message=None
    ):
        if message is None:
            message = msgpack_dumps(request)
        async with self.dbg_request(
            aio_session,
            port,
            "core-request",
            post=True,
            data=msgpack_dumps({"command": command, "message": message}),
            timeout=60,
        ) as response:
            if ok is not None:
                result = msgpack_loads(await response.read())
                assert result == {"ok": ok}, f"{result}\n{request}"

    async def assert_core_do_login(self, aio_session: aiohttp.ClientSession, port: int, email: str, password: str):
        async with self.dbg_request(
            aio_session,
            port,
            "do-login",
            post=True,
            data=msgpack_dumps({"email": email, "password": password}),
            timeout=60,
        ) as response:
            return msgpack_loads(await response.read()).get("result")

    async def assert_core_do_logout(self, aio_session: aiohttp.ClientSession, port: int):
        async with self.dbg_request(aio_session, port, "do-logout", post=True, timeout=60) as response:
            return msgpack_loads(await response.read()).get("result")

    async def assert_core_do_config_reset(self, aio_session: aiohttp.ClientSession, port: int):
        async with self.dbg_request(aio_session, port, "do-config-reset", post=True, timeout=60) as response:
            return msgpack_loads(await response.read()).get("ok")

    async def assert_core_do_drive_create(self, aio_session: aiohttp.ClientSession, port: int, path: str):
        async with self.dbg_request(
            aio_session, port, "do-drive-create", post=True, data=msgpack_dumps({"path": path}), timeout=60
        ) as response:
            return msgpack_loads(await response.read())

    async def assert_core_do_drive_delete(self, aio_session: aiohttp.ClientSession, port: int, path: str):
        async with self.dbg_request(
            aio_session, port, "do-drive-delete", post=True, data=msgpack_dumps({"path": path}), timeout=60
        ) as response:
            return msgpack_loads(await response.read())

    async def assert_core_do_drive_attach_room(
        self, aio_session: aiohttp.ClientSession, port: int, path: str, room_id: str, room_blocks: List[str]
    ):
        async with self.dbg_request(
            aio_session,
            port,
            "do-drive-attach-room",
            post=True,
            data=msgpack_dumps({"path": path, "room_id": room_id, "room_blocks": room_blocks}),
            timeout=60,
        ) as response:
            return msgpack_loads(await response.read())

    async def assert_core_do_drive_detach_room(
        self, aio_session: aiohttp.ClientSession, port: int, path: str, room_id: str
    ):
        async with self.dbg_request(
            aio_session,
            port,
            "do-drive-detach-room",
            post=True,
            data=msgpack_dumps({"path": path, "room_id": room_id}),
            timeout=60,
        ) as response:
            return msgpack_loads(await response.read())

    async def assert_core_wait_idle_indexer(
        self, aio_session: aiohttp.ClientSession, port: int, idle_after: float = 1.0, timeout: float = 20.0
    ) -> bool:
        async with self.dbg_request(
            aio_session,
            port,
            "wait-idle-indexer",
            post=True,
            data=msgpack_dumps({"idle_after": idle_after, "timeout": timeout}),
            timeout=timeout * 2,
        ) as response:
            return msgpack_loads(await response.read())["idle"]
