'''
Let's make sure them docs work yah?

'''
from contextlib import contextmanager
import itertools
import os
import signal
import sys
import subprocess
import platform
import shutil
from typing import Callable
from unittest.mock import Mock

import pytest
import tractor
from tractor._testing import (
    examples_dir,
)

_non_linux: bool = platform.system() != 'Linux'
_friggin_macos: bool = platform.system() == 'Darwin'


def _kill_proc_tree(proc: subprocess.Popen) -> None:
    '''
    Terminate an example process and its POSIX descendants.

    '''
    try:
        if platform.system() == 'Windows':
            proc.kill()
        else:
            os.killpg(proc.pid, signal.SIGKILL)
    except ProcessLookupError:
        pass


def _reap_killed_proc(
    proc: subprocess.Popen,
) -> tuple[bytes, bytes]:
    '''
    Reap a killed process without waiting on Windows descendants.

    '''
    if platform.system() != 'Windows':
        return proc.communicate()

    proc.wait(timeout=5)
    if proc.stdin:
        proc.stdin.close()
    if proc.stdout:
        proc.stdout.close()
    if proc.stderr:
        proc.stderr.close()
    return b'', b''


def _wait_for_proc(
    proc: subprocess.Popen,
    timeout: float,
    test_log: tractor.log.StackLevelAdapter,
) -> None:
    '''
    Wait for an example process and surface its captured output.

    '''
    try:
        out, err = proc.communicate(timeout=timeout)

    except subprocess.TimeoutExpired as timeout_exc:
        test_log.exception(
            f'Example failed to finish within {timeout}s ??\n'
        )
        _kill_proc_tree(proc)
        out, err = _reap_killed_proc(proc)
        if platform.system() == 'Windows':
            out = timeout_exc.output or b''
            err = timeout_exc.stderr or b''

    errmsg: str = err.decode(errors='replace')

    # XXX, ALWAYS surface the subproc's full stderr
    # whenever it exits non-zero!
    #
    # The prior impl only raised when the LAST stderr
    # line contained 'Error', swallowing any crash whose
    # traceback ends in a non-`XxxError:` line; in
    # particular EVERY `tractor` root-actor crash ends
    # with the strict-EG collapse note,
    # '( ^^^ this exc was collapsed from a group ^^^ )',
    # so ALL such failures were reduced to a bare
    # `assert 1 == 0` in CI logs.. see GH #473.
    rc: int|None = proc.returncode
    if rc:
        outmsg: str = out.decode(errors='replace')
        raise Exception(
            f'Example script exited with rc={rc} !?\n'
            f'\n'
            f'stdout:\n'
            f'{outmsg}\n'
            f'\n'
            f'stderr:\n'
            f'{errmsg}\n'
        )

    # if we get some gnarly output let's aggregate and raise
    if errmsg:
        errlines = errmsg.splitlines()
        last_error = errlines[-1]
        if (
            'Error' in last_error

            # XXX: currently we print this to console, but maybe
            # shouldn't eventually once we figure out what's
            # a better way to be explicit about aio side
            # cancels?
            and
            'asyncio.exceptions.CancelledError' not in last_error
        ):
            raise Exception(errmsg)

    assert proc.returncode == 0


def test_wait_for_failed_example_captures_output():
    '''
    Preserve diagnostics from a subprocess which already exited.

    The previous `poll()` guard skipped `communicate()` when a fast
    failure returned a non-zero status before the parent checked it.
    Its stdout and stderr were therefore reported as empty. This
    fake process begins with `returncode=1` and returns non-UTF-8
    output, proving the helper always drains both pipes and replaces
    undecodable bytes without hiding the original process failure.

    '''
    proc = Mock()
    proc.returncode = 1
    proc.communicate.return_value = (
        b'stdout\xff',
        b'stderr\xff',
    )

    with pytest.raises(Exception) as exc_info:
        _wait_for_proc(
            proc=proc,
            timeout=1,
            test_log=Mock(),
        )

    proc.communicate.assert_called_once_with(timeout=1)
    errmsg: str = str(exc_info.value)
    assert 'stdout\ufffd' in errmsg
    assert 'stderr\ufffd' in errmsg


