"""Spawn a process and stream its stdout/stderr as an async generator.

Mirrors ``strands-ts/src/sandbox/stream-process.ts``.

The shared process-supervision engine behind the shell-based backends (Docker,
SSH): a backend builds an argv, and this spawns the process, streams its output
as :class:`StreamChunk` objects, and yields a final :class:`ExecutionResult`.

Cancellation is cooperative -- cancelling the consuming task (or closing the
generator) runs the cleanup in ``finally``, which kills the process. ``timeout``
is wall-clock: measured from spawn and not reset by ongoing output.
"""

import asyncio
import contextlib
import os
import signal
from collections.abc import AsyncGenerator

from .types import ExecutionResult, StreamChunk, StreamType

_READ_CHUNK_SIZE = 65536
_SIGNAL_EXIT_BASE = 128

# The child is spawned in its own process group (start_new_session) so the whole
# tree can be killed at once. Without this, "sh -c 'cmd; sleep 60'" leaves the
# sleep child holding the pipe write-end open after the parent dies, so the
# readers never see EOF and the generator hangs.
_USE_PROCESS_GROUP = hasattr(os, "killpg")


def _kill_tree(proc: asyncio.subprocess.Process) -> None:
    """SIGKILL the entire process tree (the group on POSIX, the lone process elsewhere)."""
    with contextlib.suppress(ProcessLookupError, PermissionError):
        if _USE_PROCESS_GROUP:
            os.killpg(proc.pid, signal.SIGKILL)
        else:
            proc.kill()


async def stream_process(
    program: str,
    args: list[str],
    *,
    timeout: float | None = None,
    enoent_message: str | None = None,
) -> AsyncGenerator[StreamChunk | ExecutionResult, None]:
    """Spawn a command and stream its stdout/stderr, yielding the final result.

    Yields :class:`StreamChunk` objects as output arrives, then a single final
    :class:`ExecutionResult`. Signal termination maps to ``128 + signal`` (e.g.
    SIGKILL -> 137). Cancelling the consuming task kills the process.

    Args:
        program: The binary to spawn (e.g. ``"docker"``, ``"ssh"``).
        args: Arguments to pass to the binary.
        timeout: Maximum wall-clock execution time in seconds, measured from
            spawn (not reset by output). ``None`` means no timeout.
        enoent_message: Message to surface as ``stderr`` (with exit code 127)
            when ``program`` is not on PATH. If ``None``, the underlying
            :class:`FileNotFoundError` propagates instead.

    Yields:
        :class:`StreamChunk` objects for output, then a final
        :class:`ExecutionResult`.

    Raises:
        TimeoutError: If execution exceeds ``timeout`` seconds.
        FileNotFoundError: If ``program`` is not found and ``enoent_message`` is ``None``.
    """
    try:
        proc = await asyncio.create_subprocess_exec(
            program,
            *args,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            start_new_session=_USE_PROCESS_GROUP,
        )
    except FileNotFoundError:
        if enoent_message is None:
            raise
        yield ExecutionResult(exit_code=127, stdout="", stderr=enoent_message)
        return

    queue: asyncio.Queue[StreamChunk | None] = asyncio.Queue()
    out_buf: list[str] = []
    err_buf: list[str] = []
    timed_out = False

    async def pump(stream: asyncio.StreamReader, stream_type: StreamType, buf: list[str]) -> None:
        while data := await stream.read(_READ_CHUNK_SIZE):
            text = data.decode(errors="replace")
            buf.append(text)
            await queue.put(StreamChunk(data=text, stream_type=stream_type))

    assert proc.stdout is not None and proc.stderr is not None
    pumps = asyncio.gather(pump(proc.stdout, "stdout", out_buf), pump(proc.stderr, "stderr", err_buf))

    async def reader() -> None:
        await pumps
        await proc.wait()
        await queue.put(None)

    async def enforce_timeout(timeout: float) -> None:
        nonlocal timed_out
        await asyncio.sleep(timeout)
        if proc.returncode is None:
            timed_out = True
            _kill_tree(proc)
            await queue.put(None)  # signal the consumer loop to stop

    tasks = [asyncio.ensure_future(reader())]
    if timeout is not None:
        tasks.append(asyncio.ensure_future(enforce_timeout(timeout)))

    try:
        while (item := await queue.get()) is not None:
            yield item

        if timed_out:
            raise TimeoutError(f"Execution timed out after {timeout} seconds")

        returncode = proc.returncode if proc.returncode is not None else 1
        exit_code = _SIGNAL_EXIT_BASE - returncode if returncode < 0 else returncode
        yield ExecutionResult(exit_code=exit_code, stdout="".join(out_buf), stderr="".join(err_buf))
    finally:
        if proc.returncode is None:
            _kill_tree(proc)
        pumps.cancel()
        for task in tasks:
            task.cancel()
        await asyncio.gather(pumps, *tasks, return_exceptions=True)
