# tractor: structured concurrent "actors".
# Copyright 2018-eternity Tyler Goodlet.

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

"""
Top level routines & machinery for actor-as-process/subint spawning
over multiple backends.

"""
from __future__ import annotations
import multiprocessing as mp
import platform
from typing import (
    Any,
    Awaitable,
    Literal,
    Callable,
    TypeVar,
    TYPE_CHECKING,
)

import trio
from trio import TaskStatus

from .._exceptions import ActorFailure
from ..devx import debug
from tractor.runtime._state import (
    _runtime_vars,
)
from tractor.log import get_logger
from tractor.discovery._addr import (
    UnwrappedAddress,
)
from ._reap import unlink_uds_bind_addrs
from tractor.runtime._portal import Portal
from tractor.runtime._runtime import Actor
from tractor.msg import types as msgtypes


if TYPE_CHECKING:
    from tractor.net._bindspace import Bindspace
    from tractor.ipc import (
        _server,
        Channel,
    )
    from tractor.runtime._supervise import ActorNursery
    ProcessType = TypeVar('ProcessType', mp.Process, trio.Process)


log = get_logger('tractor')

# placeholder for an mp start context if so using that backend
_ctx: mp.context.BaseContext | None = None
SpawnMethodKey = Literal[
    'trio',  # supported on all platforms
    'mp_spawn',
    'mp_forkserver',  # posix only
]
_spawn_method: SpawnMethodKey = 'trio'


if platform.system() == 'Windows':

    _ctx = mp.get_context("spawn")

    async def proc_waiter(proc: mp.Process) -> None:
        await trio.lowlevel.WaitForSingleObject(proc.sentinel)
else:
    # *NIX systems use ``trio`` primitives as our default as well

    async def proc_waiter(proc: mp.Process) -> None:
        await trio.lowlevel.wait_readable(proc.sentinel)


async def wait_for_peer_or_proc_death(
    ipc_server: _server.Server,
    uid: tuple[str, str],
    proc_wait: Callable[[], Awaitable[int]],
    proc_repr: object = '',
) -> tuple[trio.Event, Channel]:
    '''
    Race a child handshake against process death during bootstrap.

    A child can exit before connecting to its parent. Waiting only on
    `IPCServer.wait_for_peer()` would then park its spawning task and
    leave the dead process unreaped. Run both waits in one nursery and
    let either completed result cancel its sibling.

    Return the normal peer event and channel when the handshake wins.
    Raise `ActorFailure` with the process status when death wins.

    Adapted from goodboy's Claude Code-assisted implementation in
    commit `3b0724eba85b4014170ed95773e1e41a60d5c513`.

    '''
    handshake: tuple[trio.Event, Channel]|None = None
    handshake_error: BaseException|None = None
    returncode: int|None = None
    death_error: BaseException|None = None

    async def wait_for_handshake() -> None:
        '''
        Publish a connected peer before cancelling the death waiter.

        '''
        nonlocal handshake
        nonlocal handshake_error
        try:
            handshake = await ipc_server.wait_for_peer(uid)
        except trio.Cancelled:
            if (
                returncode is not None
                or
                death_error is not None
            ):
                log.debug(
                    'Peer-handshake waiter cancelled after '
                    f'process wait completed for {uid!r}'
                )
            raise
        except BaseException as exc:
            handshake_error = exc
        nursery.cancel_scope.cancel()

    async def wait_for_death() -> None:
        '''
        Publish child exit before cancelling the handshake waiter.

        '''
        nonlocal returncode
        nonlocal death_error
        try:
            returncode = await proc_wait()
        except trio.Cancelled:
            if (
                handshake is not None
                or
                handshake_error is not None
            ):
                log.debug(
                    'Process-death waiter cancelled after '
                    f'peer wait completed for {uid!r}'
                )
            raise
        except BaseException as exc:
            death_error = exc
        nursery.cancel_scope.cancel()

    async with trio.open_nursery() as nursery:
        nursery.start_soon(wait_for_handshake)
        nursery.start_soon(wait_for_death)

    if handshake_error is not None:
        raise handshake_error
    if death_error is not None:
        raise death_error

    if returncode is not None:
        raise ActorFailure(
            f'Sub-actor {uid!r} died during boot '
            f'(rc={returncode!r}) before completing '
            f'parent-handshake.\n'
            f'  proc: {proc_repr}'
        )

    assert handshake is not None
    return handshake


