from __future__ import annotations

import codecs
import math
from contextlib import aclosing, suppress
from functools import partial
from subprocess import PIPE
from typing import TYPE_CHECKING, BinaryIO, TypedDict

import trio


if TYPE_CHECKING:
    from collections.abc import AsyncGenerator, Callable

    from typing_extensions import Unpack

    class TRunArgs(TypedDict, total=False):
        stdin: int | bytes | BinaryIO | None


class ProcessOutput:
    _send_channel: trio.MemorySendChannel[bool]
    _receive_channel: trio.MemoryReceiveChannel[bool]

    def __init__(
        self,
        command: list[str],
        timeout: float = math.inf,
        wait_terminate: float = 2.0,
        stdin: int | bytes | BinaryIO = b"",
    ):
        self.command = command
        self.timeout = timeout
        self.wait_terminate = wait_terminate
        self.stdin = stdin
        self._send_channel, self._receive_channel = trio.open_memory_channel(1)
        self._receive_max_bytes: int | None = None

    def run(self, **kwargs: Unpack[TRunArgs]) -> bool:  # pragma: no cover
        return trio.run(partial(self.arun, **kwargs))

    async def arun(self, **kwargs: Unpack[TRunArgs]) -> bool:
        stdin = kwargs.get("stdin")
        with trio.move_on_after(self.timeout):
            async with trio.open_nursery() as nursery:
                run_process = partial(
                    trio.run_process,
                    self.command,
                    check=False,
                    capture_stdout=False,
                    capture_stderr=False,
                    stdin=self.stdin if stdin is None else stdin,
                    stdout=PIPE,
                    stderr=PIPE,
                    deliver_cancel=self._deliver_cancel,
                )
                process: trio.Process = await nursery.start(run_process)

                nursery.start_soon(self._onexit, process)
                if process.stdout:  # pragma: no branch
                    nursery.start_soon(self._onoutput, self.onstdout, process.stdout)
                if process.stderr:  # pragma: no branch
                    nursery.start_soon(self._onoutput, self.onstderr, process.stderr)

                res = await self._receive_channel.receive()
                nursery.cancel_scope.cancel()
                return res

        # noinspection PyUnreachableCode
        return False

    async def _deliver_cancel(self, proc: trio.Process):
        with suppress(OSError):
            proc.terminate()
            await trio.sleep(self.wait_terminate)
            proc.kill()

    async def _onexit(self, proc: trio.Process):
        code = await proc.wait()
        result = self.onexit(code)
        await self._send_channel.send(result)

    async def _onoutput(self, callback: Callable[[int, str], bool | None], stream: trio.abc.ReceiveStream):
        idx = 0
        async with aclosing(self._line_decode(stream)) as iter_line_decode:
            async for line in iter_line_decode:
                try:
                    content = line.strip()
                    result = callback(idx, content)
                except Exception:
                    raise
                if result is not None:
                    await self._send_channel.send(bool(result))
                    break
                idx += 1

    async def _line_decode(self, stream: trio.abc.ReceiveStream) -> AsyncGenerator[str, None]:
        decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")
        buffer: str = ""
        while chunk := await stream.receive_some(max_bytes=self._receive_max_bytes):
            buffer += decoder.decode(chunk, final=False)
            while buffer:
                line, nl, buffer = buffer.partition("\n")
                if nl:
                    yield line
                else:
                    buffer = line
                    break
        if tail := decoder.decode(b"", final=True):
            buffer += tail
        if buffer:
            yield buffer

    def onexit(self, code: int) -> bool:
        return code == 0

    def onstdout(self, idx: int, line: str) -> bool | None:
        pass

    def onstderr(self, idx: int, line: str) -> bool | None:
        pass
