"""
subproc.py — single-file Unix-first controlled child process launcher.

Goals:
- Spawn a child process under a parent process with strong lifecycle control.
- Control env/cwd/stdin/stdout/stderr.
- Provide streaming output hooks and safe termination escalation.
- Unix-first (Linux/macOS). Windows is not a priority.

Design choices:
- By default, we start the child in a new session (new process group) so the parent can
  signal / terminate the entire process tree using killpg().
- Output can be captured (pipes) and streamed via callbacks, or inherited, or redirected.
- Optional PTY support for "interactive-like" behavior (useful for tools that behave
  differently when not on a TTY).

Usage example:

    from subproc import ControlledProcess

    def on_out(line): print("OUT:", line, end="")
    def on_err(line): print("ERR:", line, end="")

    p = ControlledProcess(["bash", "-lc", "echo hi; echo err 1>&2; sleep 10"], text=True)
    p.start(on_stdout=on_out, on_stderr=on_err)
    p.wait(timeout=2)  # may raise TimeoutError
    p.stop(grace=1.0)

"""

from __future__ import annotations

import errno
import os
import selectors
import signal
import subprocess
import threading
import time
from dataclasses import dataclass
from typing import (
    Any,
    Callable,
    Dict,
    List,
    Mapping,
    Optional,
    Sequence,
    Tuple,
    Union,
    IO,
)

BytesOrStr = Union[bytes, str]
LineCallback = Callable[[str], None]
RawCallback = Callable[[bytes], None]


@dataclass(frozen=True)
class Completed:
    args: Sequence[str]
    returncode: int
    stdout: Optional[Union[str, bytes]] = None
    stderr: Optional[Union[str, bytes]] = None
    pid: Optional[int] = None
    pgid: Optional[int] = None
    start_time: Optional[float] = None
    end_time: Optional[float] = None

    @property
    def duration_s(self) -> Optional[float]:
        if self.start_time is None or self.end_time is None:
            return None
        return self.end_time - self.start_time


