from contextlib import asynccontextmanager
from functools import partial
from pathlib import Path
from typing import List, Optional, Callable, Sequence, Union, Dict

import aiohttp
from clarabot_nano_common.common.solr.rest import create_solr_auth_header
from clarabot_nano_common.utility.serialize import msgpack_dumps

from clarabot_nano_tests.core.utility.common import create_children_wait_map
from clarabot_nano_tests.testing.client.manager import GraphQLClientManager
from clarabot_nano_tests.testing.core.common import clear_nano_slots
from clarabot_nano_tests.testing.core.manager import DevCoreProcessManager
from clarabot_nano_tests.testing.file import change_yaml
from clarabot_nano_tests.testing.gateway import start_nano_wait_login, start_nano_wait_gateway, start_nano_wait_ready
from clarabot_nano_tests.testing.server.manager import DevServerManager
from clarabot_nano_tests.testing.solr import wait_solr_startup


class AbstractConfigurableNanoInstanceHandler(object):
    """
    Start an instance and customize its config.
    """

    def __init__(
        self,
        config_dir: Path,
        local_config_mutators: Sequence[Union[Callable[[dict], None], None]],
        remote_config_mutators: Sequence[Union[Callable[[dict], None], None]],
    ):
        self._config_dir = config_dir
        self._local_config_mutators = [mut for mut in local_config_mutators if mut is not None]
        self._remote_config_mutators = [mut for mut in remote_config_mutators if mut is not None]

    @asynccontextmanager
    async def run(self, *exec_args, **exec_kwargs):
        await self._setup()
        if self._local_config_mutators or self._remote_config_mutators:
            async with self._prepare() as prep_args:
                await self._handle_prepared(prep_args)
            await self._adjust_config()
        async with self._execute(*exec_args, **exec_kwargs) as exec_args:
            yield exec_args

    async def _setup(self):
        raise NotImplementedError()

    @asynccontextmanager
    async def _prepare(self):
        raise NotImplementedError()

    async def _handle_prepared(self, prep_args):
        raise NotImplementedError()

    async def _adjust_config(self):
        if self._local_config_mutators:
            with change_yaml(self.local_config_path()) as config:
                for mut in self._local_config_mutators:
                    mut(config)
        if self._remote_config_mutators:
            with change_yaml(self.remote_config_path()) as config:
                for mut in self._remote_config_mutators:
                    mut(config)

    def local_config_path(self):
        return self._config_dir / "local.yml"

    def remote_config_path(self):
        return self._config_dir / "remote.yml"

    @asynccontextmanager
    async def _execute(self):
        raise NotImplementedError()


class CoreTestingNanoInstanceHandler(AbstractConfigurableNanoInstanceHandler):
    """
    Create an instance where the processes are dummies except the core itself.
    """

    def __init__(
        self,
        reset_user_app_dirs,
        dev_nano_root_dir,
        remote_config_mut: Optional[Callable[[dict], None]] = None,
        local_config_mut: Optional[Callable[[dict], None]] = None,
    ):
        self.reset_user_app_dirs = reset_user_app_dirs
        self.dev_nano_root_dir = dev_nano_root_dir
        AbstractConfigurableNanoInstanceHandler.__init__(
            self, reset_user_app_dirs[0], [local_config_mut], [remote_config_mut]
        )

    async def _setup(self):
        pass

    @asynccontextmanager
    async def _prepare(self):
        async with self._inst_run() as (pm, children_spawns):
            yield pm, children_spawns

    async def _handle_prepared(self, prep_args):
        pass

    @asynccontextmanager
    async def _execute(self):
        async with aiohttp.ClientSession() as session:
            async with self._inst_run() as (pm, children_spawns):
                dbg_config = await pm.wait_children_config(session, children_spawns)
                yield pm, session, children_spawns, dbg_config

    @asynccontextmanager
    async def _inst_run(self):
        with DevCoreProcessManager.startup(
            *self.reset_user_app_dirs,
            "clarabot_nano_tests.core.utility.process",
            env={"CBNANO_ROOT": str(self.dev_nano_root_dir)},
        ) as pm:
            children_spawns = await pm.wait_children(create_children_wait_map())
            yield pm, children_spawns


