import asyncio
import time
from typing import List

import psutil
import pytest

from clarabot_nano.common.spawn.constants import ShortNames
from clarabot_nano_tests.core.utility.common import create_children_wait_map
from clarabot_nano_tests.testing.core.instance import CoreTestingNanoInstanceHandler
from clarabot_nano_tests.testing.core.manager import DevCoreProcessManager


@pytest.mark.asyncio
async def test_core_subprocess_failures(reset_user_app_dirs, dev_nano_root_dir):
    # test death of subprocess & restart by core
    async with CoreTestingNanoInstanceHandler(reset_user_app_dirs, dev_nano_root_dir).run() as (
        pm,
        session,
        children_spawns,
        dbg_config,
    ):
        # test exiting normally
        server_pid, server_info = next(iter(children_spawns[ShortNames.GENERAL_SERVER].items()))
        server_proc = psutil.Process(server_pid)
        assert server_proc.is_running()
        async with pm.dbg_request(session, server_info["port"], "exit-normal") as r:
            pass
        server_proc.wait(5.0)

        # wait new spawn
        children_spawns = await pm.wait_children(create_children_wait_map(), 15.0)
        await pm.wait_children_config(session, children_spawns)

        # test exiting with error
        server_pid, server_info = next(iter(children_spawns[ShortNames.GENERAL_SERVER].items()))
        server_proc = psutil.Process(server_pid)
        assert server_proc.is_running()
        async with pm.dbg_request(session, server_info["port"], "exit-error") as r:
            pass
        server_proc.wait(5.0)

        # wait new spawn
        children_spawns = await pm.wait_children(create_children_wait_map(), 10.0)
        await pm.wait_children_config(session, children_spawns)

        # test hanging
        server_pid, server_info = next(iter(children_spawns[ShortNames.GENERAL_SERVER].items()))
        server_proc = psutil.Process(server_pid)
        assert server_proc.is_running()
        async with pm.dbg_request(session, server_info["port"], "set-hang") as r:
            pass
        async with pm.dbg_request(session, server_info["port"], "exit-normal") as r:
            pass
        server_proc.wait(35.0)  # 10 KA_TO + ~5 window + 10 KILL_TO

        # wait new spawn
        children_spawns = await pm.wait_children(create_children_wait_map(), 10.0)
        await pm.wait_children_config(session, children_spawns)

        # test hanging at core exit
        # NOTE: this will not trigger kill() post cleanup-wait because cleanups include a kill after a timeout
        server_pid, server_info = next(iter(children_spawns[ShortNames.GENERAL_SERVER].items()))
        server_proc = psutil.Process(server_pid)
        assert server_proc.is_running()
        async with pm.dbg_request(session, server_info["port"], "set-hang") as r:
            pass


@pytest.mark.asyncio
async def test_core_failure(reset_user_app_dirs, dev_nano_root_dir):
    # test death of core & restart by (simulated) os
    async with CoreTestingNanoInstanceHandler(reset_user_app_dirs, dev_nano_root_dir).run() as (
        pm,
        session,
        children_spawns,
        dbg_config,
    ):
        children: List[psutil.Process] = [c["_p"] for tc in children_spawns.values() for c in tc.values()]
        tight_children: List[psutil.Process] = [
            c["_p"]
            for k in [
                ShortNames.GENERAL_SERVER,
                ShortNames.FSTX_SERVER,
                ShortNames.GATEWAY,
                ShortNames.INDEXER,
                ShortNames.CONTROL,
            ]
            for c in children_spawns[k].values()
        ]

        # test core-blocking second instance
        with DevCoreProcessManager.startup(
            *reset_user_app_dirs,
            "clarabot_nano_tests.core.utility.process",
            env={"CBNANO_ROOT": str(dev_nano_root_dir)},
        ) as pm2:
            assert pm2.p.wait(3.0) == 3
        assert pm.p.poll() is None

        # check children still alive after unfortunate demise of core
        pm.p.kill()
        await asyncio.sleep(0.75)
        assert all(c.is_running() for c in children)

    # wait until core-aware children stop voluntarily
    t = time.perf_counter() + 15.0
    while time.perf_counter() < t:
        if all(not c.is_running() for c in tight_children):
            break
        await asyncio.sleep(0.25)
    else:
        assert False
    assert sum(c.is_running() for c in children) == len(children) - len(
        tight_children
    ), f"{[c.name() for c in children]}\n{[c.name() for c in tight_children]}"

    # core restart
    with DevCoreProcessManager.startup(
        *reset_user_app_dirs,
        "clarabot_nano_tests.core.utility.process",
        env={"CBNANO_ROOT": str(dev_nano_root_dir)},
    ) as pm:
        children_spawns = await pm.wait_children(create_children_wait_map())
        assert all(not c.is_running() for c in children)
        children: List[psutil.Process] = [c["_p"] for tc in children_spawns.values() for c in tc.values()]
        assert all(c.is_running() for c in children)

    # check all children are stopped with proper core exit
    assert all(not c.is_running() for c in children)