def try_set_start_method(
    key: SpawnMethodKey

) -> mp.context.BaseContext | None:
    '''
    Attempt to set the method for process starting, aka the "actor
    spawning backend".

    If the desired method is not supported this function will error.
    On Windows only the ``multiprocessing`` "spawn" method is offered
    besides the default ``trio`` which uses async wrapping around
    ``subprocess.Popen``.

    '''
    import multiprocessing as mp
    global _ctx
    global _spawn_method

    mp_methods = mp.get_all_start_methods()
    if 'fork' in mp_methods:
        # forking is incompatible with ``trio``s global task tree
        mp_methods.remove('fork')

    match key:
        case 'mp_forkserver':
            from . import _forkserver_override
            _forkserver_override.override_stdlib()
            _ctx = mp.get_context('forkserver')

        case 'mp_spawn':
            _ctx = mp.get_context('spawn')

        case 'trio':
            _ctx = None

        case _:
            raise ValueError(
                f'Spawn method `{key}` is invalid!\n'
                f'Please choose one of {SpawnMethodKey}'
            )

    _spawn_method = key
    return _ctx


async def hard_kill(
    proc: trio.Process,

    terminate_after: int = 1.6,
    # NOTE: for mucking with `.pause()`-ing inside the runtime
    # whilst also hacking on it XD
    # terminate_after: int = 99999,

    *,
    # Subactor's bind addresses + subactor record, used
    # for post-SIGKILL UDS sockpath cleanup. Optional for
    # legacy callers; new call sites should pass at least
    # `subactor` (which lets us reconstruct the sock path
    # from `aid.name + proc.pid` when `bind_addrs` is
    # empty/self-assigned). See `._reap.unlink_uds_bind_addrs()`.
    bind_addrs: list[UnwrappedAddress] | None = None,
    subactor: Actor | None = None,

) -> None:
    '''
    Un-gracefully terminate an OS level `trio.Process` after timeout.

    Used in 2 main cases:

    - "unknown remote runtime state": a hanging/stalled actor that
      isn't responding after sending a (graceful) runtime cancel
      request via an IPC msg.
    - "cancelled during spawn": a process who's actor runtime was
      cancelled before full startup completed (such that
      cancel-request-handling machinery was never fully
      initialized) and thus a "cancel request msg" is never going
      to be handled.

    '''
    log.cancel(
        'Terminating sub-proc\n'
        f'>x)\n'
        f' |_{proc}\n'
    )
    # NOTE: this timeout used to do nothing since we were shielding
    # the ``.wait()`` inside ``new_proc()`` which will pretty much
    # never release until the process exits, now it acts as
    # a hard-kill time ultimatum.
    with trio.move_on_after(terminate_after) as cs:

        # NOTE: code below was copied verbatim from the now deprecated
        # (in 0.20.0) ``trio._subrocess.Process.aclose()``, orig doc
        # string:
        #
        # Close any pipes we have to the process (both input and output)
        # and wait for it to exit. If cancelled, kills the process and
        # waits for it to finish exiting before propagating the
        # cancellation.
        #
        # This code was originally triggred by ``proc.__aexit__()``
        # but now must be called manually.
        with trio.CancelScope(shield=True):
            if proc.stdin is not None:
                await proc.stdin.aclose()
            if proc.stdout is not None:
                await proc.stdout.aclose()
            if proc.stderr is not None:
                await proc.stderr.aclose()
        try:
            await proc.wait()
        finally:
            if proc.returncode is None:
                proc.kill()
                with trio.CancelScope(shield=True):
                    await proc.wait()

    # XXX NOTE XXX: zombie squad dispatch:
    # (should ideally never, but) If we do get here it means
    # graceful termination of a process failed and we need to
    # resort to OS level signalling to interrupt and cancel the
    # (presumably stalled or hung) actor. Since we never allow
    # zombies (as a feature) we ask the OS to do send in the
    # removal swad as the last resort.
    if cs.cancelled_caught:

        # TODO? attempt at intermediary-rent-sub
        # with child in debug lock?
        # |_https://github.com/goodboy/tractor/issues/320
        #
        # if not is_root_process():
        #     log.warning(
        #         'Attempting to acquire debug-REPL-lock before zombie reap!'
        #     )
        #     with trio.CancelScope(shield=True):
        #         async with debug.acquire_debug_lock(
        #             subactor_uid=current_actor().aid.uid,
        #         ) as _ctx:
        #             log.warning(
        #                 'Acquired debug lock, child ready to be killed ??\n'
        #             )

        # TODO: toss in the skynet-logo face as ascii art?
        log.critical(
            # 'Well, the #ZOMBIE_LORD_IS_HERE# to collect\n'
            '#T-800 deployed to collect zombie B0\n'
            f'>x)\n'
            f' |_{proc}\n'
        )
        proc.kill()

    # Post-mortem UDS sockpath cleanup. SIGKILL bypassed
    # the subactor's normal `os.unlink(addr.sockpath)` in
    # `_serve_ipc_eps`'s `finally:`; the parent has the
    # bind addrs (or can reconstruct from name + pid) so
    # we do it here. Runs UNCONDITIONALLY (graceful-exit
    # case is a no-op via `FileNotFoundError` skip in the
    # helper) so the cleanup also covers the "cancelled
    # during spawn" path where the subactor never reached
    # its IPC server finally block.
    unlink_uds_bind_addrs(
        proc,
        bind_addrs=bind_addrs,
        subactor=subactor,
    )