class _SolrDependencyMixin(object):
    def __init__(self):
        self._previous_solr_core = None
        self._previous_solr_port = None
        self._previous_solr_key = None

    async def _wait_for_solr_and_indexer(self, testing_server_manager, pm, same_core=True):
        # check spawned solr (is new and not old config)
        state = await pm.wait_dbg_states(
            [
                ("solr_port", lambda v: isinstance(v, int) and v != self._previous_solr_port),
                ("solr_key", lambda v: isinstance(v, str) and v != self._previous_solr_key),
                (
                    "solr_core",
                    lambda v: isinstance(v, str)
                    and (
                        v == self._previous_solr_core
                        if same_core and self._previous_solr_core
                        else v != self._previous_solr_core
                    ),
                ),
            ],
            timeout=15.0,
        )
        session = await testing_server_manager.session()
        await wait_solr_startup(
            session,
            state["solr_port"],
            create_solr_auth_header(state["solr_key"]),
            state["solr_core"],
        )


class ConfigTestingNanoInstanceHandler(AbstractConfigurableNanoInstanceHandler, _SolrDependencyMixin):
    """
    Create an instance with an adjusted config.
    """

    def __init__(
        self,
        reset_user_app_dirs,
        dev_nano_root_dir,
        testing_server_manager: DevServerManager,
        remote_config_mut: Optional[Union[Callable[[dict], None], partial]] = None,
        local_config_mut: Optional[Union[Callable[[dict], None], partial]] = None,
    ):
        self.reset_user_app_dirs = reset_user_app_dirs
        self.dev_nano_root_dir = dev_nano_root_dir
        self.testing_server_manager = testing_server_manager
        AbstractConfigurableNanoInstanceHandler.__init__(
            self, reset_user_app_dirs[0], [local_config_mut], [remote_config_mut]
        )
        _SolrDependencyMixin.__init__(self)

    async def _setup(self):
        pass

    @asynccontextmanager
    async def _prepare(self):
        async with self._inst_context() as (pm, debug_info):
            yield pm, debug_info

    async def _handle_prepared(self, prep_args):
        pm, debug_info = prep_args
        await self._wait_for_solr_and_indexer(self.testing_server_manager, pm)

    @asynccontextmanager
    async def _execute(self, core_args=(), with_control_api=False, with_document_server=False):
        async with self._inst_context(core_args, with_control_api, with_document_server) as (pm, debug_info):
            await self._wait_for_solr_and_indexer(self.testing_server_manager, pm)
            session = await self.testing_server_manager.session()
            dbg_config = await pm.get_children_config(session, debug_info["port"])
            yield pm, debug_info, session, dbg_config

    def _inst_context(self, core_args=(), with_control_api=False, with_document_server=False):
        return start_nano_wait_ready(
            self.testing_server_manager,
            self.reset_user_app_dirs,
            self.dev_nano_root_dir,
            core_args=core_args,
            with_control_api=with_control_api,
            with_document_server=with_document_server,
        )