@pytest.mark.skipif(
    platform.system() == 'Windows',
    reason='POSIX process groups are unavailable on Windows',
)
def test_wait_for_timed_out_example_reaps_group(
    monkeypatch: pytest.MonkeyPatch,
):
    '''
    Kill the example process group and reap its leader on timeout.

    The old timeout branch killed only the immediate process and
    never drained it. Actor descendants could retain the capture
    pipes while the leader remained unreaped, hanging CI until its
    job timeout. This fake process raises `TimeoutExpired` on the
    timed wait and completes on the second `communicate()` call;
    the assertions prove group-directed `SIGKILL` precedes that
    final drain and leaves a concrete non-zero return code.

    '''
    proc = Mock()
    proc.pid = 1234

    def communicate(timeout=None):
        if timeout is not None:
            raise subprocess.TimeoutExpired('example', timeout)
        proc.returncode = -signal.SIGKILL
        return b'', b'timed out'

    proc.communicate.side_effect = communicate
    killpg = Mock()
    monkeypatch.setattr(os, 'killpg', killpg)

    with pytest.raises(Exception, match='timed out'):
        _wait_for_proc(
            proc=proc,
            timeout=.01,
            test_log=Mock(),
        )

    killpg.assert_called_once_with(1234, signal.SIGKILL)
    assert proc.communicate.call_count == 2
    assert proc.returncode == -signal.SIGKILL


@pytest.fixture
def run_example_in_subproc(
    loglevel: str,
    testdir: pytest.Pytester,
    reg_addr: tuple[str, int],
):

    @contextmanager
    def run(script_code):
        kwargs = dict()

        if platform.system() == 'Windows':
            # on windows we need to create a special __main__.py which will
            # be executed with ``python -m <modulename>`` on windows..
            shutil.copyfile(
                examples_dir() / '__main__.py',
                str(testdir / '__main__.py'),
            )

            # drop the ``if __name__ == '__main__'`` guard onwards from
            # the *NIX version of each script
            windows_script_lines = itertools.takewhile(
                lambda line: "if __name__ ==" not in line,
                script_code.splitlines()
            )
            script_code = '\n'.join(windows_script_lines)
            script_file = testdir.makefile('.py', script_code)

            # without this, tests hang on windows forever
            kwargs['creationflags'] = subprocess.CREATE_NEW_PROCESS_GROUP

            # run the testdir "libary module" as a script
            cmdargs = [
                sys.executable,
                '-m',
                # use the "module name" of this "package"
                'test_example'
            ]
        else:
            script_file = testdir.makefile('.py', script_code)
            kwargs['start_new_session'] = True
            cmdargs = [
                sys.executable,
                str(script_file),
            ]

        # XXX: BE FOREVER WARNED: if you enable lots of tractor logging
        # in the subprocess it may cause infinite blocking on the pipes
        # due to backpressure!!!
        proc = testdir.popen(
            cmdargs,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            **kwargs,
        )
        assert not proc.returncode
        try:
            yield proc
        except BaseException:
            if proc.poll() is None:
                try:
                    _kill_proc_tree(proc)
                    _reap_killed_proc(proc)
                except Exception:
                    pass
            raise
        else:
            if proc.poll() is None:
                _kill_proc_tree(proc)
                _reap_killed_proc(proc)

    yield run


@pytest.mark.parametrize(
    'example_script',

    # walk yields: (dirpath, dirnames, filenames)
    [
        (p[0], f)
        for p in os.walk(examples_dir())
        for f in p[2]

        if (
            '__' not in f  # ignore any pkg-mods
            # ignore any `__pycache__` subdir
            and '__pycache__' not in str(p[0])
            and f[0] != '_'  # ignore any WIP "examplel mods"
            and 'debugging' not in p[0]
            and 'integration' not in p[0]
            and 'advanced_faults' not in p[0]
            and 'multihost' not in p[0]
            and 'trio' not in p[0]
        )
    ],
    ids=lambda t: t[1],
)
def test_example(
    run_example_in_subproc: Callable,
    example_script: str,
    test_log: tractor.log.StackLevelAdapter,
    ci_env: bool,
):
    '''
    Load and run scripts from this repo's ``examples/`` dir as a user
    would copy and pasing them into their editor.

    On windows a little more "finessing" is done to make
    ``multiprocessing`` play nice: we copy the ``__main__.py`` into the
    test directory and invoke the script as a module with ``python -m
    test_example``.

    '''
    ex_file: str = os.path.join(*example_script)

    if (
        'rpc_bidir_streaming' in ex_file
        and
        sys.version_info < (3, 9)
    ):
        pytest.skip("2-way streaming example requires py3.9 async with syntax")

    if (
        'full_fledged_streaming_service' in ex_file
        and
        _friggin_macos
        and
        ci_env
    ):
        pytest.skip(
            'Streaming example is too flaky in CI\n'
            'AND their competitor runs this CI service..\n'
            'This test does run just fine "in person" however..'
        )

    from .conftest import cpu_perf_headroom

    timeout: float = (
        60
        if ci_env and _non_linux
        else 16
    )

    # add latency headroom for CPU freq scaling/throttle
    # (auto-cpufreq et al.)
    headroom: float = cpu_perf_headroom()
    if headroom != 1.:
        timeout *= headroom

    with open(ex_file, 'r') as ex:
        code = ex.read()

        with run_example_in_subproc(code) as proc:
            _wait_for_proc(
                proc=proc,
                timeout=timeout,
                test_log=test_log,
            )