class ControlledProcess:
    """
    A controlled child process wrapper.

    Key features (Unix-first):
    - start_new_session=True by default -> child runs in a new process group/session.
    - killpg() to terminate the entire process group.
    - stream stdout/stderr to callbacks using a selector thread.
    - optional capture of stdout/stderr to buffers.
    - write to stdin (bytes or str) and close stdin when done.
    """

    def __init__(
        self,
        args: Sequence[str],
        *,
        cwd: Optional[str] = None,
        env: Optional[Mapping[str, str]] = None,
        inherit_env: bool = True,
        stdin: Union[int, IO[Any], None] = subprocess.PIPE,
        stdout: Union[int, IO[Any], None] = subprocess.PIPE,
        stderr: Union[int, IO[Any], None] = subprocess.PIPE,
        text: bool = True,
        encoding: str = "utf-8",
        errors: str = "replace",
        bufsize: int = 0,
        start_new_session: bool = True,
        pty: bool = False,
        # If capture is True, we'll store streamed output in memory as well.
        capture: bool = True,
        # If merge_stderr is True, stderr is redirected to stdout.
        merge_stderr: bool = False,
    ) -> None:
        self.args = list(args)
        self.cwd = cwd
        self._user_env = dict(env) if env is not None else None
        self.inherit_env = inherit_env

        self.stdin_spec = stdin
        self.stdout_spec = stdout
        self.stderr_spec = stderr

        self.text = text
        self.encoding = encoding
        self.errors = errors
        self.bufsize = bufsize

        self.start_new_session = start_new_session
        self.pty = pty
        self.capture = capture
        self.merge_stderr = merge_stderr

        self._proc: Optional[subprocess.Popen] = None
        self._pgid: Optional[int] = None

        self._start_time: Optional[float] = None
        self._end_time: Optional[float] = None

        self._stdout_buf: bytearray = bytearray()
        self._stderr_buf: bytearray = bytearray()

        self._io_thread: Optional[threading.Thread] = None
        self._io_stop = threading.Event()

        self._pty_master_fd: Optional[int] = None
        self._pty_slave_fd: Optional[int] = None

        self._lock = threading.RLock()

    # -------------------------
    # Lifecycle & properties
    # -------------------------

    @property
    def proc(self) -> subprocess.Popen:
        if self._proc is None:
            raise RuntimeError("Process not started")
        return self._proc

    @property
    def pid(self) -> Optional[int]:
        return None if self._proc is None else self._proc.pid

    @property
    def pgid(self) -> Optional[int]:
        return self._pgid

    @property
    def returncode(self) -> Optional[int]:
        return None if self._proc is None else self._proc.returncode

    def is_running(self) -> bool:
        return self._proc is not None and self._proc.poll() is None

    # -------------------------
    # Spawn / IO management
    # -------------------------

    def start(
        self,
        *,
        on_stdout: Optional[LineCallback] = None,
        on_stderr: Optional[LineCallback] = None,
        on_stdout_bytes: Optional[RawCallback] = None,
        on_stderr_bytes: Optional[RawCallback] = None,
    ) -> "ControlledProcess":
        """
        Start the process and (if piping) begin streaming output.

        Callbacks:
        - If text=True and on_stdout/on_stderr are provided, callbacks receive decoded text lines.
        - If on_*_bytes is provided, callback receives raw bytes chunks.
        You can provide both line and bytes callbacks; bytes callback is called first.
        """
        with self._lock:
            if self._proc is not None:
                raise RuntimeError("Process already started")

            env = self._build_env()

            stdin_spec, stdout_spec, stderr_spec = self._build_stdio_specs()

            self._start_time = time.time()

            # For PTY: both stdout and stderr go to PTY slave (tty-like); stdin from PTY slave too.
            # This is commonly used to convince programs they're running in a terminal.
            self._proc = subprocess.Popen(
                self.args,
                cwd=self.cwd,
                env=env,
                stdin=stdin_spec,
                stdout=stdout_spec,
                stderr=stderr_spec,
                bufsize=self.bufsize,
                start_new_session=self.start_new_session,
                text=False,  # we always read bytes internally; decode ourselves if needed
                close_fds=True,
            )

            # Determine process group id for later killpg.
            try:
                self._pgid = os.getpgid(self._proc.pid)
            except Exception:
                # If start_new_session=False, PGID might still exist; or on some systems this can race.
                self._pgid = None

            # Start IO thread if we have pipes/pty master to read.
            self._start_io_thread(
                on_stdout=on_stdout,
                on_stderr=on_stderr,
                on_stdout_bytes=on_stdout_bytes,
                on_stderr_bytes=on_stderr_bytes,
            )

        return self

    def _build_env(self) -> Dict[str, str]:
        if self.inherit_env:
            base = dict(os.environ)
        else:
            base = {}
        if self._user_env:
            base.update({k: str(v) for k, v in self._user_env.items()})
        return base

    def _build_stdio_specs(self) -> Tuple[Any, Any, Any]:
        if self.pty:
            # Create a PTY pair; child uses slave as stdio, parent reads/writes master.
            import pty  # Unix-only

            master, slave = pty.openpty()
            self._pty_master_fd, self._pty_slave_fd = master, slave

            # child stdio:
            stdin_spec = slave
            stdout_spec = slave
            if self.merge_stderr:
                stderr_spec = subprocess.STDOUT
            else:
                stderr_spec = slave
            return stdin_spec, stdout_spec, stderr_spec

        # Non-PTY behavior: allow caller to specify stdin/stdout/stderr.
        stdin_spec = self.stdin_spec
        stdout_spec = self.stdout_spec
        if self.merge_stderr:
            stderr_spec = subprocess.STDOUT
        else:
            stderr_spec = self.stderr_spec
        return stdin_spec, stdout_spec, stderr_spec

    def _start_io_thread(
        self,
        *,
        on_stdout: Optional[LineCallback],
        on_stderr: Optional[LineCallback],
        on_stdout_bytes: Optional[RawCallback],
        on_stderr_bytes: Optional[RawCallback],
    ) -> None:
        # If stdout/stderr are not pipes (or PTY), nothing to stream.
        if self._proc is None:
            return

        if self.pty:
            # We read from the PTY master only; stdout/stderr are merged in PTY mode.
            if self._pty_master_fd is None:
                return
            self._io_stop.clear()
            t = threading.Thread(
                target=self._io_loop_pty,
                name="ControlledProcess-io-pty",
                daemon=True,
                kwargs=dict(
                    master_fd=self._pty_master_fd,
                    on_bytes=on_stdout_bytes,
                    on_line=on_stdout,  # treat all output as "stdout"
                ),
            )
            self._io_thread = t
            t.start()
            return

        stdout_pipe = self._proc.stdout
        stderr_pipe = self._proc.stderr

        # If merging stderr to stdout, stderr_pipe may be None.
        has_any = stdout_pipe is not None or stderr_pipe is not None
        if not has_any:
            return

        self._io_stop.clear()
        t = threading.Thread(
            target=self._io_loop_pipes,
            name="ControlledProcess-io",
            daemon=True,
            kwargs=dict(
                stdout_pipe=stdout_pipe,
                stderr_pipe=stderr_pipe,
                on_stdout=on_stdout,
                on_stderr=on_stderr,
                on_stdout_bytes=on_stdout_bytes,
                on_stderr_bytes=on_stderr_bytes,
            ),
        )
        self._io_thread = t
        t.start()

    def _io_loop_pty(
        self,
        *,
        master_fd: int,
        on_bytes: Optional[RawCallback],
        on_line: Optional[LineCallback],
    ) -> None:
        sel = selectors.DefaultSelector()
        try:
            os.set_blocking(master_fd, False)
        except Exception:
            pass

        sel.register(master_fd, selectors.EVENT_READ)

        pending = bytearray()
        try:
            while not self._io_stop.is_set():
                events = sel.select(timeout=0.1)
                if not events:
                    if self._proc and self._proc.poll() is not None:
                        break
                    continue
                for key, _mask in events:
                    fd = key.fileobj
                    try:
                        chunk = os.read(fd, 4096)
                    except OSError as e:
                        if e.errno in (errno.EAGAIN, errno.EWOULDBLOCK):
                            continue
                        return
                    if not chunk:
                        return

                    if self.capture:
                        with self._lock:
                            self._stdout_buf.extend(chunk)

                    if on_bytes:
                        on_bytes(chunk)

                    if on_line and self.text:
                        pending.extend(chunk)
                        self._emit_lines(pending, on_line)
        finally:
            try:
                sel.unregister(master_fd)
            except Exception:
                pass
            try:
                sel.close()
            except Exception:
                pass

    def _io_loop_pipes(
        self,
        *,
        stdout_pipe: Optional[Any],
        stderr_pipe: Optional[Any],
        on_stdout: Optional[LineCallback],
        on_stderr: Optional[LineCallback],
        on_stdout_bytes: Optional[RawCallback],
        on_stderr_bytes: Optional[RawCallback],
    ) -> None:
        sel = selectors.DefaultSelector()

        def register(pipe: Any) -> None:
            if pipe is None:
                return
            try:
                os.set_blocking(pipe.fileno(), False)
            except Exception:
                pass
            sel.register(pipe, selectors.EVENT_READ)

        register(stdout_pipe)
        register(stderr_pipe)

        pending_out = bytearray()
        pending_err = bytearray()

        try:
            while not self._io_stop.is_set():
                events = sel.select(timeout=0.1)
                if not events:
                    if self._proc and self._proc.poll() is not None:
                        # Drain remaining data quickly
                        self._drain_pipe(
                            stdout_pipe,
                            pending_out,
                            on_stdout,
                            on_stdout_bytes,
                            is_stdout=True,
                        )
                        self._drain_pipe(
                            stderr_pipe,
                            pending_err,
                            on_stderr,
                            on_stderr_bytes,
                            is_stdout=False,
                        )
                        break
                    continue

                for key, _mask in events:
                    pipe = key.fileobj
                    is_stdout = stdout_pipe is not None and pipe == stdout_pipe

                    try:
                        chunk = pipe.read(4096)  # bytes because text=False in Popen
                    except Exception:
                        continue
                    if not chunk:
                        continue

                    if self.capture:
                        with self._lock:
                            if is_stdout:
                                self._stdout_buf.extend(chunk)
                            else:
                                self._stderr_buf.extend(chunk)

                    if is_stdout:
                        if on_stdout_bytes:
                            on_stdout_bytes(chunk)
                        if on_stdout and self.text:
                            pending_out.extend(chunk)
                            self._emit_lines(pending_out, on_stdout)
                    else:
                        if on_stderr_bytes:
                            on_stderr_bytes(chunk)
                        if on_stderr and self.text:
                            pending_err.extend(chunk)
                            self._emit_lines(pending_err, on_stderr)
        finally:
            try:
                if stdout_pipe is not None:
                    sel.unregister(stdout_pipe)
            except Exception:
                pass
            try:
                if stderr_pipe is not None:
                    sel.unregister(stderr_pipe)
            except Exception:
                pass
            try:
                sel.close()
            except Exception:
                pass

    def _drain_pipe(
        self,
        pipe: Optional[Any],
        pending: bytearray,
        on_line: Optional[LineCallback],
        on_bytes: Optional[RawCallback],
        *,
        is_stdout: bool,
    ) -> None:
        if pipe is None:
            return
        try:
            while True:
                chunk = pipe.read(4096)
                if not chunk:
                    break
                if self.capture:
                    with self._lock:
                        if is_stdout:
                            self._stdout_buf.extend(chunk)
                        else:
                            self._stderr_buf.extend(chunk)
                if on_bytes:
                    on_bytes(chunk)
                if on_line and self.text:
                    pending.extend(chunk)
                    self._emit_lines(pending, on_line)
        except Exception:
            return

    def _emit_lines(self, pending: bytearray, cb: LineCallback) -> None:
        """
        Emit complete lines from pending buffer; keep partial line buffered.
        """
        while True:
            i = pending.find(b"\n")
            if i < 0:
                return
            line_bytes = pending[: i + 1]
            del pending[: i + 1]
            try:
                line = line_bytes.decode(self.encoding, errors=self.errors)
            except Exception:
                line = line_bytes.decode("utf-8", errors="replace")
            cb(line)

    # -------------------------
    # Parent -> child input
    # -------------------------

    def write(self, data: BytesOrStr) -> None:
        """
        Write to the child's stdin.

        - In PTY mode, writes to PTY master.
        - In pipe mode, writes to proc.stdin.
        """
        with self._lock:
            if not self.is_running():
                raise RuntimeError("Process not running")

            if self.pty:
                if self._pty_master_fd is None:
                    raise RuntimeError("PTY not initialized")
                b = (
                    data.encode(self.encoding, errors=self.errors)
                    if isinstance(data, str)
                    else data
                )
                os.write(self._pty_master_fd, b)
                return

            stdin = self.proc.stdin
            if stdin is None:
                raise RuntimeError("stdin is not a pipe")
            b = (
                data.encode(self.encoding, errors=self.errors)
                if isinstance(data, str)
                else data
            )
            try:
                stdin.write(b)
                stdin.flush()
            except BrokenPipeError:
                # Child closed stdin; ignore.
                return

    def close_stdin(self) -> None:
        with self._lock:
            if self._proc is None:
                return
            if self.pty:
                # PTY stdin close: close master write side (closing master fd suffices).
                if self._pty_master_fd is not None:
                    try:
                        os.close(self._pty_master_fd)
                    except Exception:
                        pass
                    self._pty_master_fd = None
                return

            if self._proc.stdin is not None:
                try:
                    self._proc.stdin.close()
                except Exception:
                    pass

    # -------------------------
    # Waiting / completion
    # -------------------------

    def wait(self, timeout: Optional[float] = None) -> int:
        """
        Wait for completion. Returns returncode. Raises TimeoutError on timeout.
        """
        if self._proc is None:
            raise RuntimeError("Process not started")
        try:
            rc = self._proc.wait(timeout=timeout)
        except subprocess.TimeoutExpired as e:
            raise TimeoutError(str(e)) from e
        finally:
            if self._proc and self._proc.poll() is not None:
                self._mark_ended()
        return rc

    def poll(self) -> Optional[int]:
        if self._proc is None:
            return None
        rc = self._proc.poll()
        if rc is not None:
            self._mark_ended()
        return rc

    def _mark_ended(self) -> None:
        with self._lock:
            if self._end_time is None:
                self._end_time = time.time()

    def communicate(
        self,
        input: Optional[BytesOrStr] = None,
        timeout: Optional[float] = None,
    ) -> Tuple[Optional[Union[str, bytes]], Optional[Union[str, bytes]]]:
        """
        Like subprocess.communicate, but respects our internal capture and PTY mode.

        If you use communicate(), prefer NOT to also use streaming callbacks, since both read stdout.
        (This module supports both patterns, but mixing them can be confusing.)
        """
        if self._proc is None:
            raise RuntimeError("Process not started")

        if self.pty:
            # In PTY mode, we don't have separate stdout/stderr pipes.
            if input is not None:
                self.write(input)
                self.close_stdin()
            rc = self.wait(timeout=timeout)
            _ = rc
            out = self.stdout_text() if self.text else self.stdout_bytes()
            return out, None

        b_in: Optional[bytes]
        if input is None:
            b_in = None
        else:
            b_in = (
                input.encode(self.encoding, errors=self.errors)
                if isinstance(input, str)
                else input
            )

        try:
            out, err = self._proc.communicate(input=b_in, timeout=timeout)
        except subprocess.TimeoutExpired as e:
            raise TimeoutError(str(e)) from e
        finally:
            self._mark_ended()

        if self.text:
            out_s = (
                out.decode(self.encoding, errors=self.errors)
                if out is not None
                else None
            )
            err_s = (
                err.decode(self.encoding, errors=self.errors)
                if err is not None
                else None
            )
            return out_s, err_s
        return out, err

    # -------------------------
    # Signals / termination
    # -------------------------

    def send_signal(self, sig: int) -> None:
        with self._lock:
            if not self.is_running():
                return
            self.proc.send_signal(sig)

    def terminate(self) -> None:
        with self._lock:
            if not self.is_running():
                return
            self.proc.terminate()

    def kill(self) -> None:
        with self._lock:
            if not self.is_running():
                return
            self.proc.kill()

    def killpg(self, sig: int = signal.SIGTERM) -> None:
        """
        Send a signal to the entire process group (Unix). Requires we know pgid.
        """
        with self._lock:
            if not self.is_running():
                return
            if self._pgid is None:
                # Fallback: signal just the process.
                self.proc.send_signal(sig)
                return
            try:
                os.killpg(self._pgid, sig)
            except ProcessLookupError:
                return

    def stop(
        self,
        *,
        grace: float = 3.0,
        sigterm: int = signal.SIGTERM,
        sigkill: int = signal.SIGKILL,
        use_process_group: bool = True,
    ) -> None:
        """
        Stop the process gracefully, then escalate.

        - Sends SIGTERM (to group if use_process_group else to pid)
        - Waits `grace` seconds
        - Sends SIGKILL if still running
        """
        if self._proc is None:
            return

        if not self.is_running():
            self._finalize()
            return

        if use_process_group:
            self.killpg(sigterm)
        else:
            self.send_signal(sigterm)

        try:
            self.wait(timeout=grace)
            self._finalize()
            return
        except TimeoutError:
            pass

        if use_process_group:
            self.killpg(sigkill)
        else:
            self.send_signal(sigkill)

        try:
            self.wait(timeout=grace)
        except TimeoutError:
            # If kill didn't work, nothing else we can do here.
            pass
        finally:
            self._finalize()

    # -------------------------
    # Output access
    # -------------------------

    def stdout_bytes(self) -> bytes:
        with self._lock:
            return bytes(self._stdout_buf)

    def stderr_bytes(self) -> bytes:
        with self._lock:
            return bytes(self._stderr_buf)

    def stdout_text(self) -> str:
        b = self.stdout_bytes()
        return b.decode(self.encoding, errors=self.errors)

    def stderr_text(self) -> str:
        b = self.stderr_bytes()
        return b.decode(self.encoding, errors=self.errors)

    def result(self) -> Completed:
        if self._proc is None:
            raise RuntimeError("Process not started")
        if self._proc.poll() is None:
            raise RuntimeError("Process still running")
        self._mark_ended()

        if self.capture:
            out: Optional[Union[str, bytes]] = (
                self.stdout_text() if self.text else self.stdout_bytes()
            )
            err: Optional[Union[str, bytes]]
            if self.merge_stderr or self.pty:
                err = None
            else:
                err = self.stderr_text() if self.text else self.stderr_bytes()
        else:
            out = None
            err = None

        return Completed(
            args=self.args,
            returncode=int(self._proc.returncode),
            stdout=out,
            stderr=err,
            pid=self.pid,
            pgid=self.pgid,
            start_time=self._start_time,
            end_time=self._end_time,
        )

    # -------------------------
    # Cleanup
    # -------------------------

    def _finalize(self) -> None:
        # Stop IO thread, close pipes/pty fds.
        self._io_stop.set()

        t = self._io_thread
        if t is not None and t.is_alive():
            t.join(timeout=0.5)

        with self._lock:
            # Close PTY fds
            if self._pty_slave_fd is not None:
                try:
                    os.close(self._pty_slave_fd)
                except Exception:
                    pass
                self._pty_slave_fd = None

            if self._pty_master_fd is not None:
                try:
                    os.close(self._pty_master_fd)
                except Exception:
                    pass
                self._pty_master_fd = None

            # Close proc pipes
            if self._proc is not None:
                for stream_name in ("stdin", "stdout", "stderr"):
                    s = getattr(self._proc, stream_name, None)
                    if s is not None:
                        try:
                            s.close()
                        except Exception:
                            pass

    def __enter__(self) -> "ControlledProcess":
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        try:
            self.stop()
        except Exception:
            pass


# -------------------------
# Convenience helpers
# -------------------------


def run(
    args: Sequence[str],
    *,
    cwd: Optional[str] = None,
    env: Optional[Mapping[str, str]] = None,
    inherit_env: bool = True,
    input: Optional[BytesOrStr] = None,
    timeout: Optional[float] = None,
    text: bool = True,
    encoding: str = "utf-8",
    errors: str = "replace",
    pty: bool = False,
    merge_stderr: bool = False,
) -> Completed:
    """
    Convenience function: spawn, optionally send input, wait, return Completed.
    """
    p = ControlledProcess(
        args,
        cwd=cwd,
        env=env,
        inherit_env=inherit_env,
        text=text,
        encoding=encoding,
        errors=errors,
        pty=pty,
        merge_stderr=merge_stderr,
        capture=True,
    )
    p.start()
    try:
        p.communicate(input=input, timeout=timeout)
    finally:
        if p.is_running():
            p.stop()
    return p.result()