class RequestTestingNanoInstanceHandler(AbstractConfigurableNanoInstanceHandler, _SolrDependencyMixin):
    """
    Create an instance that logs in and is available through the server.
    """

    def __init__(
        self,
        reset_user_app_dirs,
        dev_nano_root_dir,
        testing_server_manager: DevServerManager,
        testing_server_user0,
        testing_client_manager: GraphQLClientManager,
        drives: List[dict],
        room_blocks: Dict[str, List[str]],
        clear_slots=True,
        remote_config_mut: Optional[Union[Callable[[dict], None], partial]] = None,
    ):
        self.reset_user_app_dirs = reset_user_app_dirs
        self.dev_nano_root_dir = dev_nano_root_dir
        self.testing_server_manager = testing_server_manager
        self.user_data = testing_server_user0
        self.testing_client_manager = testing_client_manager
        self.drives = drives
        self.room_blocks = room_blocks
        self.clear_slots = clear_slots

        remote_conf_muts = [remote_config_mut]

        if self.drives:

            def _mutate_drives(c):
                assert "drives" in c
                for drive in self.drives:
                    c["drives"].append(drive)

            remote_conf_muts.append(_mutate_drives)

        if self.room_blocks:

            def _mutate_room_blocks(c):
                assert "room_blocks" in c
                c["room_blocks"].update(self.room_blocks)

            remote_conf_muts.append(_mutate_room_blocks)

        AbstractConfigurableNanoInstanceHandler.__init__(self, reset_user_app_dirs[0], [], remote_conf_muts)
        _SolrDependencyMixin.__init__(self)
        self.packed_room_ids = set()
        self.room_ids = set()
        self.uwsc = None
        self.nano_id = None
        self.session_id = None

    async def _setup(self):
        self.packed_room_ids = set()
        for drive in self.drives:
            self.packed_room_ids.update(drive.get("rooms", []))
        self.room_ids = set(
            [
                await self.testing_server_manager.safe_radix32(packed_room_id, pack=False)
                for packed_room_id in self.packed_room_ids
            ]
        )
        await self.testing_server_manager.call_nano_clear_sessions(list(self.packed_room_ids))
        if self.clear_slots:
            await clear_nano_slots(self.testing_server_manager, [self.user_data["account_id"]])

        self.uwsc = await self.testing_client_manager.create_logged_in_wsc(self.user_data)
        assert not self.uwsc.ws.closed

    @asynccontextmanager
    async def _prepare(self):
        async with start_nano_wait_login(
            self.testing_server_manager,
            self.reset_user_app_dirs,
            self.dev_nano_root_dir,
            self.user_data,
        ) as (pm, debug_info, gateway_port, gateway_key, session_id):
            yield pm, debug_info, gateway_port, gateway_key, session_id

    async def _handle_prepared(self, prep_args):
        pm, debug_info, gateway_port, gateway_key, self.session_id = prep_args
        assert self.session_id is not None
        self.nano_id = await self.testing_server_manager.get_nano_id_by_nano_session_id(
            self.user_data["account_id"], self.session_id
        )
        assert self.nano_id is not None
        # wait until indexer startup
        await self._wait_for_solr_and_indexer(self.testing_server_manager, pm)
        # admin requests own nano
        await self.testing_client_manager.nano_request_responses(
            msgpack_dumps(
                await self.testing_client_manager.prepare_request_with_mutation_id(
                    {"type": "debug-echo-sp", "hello": "me"},
                    "dbg0",
                    None,
                    self.nano_id,
                    self.uwsc,
                    self.user_data["_acc_kr"],
                    self.user_data["_acc_kr"],
                )
            ),
            None,
            self.nano_id,
            self.uwsc,
            self.user_data["_acc_kr"],
            self.user_data["_acc_kr"],
        )

    @asynccontextmanager
    async def _execute(self):
        async with start_nano_wait_gateway(
            self.testing_server_manager,
            self.reset_user_app_dirs,
            self.dev_nano_root_dir,
            "active",
        ) as (pm, debug_info, gateway_port, gateway_key):
            assert pm.get_session_id_from_continue_token() == self.session_id
            await self.testing_server_manager.wait_nano_bind(
                self.user_data["account_id"], {self.session_id: self.room_ids}
            )
            # wait until new indexer startup (solr & server consistency)
            await self._wait_for_solr_and_indexer(self.testing_server_manager, pm)
            remote_yaml = self.remote_config_path()
            yield pm, debug_info, gateway_port, gateway_key, self.session_id, remote_yaml, self.nano_id, self.uwsc