async def soft_kill(
    proc: ProcessType,
    wait_func: Callable[
        [ProcessType],
        Awaitable,
    ],
    portal: Portal,

) -> None:
    '''
    Wait for proc termination but **don't yet** teardown
    std-streams since it will clobber any ongoing pdb REPL
    session.

    This is our "soft"/graceful, and thus itself also cancellable,
    join/reap on an actor-runtime-in-process shutdown; it is
    **not** the same as a "hard kill" via an OS signal (for that
    see `.hard_kill()`).

    '''
    chan: Channel = portal.channel
    peer_aid: msgtypes.Aid = chan.aid
    try:
        log.cancel(
            f'Soft killing sub-actor via portal request\n'
            f'\n'
            f'c)=> {peer_aid.reprol()}@[{chan.maddr}]\n'
            f'   |_{proc}\n'
        )
        # wait on sub-proc to signal termination
        await wait_func(proc)

    except trio.Cancelled:
        with trio.CancelScope(shield=True):
            await debug.maybe_wait_for_debugger(
                child_in_debug=_runtime_vars.get(
                    '_debug_mode', False
                ),
                header_msg=(
                    'Delaying `soft_kill()` subproc reaper while debugger locked..\n'
                ),
                # TODO: need a diff value then default?
                # poll_steps=9999999,
            )

        # if cancelled during a soft wait, cancel the child
        # actor before entering the hard reap sequence
        # below. This means we try to do a graceful teardown
        # via sending a cancel message before getting out
        # zombie killing tools.
        async with trio.open_nursery() as n:
            n.cancel_scope.shield = True

            async def cancel_on_proc_deth():
                '''
                "Cancel-the-cancel" request: if we detect that the
                underlying sub-process exited prior to
                a `Portal.cancel_actor()` call completing .

                '''
                await wait_func(proc)
                n.cancel_scope.cancel()

            # start a task to wait on the termination of the
            # process by itself waiting on a (caller provided) wait
            # function which should unblock when the target process
            # has terminated.
            n.start_soon(cancel_on_proc_deth)

            # send the actor-runtime a cancel request.
            await portal.cancel_actor()

            if proc.poll() is None:  # type: ignore
                log.warning(
                    'Subactor still alive after cancel request?\n\n'
                    f'uid: {peer_aid}\n'
                    f'|_{proc}\n'
                )
                n.cancel_scope.cancel()
        raise


async def new_proc(
    name: str,
    actor_nursery: ActorNursery,
    subactor: Actor,
    errors: dict[tuple[str, str], Exception],

    # passed through to actor main
    bind_addrs: list[UnwrappedAddress],
    parent_addr: UnwrappedAddress,
    _runtime_vars: dict[str, Any],  # serialized and sent to _child

    *,

    bindspace: Bindspace|None = None,
    infect_asyncio: bool = False,
    task_status: TaskStatus[Portal] = trio.TASK_STATUS_IGNORED,
    proc_kwargs: dict[str, any] = {}

) -> None:

    # lookup backend spawning target
    target: Callable = _methods[_spawn_method]

    # mark the new actor with the global spawn method
    subactor._spawn_method = _spawn_method

    try:
        await target(
            name,
            actor_nursery,
            subactor,
            errors,
            bind_addrs,
            parent_addr,
            _runtime_vars,  # run time vars
            bindspace=bindspace,
            infect_asyncio=infect_asyncio,
            task_status=task_status,
            proc_kwargs=proc_kwargs
        )
    finally:
        actor_nursery._mark_child_reaped(subactor.aid)


# NOTE: bottom-of-module to avoid a circular import since the
# backend submodules pull `soft_kill`/`hard_kill`/`proc_waiter`
# from this module.
from ._trio import trio_proc
from ._mp import mp_proc


# proc spawning backend target map
_methods: dict[SpawnMethodKey, Callable] = {
    'trio': trio_proc,
    'mp_spawn': mp_proc,
    'mp_forkserver': mp_proc,
}