# 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/>.

'''
Root actor runtime ignition(s).

'''
from contextlib import (
    asynccontextmanager as acm,
)
from functools import partial
import importlib
import inspect
import logging
import os
import signal
import sys
from typing import (
    Any,
    Callable,
)
import warnings


import trio

from .runtime import _runtime
from .discovery._registry import Registrar
from .devx import (
    debug,
    _frame_stack,
    pformat as _pformat,
)
from .spawn import _spawn
from .runtime import _state
from . import log
from .ipc import (
    _connect_chan,
)
from .discovery._addr import (
    Address,
    UnwrappedAddress,
    default_lo_addrs,
    mk_uuid,
    wrap_address,
)
from .trionics import (
    is_multi_cancelled,
    collapse_eg,
)
from ._exceptions import (
    RuntimeFailure,
)


logger = log.get_logger('tractor')


# TODO: stick this in a `@acm` defined in `devx.debug`?
# -[ ] also maybe consider making this a `wrapt`-deco to
#     save an indent level?
#
@acm
async def maybe_block_bp(
    debug_mode: bool,
    maybe_enable_greenback: bool,
) -> bool:
    # Override the global debugger hook to make it play nice with
    # ``trio``, see much discussion in:
    # https://github.com/python-trio/trio/issues/1155#issuecomment-742964018
    builtin_bp_handler: Callable = sys.breakpointhook
    orig_bp_path: str|None = os.environ.get(
        'PYTHONBREAKPOINT',
        None,
    )
    bp_blocked: bool
    if (
        debug_mode
        and
        maybe_enable_greenback
        and (
            maybe_mod := await debug.maybe_init_greenback(
                raise_not_found=False,
            )
        )
    ):
        logger.info(
            f'Found `greenback` installed @ {maybe_mod}\n'
            f'Enabling `tractor.pause_from_sync()` support!\n'
        )
        os.environ['PYTHONBREAKPOINT'] = (
            'tractor.devx.debug._sync_pause_from_builtin'
        )
        _state._runtime_vars['use_greenback'] = True
        bp_blocked = False

    else:
        # TODO: disable `breakpoint()` by default (without
        # `greenback`) since it will break any multi-actor
        # usage by a clobbered TTY's stdstreams!
        def block_bps(*args, **kwargs):
            raise RuntimeError(
                'Trying to use `breakpoint()` eh?\n\n'
                'Welp, `tractor` blocks `breakpoint()` built-in calls by default!\n'
                'If you need to use it please install `greenback` and set '
                '`debug_mode=True` when opening the runtime '
                '(either via `.open_nursery()` or `open_root_actor()`)\n'
            )

        sys.breakpointhook = block_bps
        # lol ok,
        # https://docs.python.org/3/library/sys.html#sys.breakpointhook
        os.environ['PYTHONBREAKPOINT'] = "0"
        bp_blocked = True

    try:
        yield bp_blocked
    finally:
        # restore any prior built-in `breakpoint()` hook state
        if builtin_bp_handler is not None:
            sys.breakpointhook = builtin_bp_handler

        if orig_bp_path is not None:
            os.environ['PYTHONBREAKPOINT'] = orig_bp_path

        else:
            # clear env back to having no entry
            os.environ.pop('PYTHONBREAKPOINT', None)



