from __future__ import annotations

import queue
import subprocess
import threading
import time
from typing import Iterable, Sequence


__all__ = ["Task", "TaskTimeout"]


class TaskTimeout(Exception):
    def __init__(self, timeout):
        super().__init__(f"{timeout}s elapsed")


class Task:
    """
    A non-blocking subprocess runner.

    Example:

        >>> task = Task(["sleep", "1"])
        >>> for i in task.loop(sleep=0.4):
        ...     print(i)
        0
        1
        2

    Uses some of the insights from the following thread:
    https://stackoverflow.com/q/375427
    """

    args: Sequence | str | bytes
    code: int | None
    stdout: bytes | None
    stderr: bytes | None
    _process: subprocess.Popen | None
    _thread: threading.Thread | None
    _queue: queue.Queue

    def __init__(self, args: Sequence | str | bytes) -> None:
        self.args = args
        self.code = None
        self.stdout = None
        self.stderr = None
        self._process = None
        self._thread = None
        self._queue = queue.Queue()

    def loop(self, sleep: float | int | None = None) -> Iterable[int]:
        """
        Returns an iterator of integers starting at 0, ending when the
        subprocess exits.

        If sleep is specified, at the end of each loop, this sleeps for sleep
        seconds.
        """
        i = 0
        self.start()
        while i == 0 or self.running():
            yield i
            i += 1
            if sleep is not None and self.running():
                time.sleep(sleep)

    def poll(self) -> int | None:
        """
        Returns the process exit status or None.

        This will set code, stderr, stdout if the subprocess has exited.
        """
        if self._process:
            if self.code is None:
                try:
                    self.wait(0)
                except TaskTimeout:
                    return None
            return self.code
        else:
            return None

    def running(self) -> bool:
        """Returns whether or not the subprocess is running.

        This will set code, stderr, stdout if the subprocess has exited.
        """
        if self._process:
            return self.poll() is None
        else:
            return False

    def start(self) -> None:
        """Starts the subprocess if it hasn't already been started."""

        if self._process:
            return

        self._process = subprocess.Popen(
            self.args,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            shell=isinstance(self.args, (str, bytes)),
        )

        def target():
            stdout, stderr = b"", b""
            try:
                stdout, stderr = self._process.communicate()
            finally:
                self._queue.put((stdout, stderr))

        self._thread = threading.Thread(target=target, daemon=True)
        self._thread.start()

    def wait(self, timeout: float | int | None = None) -> int:
        """
        Waits for the subprocess to exit and returns its exit status.

        This will set code, stderr, stdout if the subprocess has exited.

        If the subprocess wasn't already started, it will be started.

        If timeout is not specified, then this blocks until the subprocess has
        exited.

        If timeout is specified, then this blocks at most timeout seconds. If
        the subprocess hasn't exited, then this raises a TimeoutExpired
        exception.
        """

        self.start()
        if not self._process:
            raise RuntimeError("expected self._process to be set")

        if self.code is None:
            block = timeout is None or timeout > 0
            try:
                self.stdout, self.stderr = self._queue.get(block=block, timeout=timeout)
            except queue.Empty:
                raise TaskTimeout(timeout)

            self.code = self._process.returncode
            if self.code is None:  # pragma: no cover
                self.code = -255
                raise RuntimeError("expected subprocess to set returncode")

        return self.code
