from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timezone
import os
from pathlib import Path
import signal
import subprocess
import time
from typing import Sequence

from .models import ErrorCode


@dataclass(frozen=True)
class ProcessResult:
    argv_redacted: tuple[str, ...]
    started: bool
    pid: int | None
    exit_code: int | None
    timed_out: bool
    terminate_issued: bool
    exited: bool
    liveness_unknown: bool
    stdout_bytes: bytes
    stderr_bytes: bytes
    stdout_truncated: bool
    stderr_truncated: bool
    started_at_utc: str | None
    ended_at_utc: str | None
    elapsed_ms: int | None
    error_code: ErrorCode | None


class ProcessSupervisor:
    def __init__(self, *, max_stream_bytes: int = 1_048_576, kill_wait_ms: int = 5_000):
        self.max_stream_bytes = max_stream_bytes
        self.kill_wait_ms = kill_wait_ms

    @staticmethod
    def _utc_now() -> str:
        return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")

    def run(self, argv: Sequence[str], *, timeout_ms: int, cwd: Path | None = None,
            stdin_bytes: bytes | None = None, redacted_argv: Sequence[str] | None = None) -> ProcessResult:
        if not argv or timeout_ms <= 0:
            raise ValueError("argv and positive timeout required")
        shown = tuple(redacted_argv or argv)
        started_at = self._utc_now()
        begin = time.monotonic()
        deadline = begin + timeout_ms / 1000
        creationflags = subprocess.CREATE_NEW_PROCESS_GROUP if os.name == "nt" else 0
        try:
            proc = subprocess.Popen(
                tuple(argv), cwd=str(cwd) if cwd else None, stdin=subprocess.PIPE if stdin_bytes is not None else subprocess.DEVNULL,
                stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=False, creationflags=creationflags,
                start_new_session=(os.name != "nt"),
            )
        except OSError:
            return ProcessResult(shown, False, None, None, False, False, True, False, b"", b"", False,
                                 False, started_at, self._utc_now(), round((time.monotonic()-begin)*1000),
                                 ErrorCode.PROCESS_START_FAILED)
        timed_out = False
        terminate_issued = False
        liveness_unknown = False
        try:
            termination_reserve_ms = min(self.kill_wait_ms, max(100, min(1_000, timeout_ms // 5)),
                                         max(1, timeout_ms - 1))
            normal_wait = max(0.001, (timeout_ms - termination_reserve_ms) / 1000)
            out, err = proc.communicate(stdin_bytes, timeout=normal_wait)
        except subprocess.TimeoutExpired:
            timed_out = True
            terminate_issued = True
            self._terminate_tree(proc, max(0.001, deadline - time.monotonic()))
            try:
                kill_remaining = max(0.0, deadline - time.monotonic())
                if kill_remaining <= 0:
                    raise subprocess.TimeoutExpired(tuple(argv), timeout_ms / 1000)
                out, err = proc.communicate(timeout=min(self.kill_wait_ms / 1000, kill_remaining))
            except subprocess.TimeoutExpired:
                out, err = b"", b""
                liveness_unknown = True
                try:
                    proc.kill()
                except OSError:
                    pass
        elapsed = round((time.monotonic() - begin) * 1000)
        out_trim, out_trunc = self._cap(out)
        err_trim, err_trunc = self._cap(err)
        exited = proc.poll() is not None
        code = None
        if liveness_unknown:
            code = ErrorCode.PROCESS_LIVENESS_UNKNOWN
        elif timed_out:
            code = ErrorCode.PROCESS_TIMEOUT
        elif out_trunc or err_trunc:
            code = ErrorCode.PROCESS_OUTPUT_LIMIT
        return ProcessResult(shown, True, proc.pid, proc.returncode if exited else None, timed_out, terminate_issued,
                             exited, liveness_unknown, out_trim, err_trim, out_trunc, err_trunc,
                             started_at, self._utc_now(), elapsed, code)

    def _cap(self, data: bytes) -> tuple[bytes, bool]:
        if len(data) <= self.max_stream_bytes:
            return data, False
        return data[:self.max_stream_bytes], True

    @staticmethod
    def _terminate_tree(proc: subprocess.Popen[bytes], timeout_seconds: float = 5.0) -> None:
        try:
            if os.name == "nt":
                subprocess.run(("taskkill", "/PID", str(proc.pid), "/T", "/F"), shell=False,
                               stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                               timeout=max(0.001, min(5.0, timeout_seconds)))
            else:
                os.killpg(proc.pid, signal.SIGKILL)
        except (OSError, subprocess.SubprocessError):
            try:
                proc.kill()
            except OSError:
                pass