@acm
async def open_root_actor(
    *,
    tpt_bind_addrs: list[
        Address  # `Address.get_random()` case
        |UnwrappedAddress  # registrar case `= uw_reg_addrs`
    ]|None = None,

    # defaults are above
    registry_addrs: list[
        Address
        |UnwrappedAddress
    ]|None = None,

    enable_transports: list[
        # TODO, this should eventually be the pairs as
        # defined by (codec, proto) as on `MsgTransport.
        _state.TransportProtocolKey,
    ]|None = None,

    name: str|None = 'root',

    # either the `multiprocessing` start method:
    # https://docs.python.org/3/library/multiprocessing.html#contexts-and-start-methods
    # OR `trio` (the new default).
    start_method: _spawn.SpawnMethodKey|None = None,

    # enables the multi-process debugger support
    debug_mode: bool = False,
    maybe_enable_greenback: bool = False,  # `.pause_from_sync()/breakpoint()` support
    # ^XXX NOTE^ the perf implications of use,
    # https://greenback.readthedocs.io/en/latest/principle.html#performance
    enable_stack_on_sig: bool = False,

    # internal logging
    loglevel: str|None = None,

    enable_modules: list|None = None,
    rpc_module_paths: list|None = None,

    # NOTE: allow caller to ensure that only one registry exists
    # and that this call creates it.
    ensure_registry: bool = False,

    hide_tb: bool = True,

    # XXX, proxied directly to `.devx.debug._maybe_enter_pm()`
    # for REPL-entry logic.
    debug_filter: Callable[
        [BaseException|BaseExceptionGroup],
        bool,
    ] = lambda err: not is_multi_cancelled(err),

    # TODO, a way for actors to augment passing derived
    # read-only state to sublayers?
    # extra_rt_vars: dict|None = None,

) -> _runtime.Actor:
    '''
    Initialize the `tractor` runtime by starting a "root actor" in
    a parent-most Python process.

    All (disjoint) actor-process-trees-as-programs are created via
    this entrypoint.

    '''
    # XXX NEVER allow nested actor-trees!
    if already_actor := _state.current_actor(
        err_on_no_runtime=False,
    ):
        rtvs: dict[str, Any] = _state._runtime_vars
        root_mailbox: list[str, int] = rtvs['_root_mailbox']
        registry_addrs: list[list[str, int]] = rtvs['_registry_addrs']
        raise RuntimeFailure(
            f'A current actor already exists !?\n'
            f'({already_actor}\n'
            f'\n'
            f'You can NOT open a second root actor from within '
            f'an existing tree and the current root of this '
            f'already exists !!\n'
            f'\n'
            f'_root_mailbox: {root_mailbox!r}\n'
            f'_registry_addrs: {registry_addrs!r}\n'
        )

    # debug.mk_pdb().set_trace()
    async with maybe_block_bp(
        debug_mode=debug_mode,
        maybe_enable_greenback=maybe_enable_greenback,
    ):
        if enable_transports is None:
            enable_transports: list[str] = _state.current_ipc_protos()
        else:
            _state._runtime_vars['_enable_tpts'] = enable_transports

        # TODO! support multi-tpts per actor!
        # Bo
        if not len(enable_transports) == 1:
            raise RuntimeError(
                f'No multi-tpt support yet!\n'
                f'enable_transports={enable_transports!r}\n'
            )

        _frame_stack.hide_runtime_frames()
        __tracebackhide__: bool = hide_tb

        # attempt to retreive ``trio``'s sigint handler and stash it
        # on our debugger lock state.
        debug.DebugStatus._trio_handler = signal.getsignal(signal.SIGINT)

        # mark top most level process as root actor
        _state._runtime_vars['_is_root'] = True

        # caps based rpc list
        enable_modules = (
            enable_modules
            or
            []
        )

        if rpc_module_paths:
            warnings.warn(
                "`rpc_module_paths` is now deprecated, use "
                " `enable_modules` instead.",
                DeprecationWarning,
                stacklevel=2,
            )
            enable_modules.extend(rpc_module_paths)

        # `TRACTOR_LOGLEVEL` env-var wins over any caller-passed
        # `loglevel` so devs/test-runs can crank (or silence)
        # console verbosity without touching application code.
        env_ll_report: str = ''
        if env_ll := os.environ.get('TRACTOR_LOGLEVEL'):
            loglevel = env_ll
            env_ll_report: str = (
                f'Detected env-var setting,\n'
                f'TRACTOR_LOGLEVEL={env_ll!r}\n'
                f'\n'
                f'Setting console loglevel per,\n'
                f'loglevel={loglevel!r}\n'
            )
            if (
                loglevel
                and
                loglevel.upper() != env_ll.upper()
            ):
                env_ll_report += (
                    f'\n'
                    f'NOTE env-var OVERRIDES caller-passed,\n'
                    f'loglevel={loglevel!r}\n'
                )

        loglevel: str = (
            loglevel
            or
            log._default_loglevel
        )
        loglevel: str = loglevel.upper()

        assert loglevel
        _log = log.get_console_log(
            level=loglevel,
            name='tractor',
            logger=logger,
        )
        assert _log
        if env_ll_report:
            _log.info(env_ll_report)

        # `TRACTOR_SPAWN_METHOD` env-var wins over any caller-passed
        # `start_method` so devs/test-runs can swap the actor spawn
        # backend without touching application code (e.g. driving
        # the `examples/debugging/<script>.py` suite under each
        # backend from `tests/devx/conftest.py::spawn`).
        if env_sm := os.environ.get('TRACTOR_SPAWN_METHOD'):
            start_method: str = env_sm
            env_sm_report: str = (
                f'Detected env-var setting,\n'
                f'TRACTOR_SPAWN_METHOD={env_sm!r}\n'
                f'\n'
                f'Setting spawn backend as,\n'
                f'start_method={env_sm!r}\n'
            )
            if (
                start_method
                and
                start_method != env_sm
            ):
                _log.warning(
                    env_sm_report
                    +
                    f'NOTE env-var OVERRIDES caller-passed,\n'
                    f'`start_method={start_method!r}`\n'
                )
            else:
                _log.info(env_sm_report)

        if start_method is not None:
            _spawn.try_set_start_method(start_method)

        # XXX expect pre-unwrapped registrar addrs.
        uw_reg_addrs = registry_addrs
        if not uw_reg_addrs:
            uw_reg_addrs: list[UnwrappedAddress] = default_lo_addrs(
                enable_transports
            )

        # must exist by now since all below code is dependent
        assert uw_reg_addrs
        registry_addrs: list[Address] = [
            wrap_address(uw_addr)
            for uw_addr in uw_reg_addrs
        ]

        # fail-fast on `enable_transports` / `registry_addrs` proto
        # mismatch — historically this caused a silent indefinite
        # hang during the registrar handshake (registry was reachable
        # only via a transport not in `enable_transports`, so the
        # actor could never connect to register/discover). See
        # `tests/ipc/test_multi_tpt.py::test_root_passes_tpt_to_sub`
        # for the foot-gun case + its layer-1 skip-guard.
        bad_addrs: list[tuple[str, Address]] = [
            (addr.proto_key, addr)
            for addr in registry_addrs
            if addr.proto_key not in enable_transports
        ]
        if bad_addrs:
            mismatch_lines: str = '\n'.join(
                f'  - proto_key={pk!r}  addr={a!r}'
                for pk, a in bad_addrs
            )
            raise ValueError(
                f'`registry_addrs` contains addr(s) whose proto is '
                f'not in `enable_transports`!\n'
                f'enable_transports: {enable_transports!r}\n'
                f'mismatched_addrs:\n'
                f'{mismatch_lines}\n'
                f'\n'
                f'Either add the missing proto to '
                f'`enable_transports`, or remove the addr from '
                f'`registry_addrs`.'
            )

        if (
            debug_mode
            and
            _spawn._spawn_method == 'trio'
        ):
            _state._runtime_vars['_debug_mode'] = True

            # expose internal debug module to every actor allowing for
            # use of ``await tractor.pause()``
            enable_modules.append('tractor.devx.debug._tty_lock')

            # if debug mode get's enabled *at least* use that level of
            # logging for some informative console prompts.
            if (
                logging.getLevelName(
                    # lul, need the upper case for the -> int map?
                    # sweet "dynamic function behaviour" stdlib...
                    loglevel,
                ) > logging.getLevelName('PDB')
            ):
                loglevel = 'PDB'


        elif debug_mode:
            raise RuntimeError(
                "Debug mode is only supported for the `trio` backend!"
            )

        # TODO: factor this into `.devx._stackscope`!!
        #
        # NOTE, intentionally NOT gated on `debug_mode` so SIGUSR1
        # task-tree dumps work in plain (non-pdb) runs too — esp.
        # in infected-`asyncio` root processes where the default
        # SIGUSR1 action would otherwise terminate the proc.
        if enable_stack_on_sig:
            from .devx._stackscope import enable_stack_on_sig
            enable_stack_on_sig()

        # closed into below ping task-func
        ponged_addrs: list[Address] = []

        async def ping_tpt_socket(
            addr: Address,
            timeout: float = 1,
        ) -> None:
            '''
            Attempt temporary connection to see if a registry is
            listening at the requested address by a tranport layer
            ping.

            If a connection can't be made quickly we assume none no
            server is listening at that addr.

            '''
            try:
                # TODO: this connect-and-bail forces us to have to
                # carefully rewrap TCP 104-connection-reset errors as
                # EOF so as to avoid propagating cancel-causing errors
                # to the channel-msg loop machinery. Likely it would
                # be better to eventually have a "discovery" protocol
                # with basic handshake instead?
                with trio.move_on_after(timeout):
                    async with _connect_chan(addr.unwrap()):
                        ponged_addrs.append(addr)

            except OSError:
                # ?TODO, make this a "discovery" log level?
                logger.info(
                    f'No root-actor registry found @ {addr!r}\n'
                )

        # !TODO, this is basically just another (abstract)
        # happy-eyeballs, so we should try for formalize it somewhere
        # in a `.[_]discovery` ya?
        #
        async with trio.open_nursery() as tn:
            for uw_addr in uw_reg_addrs:
                addr: Address = wrap_address(uw_addr)
                tn.start_soon(
                    ping_tpt_socket,
                    addr,
                )

        if tpt_bind_addrs is None:
            tpt_bind_addrs: list[Address] = []
        else:
            input_bind_addrs = list(tpt_bind_addrs)
            tpt_bind_addrs: list[Address] = []
            for addr in input_bind_addrs:
                addr: Address = wrap_address(addr)
                tpt_bind_addrs.append(addr)

        # ------ NON-REGISTRAR ------
        # create a new root-actor instance.
        if ponged_addrs:
            if ensure_registry:
                raise RuntimeError(
                     f'Failed to open `{name}`@{ponged_addrs}: '
                    'registry socket(s) already bound'
                )

            # we were able to connect to a registrar
            logger.info(
                f'Registry(s) seem(s) to exist @ {ponged_addrs}'
            )

            actor = _runtime.Actor(
                name=name or 'anonymous',
                uuid=mk_uuid(),
                registry_addrs=ponged_addrs,
                loglevel=loglevel,
                enable_modules=enable_modules,
            )
            # **DO NOT** use the registry_addrs as the
            # ipc-transport-server's bind-addrs as this is
            # a new NON-registrar, ROOT-actor.
            #
            # XXX INSTEAD, bind random addrs using the same tpt
            # proto if not already provided.
            if not tpt_bind_addrs:
                for addr in ponged_addrs:
                    tpt_bind_addrs.append(
                        # XXX, these are `Address` NOT `UnwrappedAddress`.
                        #
                        # NOTE, in the case of posix/berkley socket
                        # protos we allocate port=0 such that the system
                        # allocates a random value at bind time; this
                        # happens in the `.ipc.*` stack's backend.
                        addr.get_random(
                            bindspace=addr.bindspace,
                        )
                    )

            header: str = '-> Contacting existing registry @ '

        # ------ REGISTRAR ------
        # create a new "registry providing" root-actor instance.
        #
        # Start this local actor as the "registrar", aka a regular
        # actor who manages the local registry of "mailboxes" of
        # other process-tree-local sub-actors.
        else:
            # NOTE that if the current actor IS THE REGISTAR, the
            # following init steps are taken:
            # - the tranport layer server is bound to each addr
            #   pair defined in provided registry_addrs, or the default.
            tpt_bind_addrs = list(set(
                tpt_bind_addrs
                +
                [wrap_address(a) for a in uw_reg_addrs]
            ))

            # - it is normally desirable for any registrar to stay up
            #   indefinitely until either all registered (child/sub)
            #   actors are terminated (via SC supervision) or,
            #   a re-election process has taken place.
            # NOTE: all of ^ which is not implemented yet - see:
            # https://github.com/goodboy/tractor/issues/216
            # https://github.com/goodboy/tractor/pull/348
            # https://github.com/goodboy/tractor/issues/296

            actor = Registrar(
                name=name or 'registrar',
                uuid=mk_uuid(),
                registry_addrs=registry_addrs,
                loglevel=loglevel,
                enable_modules=enable_modules,
            )
            # XXX, in case the root actor runtime was actually run from
            # `tractor.to_asyncio.run_as_asyncio_guest()` and NOT
            # `.trio.run()`.
            actor._infected_aio = _state._runtime_vars['_is_infected_aio']
            header: str = '-> Opening new registry @ '

        # Start up main task set via core actor-runtime nurseries.
        try:
            # assign process-local actor
            _state._current_actor = actor

            # start local channel-server and fake the portal API
            # NOTE: this won't block since we provide the nursery
            report: str = f'Starting actor-runtime for {actor.aid.reprol()!r}\n'
            if reg_addrs := actor.registry_addrs:
                report += (
                    header
                    +
                    '\n'.join(
                        f'{addr}' for addr in reg_addrs
                    )
                )
            logger.info(f'{report}\n')

            # start runtime in a bg sub-task, yield to caller.
            async with (
                collapse_eg(),
                trio.open_nursery() as root_tn,

                # ?TODO? finally-footgun below?
                # -> see note on why shielding.
                # maybe_raise_from_masking_exc(),
            ):
                actor._root_tn = root_tn
                # `_runtime.async_main()` creates an internal nursery
                # and blocks here until any underlying actor(-process)
                # tree has terminated thereby conducting so called
                # "end-to-end" structured concurrency throughout an
                # entire hierarchical python sub-process set; all
                # "actor runtime" primitives are SC-compat and thus all
                # transitively spawned actors/processes must be as
                # well.
                accept_addrs: list[UnwrappedAddress]
                reg_addrs: list[UnwrappedAddress]
                (
                    accept_addrs,
                    reg_addrs,
                ) = await root_tn.start(
                    partial(
                        _runtime.async_main,
                        actor,
                        accept_addrs=tpt_bind_addrs,
                        parent_addr=None
                    )
                )
                # NOTE, only set a local-host addr (i.e. like
                # `lo`-loopback for TCP) for the process-tree-global
                # "root"-process (its tree-wide "mailbox") since all
                # sub-actors should be able to speak to their root
                # actor over that channel.
                #
                # ?TODO, per-OS non-network-proto alt options?
                # -[ ] on linux we should be able to always use UDS?
                #
                raddrs: list[UnwrappedAddress] = _state._runtime_vars['_root_addrs']
                raddrs.extend(
                    accept_addrs,
                )
                # TODO, remove once we have also removed all usage;
                # eventually all (root-)registry apis should expect > 1 addr.
                _state._runtime_vars['_root_mailbox'] = raddrs[0]
                # if 'chart' in actor.aid.name:
                #     from tractor.devx import mk_pdb
                #     mk_pdb().set_trace()

                try:
                    yield actor
                except (
                    Exception,
                    BaseExceptionGroup,
                ) as err:

                    # TODO, in beginning to handle the subsubactor with
                    # crashed grandparent cases..
                    #
                    # was_locked: bool = await debug.maybe_wait_for_debugger(
                    #     child_in_debug=True,
                    # )
                    # XXX NOTE XXX see equiv note inside
                    # `._runtime.Actor._stream_handler()` where in the
                    # non-root or root-that-opened-this-mahually case we
                    # wait for the local actor-nursery to exit before
                    # exiting the transport channel handler.
                    entered: bool = await debug._maybe_enter_pm(
                        err,
                        api_frame=inspect.currentframe(),
                        debug_filter=debug_filter,

                        # XXX NOTE, required to debug root-actor
                        # crashes under cancellation conditions; so
                        # most of them!
                        shield=root_tn.cancel_scope.cancel_called,
                    )

                    if (
                        not entered
                        and
                        not is_multi_cancelled(
                            err,
                        )
                    ):
                        logger.exception(
                            'Root actor crashed\n'
                            f'>x)\n'
                            f' |_{actor}\n'
                        )

                    # ALWAYS re-raise any error bubbled up from the
                    # runtime!
                    raise

                finally:
                    # NOTE/TODO?, not sure if we'll ever need this but it's
                    # possibly better for even more determinism?
                    # logger.cancel(
                    #     f'Waiting on {len(nurseries)} nurseries in root..')
                    # nurseries = actor._actoruid2nursery.values()
                    # async with trio.open_nursery() as tempn:
                    #     for an in nurseries:
                    #         tempn.start_soon(an.exited.wait)

                    op_nested_actor_repr: str = _pformat.nest_from_op(
                        input_op='>) ',
                        text=actor.pformat(),
                        nest_prefix='|_',
                    )
                    logger.info(
                        f'Closing down root actor\n'
                        f'{op_nested_actor_repr}'
                    )
                    # XXX, THIS IS A *finally-footgun*!
                    # (also mentioned in with-block above)
                    # -> though already shields iternally it can
                    # taskc here and mask underlying errors raised in
                    # the try-block above?
                    with trio.CancelScope(shield=True):
                        await actor.cancel(None)  # self cancel
        finally:
            # revert all process-global runtime state
            if (
                debug_mode
                and
                _spawn._spawn_method == 'trio'
            ):
                _state._runtime_vars['_debug_mode'] = False

            # !XXX, clear ALL prior contact info state, this is MEGA
            # important if you are opening the runtime multiple times
            # from the same parent process (like in our test
            # harness)!
            _state._runtime_vars['_root_addrs'].clear()
            _state._runtime_vars['_root_mailbox'] = None

            _state._current_actor = None
            _state._last_actor_terminated = actor

            sclang_repr: str = _pformat.nest_from_op(
                input_op=')>',
                text=actor.pformat(),
                nest_prefix='|_',
                nest_indent=1,
            )

            logger.info(
                f'Root actor terminated\n'
                f'{sclang_repr}'
            )


def run_daemon(
    enable_modules: list[str],

    # runtime kwargs
    name: str | None = 'root',
    registry_addrs: list[UnwrappedAddress]|None = None,

    start_method: str | None = None,
    debug_mode: bool = False,

    # TODO, support `infected_aio=True` mode by,
    # - calling the appropriate entrypoint-func from `.to_asyncio`
    # - maybe init-ing `greenback` as done above in
    #   `open_root_actor()`.

    **kwargs

) -> None:
    '''
    Spawn a root (daemon) actor which will respond to RPC; the main
    task simply starts the runtime and then blocks via embedded
    `trio.sleep_forever()`.

    This is a very minimal convenience wrapper around starting
    a "run-until-cancelled" root actor which can be started with a set
    of enabled modules for RPC request handling.

    '''
    kwargs['enable_modules'] = list(enable_modules)

    for path in enable_modules:
        importlib.import_module(path)

    async def _main():
        async with open_root_actor(
            registry_addrs=registry_addrs,
            name=name,
            start_method=start_method,
            debug_mode=debug_mode,
            **kwargs,
        ):
            return await trio.sleep_forever()

    return trio.run(_main)