from __future__ import annotations

import os
import re
import signal
import subprocess
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Final

from k1link.artifacts import write_json_atomic

PROCESS_ID_PATTERN: Final = re.compile(r"^[a-z0-9][a-z0-9-]{0,63}$")


class ProcessSupervisorError(RuntimeError):
    """A provider process graph could not be owned or stopped safely."""


@dataclass(frozen=True, slots=True)
class ProcessSpec:
    process_id: str
    argv: tuple[str, ...]
    start_order: int
    shutdown_order: int
    environment: tuple[tuple[str, str], ...] = ()
    startup_grace_seconds: float = 0.05
    ready_log_patterns: tuple[str, ...] = ()
    health_timeout_seconds: float = 0.0
    health_poll_interval_seconds: float = 0.05
    keep_stdin_open: bool = False
    interrupt_timeout_seconds: float = 2.0
    terminate_timeout_seconds: float = 1.0

    def __post_init__(self) -> None:
        if not PROCESS_ID_PATTERN.fullmatch(self.process_id):
            raise ValueError("process id must be a safe lowercase identifier")
        if not self.argv or any(not item or "\x00" in item for item in self.argv):
            raise ValueError("process argv must contain nonempty NUL-free values")
        if self.start_order < 0 or self.shutdown_order < 0:
            raise ValueError("process order must not be negative")
        if (
            self.startup_grace_seconds < 0
            or self.health_timeout_seconds < 0
            or self.health_poll_interval_seconds <= 0
            or self.interrupt_timeout_seconds < 0
            or self.terminate_timeout_seconds < 0
        ):
            raise ValueError("process timeouts must not be negative")
        if any(not pattern or "\x00" in pattern for pattern in self.ready_log_patterns):
            raise ValueError("ready log patterns must be nonempty and NUL-free")
        if len(self.ready_log_patterns) != len(set(self.ready_log_patterns)):
            raise ValueError("ready log patterns must be unique")
        if self.ready_log_patterns and self.health_timeout_seconds <= 0:
            raise ValueError("ready log patterns require a positive health timeout")
        keys = [key for key, _ in self.environment]
        if len(keys) != len(set(keys)) or any(
            not key or "=" in key or "\x00" in key or "\x00" in value
            for key, value in self.environment
        ):
            raise ValueError("process environment must contain unique safe entries")


@dataclass(frozen=True, slots=True)
class OwnedProcess:
    process_id: str
    pid: int
    pgid: int
    start_order: int
    shutdown_order: int
    stdout_path: Path
    stderr_path: Path

    def to_dict(self) -> dict[str, object]:
        return {
            "process_id": self.process_id,
            "pid": self.pid,
            "pgid": self.pgid,
            "start_order": self.start_order,
            "shutdown_order": self.shutdown_order,
            "stdout_path": self.stdout_path.name,
            "stderr_path": self.stderr_path.name,
        }


@dataclass(frozen=True, slots=True)
class ProcessStopResult:
    stopped_process_ids: tuple[str, ...]
    residue_process_ids: tuple[str, ...]


class PosixProcessSupervisor:
    """Own provider PGIDs and stop them in explicit reverse dependency order."""

    def __init__(
        self,
        runtime_root: Path,
        *,
        base_environment: dict[str, str] | None = None,
    ) -> None:
        if os.name != "posix":
            raise ProcessSupervisorError("S1 process supervisor requires a POSIX worker")
        self.runtime_root = runtime_root.expanduser().resolve()
        self.runtime_root.mkdir(mode=0o700, parents=True, exist_ok=True)
        self.base_environment = dict(base_environment or {})
        self._lock = threading.RLock()
        self._run_id: str | None = None
        self._processes: dict[str, tuple[ProcessSpec, subprocess.Popen[bytes]]] = {}

    def start(self, run_id: str, specs: tuple[ProcessSpec, ...]) -> tuple[OwnedProcess, ...]:
        if not PROCESS_ID_PATTERN.fullmatch(run_id):
            raise ProcessSupervisorError("run id must be a safe lowercase identifier")
        if not specs:
            raise ProcessSupervisorError("provider graph must not be empty")
        identifiers = [spec.process_id for spec in specs]
        if len(identifiers) != len(set(identifiers)):
            raise ProcessSupervisorError("provider process ids must be unique")
        start_orders = [spec.start_order for spec in specs]
        shutdown_orders = [spec.shutdown_order for spec in specs]
        if len(start_orders) != len(set(start_orders)):
            raise ProcessSupervisorError("provider start orders must be unique")
        if len(shutdown_orders) != len(set(shutdown_orders)):
            raise ProcessSupervisorError("provider shutdown orders must be unique")
        with self._lock:
            if self._processes:
                raise ProcessSupervisorError(
                    f"process authority is already owned by {self._run_id}"
                )
            run_root = self.runtime_root / run_id
            if run_root.exists():
                raise ProcessSupervisorError("run-scoped process directory already exists")
            run_root.mkdir(mode=0o700)
            self._run_id = run_id
            try:
                for spec in sorted(specs, key=lambda item: item.start_order):
                    self._start_one(run_root, spec)
                records = self.snapshot()
                self._persist_registry(run_root, records, terminal=False)
                return records
            except BaseException:
                self.stop()
                raise

    def snapshot(self) -> tuple[OwnedProcess, ...]:
        with self._lock:
            records: list[OwnedProcess] = []
            for process_id, (spec, process) in self._processes.items():
                return_code = process.poll()
                if return_code is not None:
                    raise ProcessSupervisorError(f"provider {process_id} exited with {return_code}")
                records.append(
                    OwnedProcess(
                        process_id=process_id,
                        pid=process.pid,
                        pgid=os.getpgid(process.pid),
                        start_order=spec.start_order,
                        shutdown_order=spec.shutdown_order,
                        stdout_path=self.runtime_root
                        / str(self._run_id)
                        / f"{process_id}.stdout.log",
                        stderr_path=self.runtime_root
                        / str(self._run_id)
                        / f"{process_id}.stderr.log",
                    )
                )
            return tuple(sorted(records, key=lambda item: item.start_order))

    def stop(self) -> ProcessStopResult:
        with self._lock:
            stopped: list[str] = []
            residue: list[str] = []
            remaining: dict[str, tuple[ProcessSpec, subprocess.Popen[bytes]]] = {}
            run_id = self._run_id
            for process_id, (spec, process) in sorted(
                self._processes.items(),
                key=lambda item: item[1][0].shutdown_order,
            ):
                pgid = process.pid
                if self._stop_group(process, pgid, spec):
                    stopped.append(process_id)
                else:
                    residue.append(process_id)
                    remaining[process_id] = (spec, process)
            if run_id is not None:
                run_root = self.runtime_root / run_id
                records = tuple(
                    OwnedProcess(
                        process_id=process_id,
                        pid=process.pid,
                        pgid=process.pid,
                        start_order=spec.start_order,
                        shutdown_order=spec.shutdown_order,
                        stdout_path=run_root / f"{process_id}.stdout.log",
                        stderr_path=run_root / f"{process_id}.stderr.log",
                    )
                    for process_id, (spec, process) in self._processes.items()
                )
                self._persist_registry(
                    run_root,
                    records,
                    terminal=True,
                    residue=tuple(residue),
                )
            self._processes = remaining
            if not remaining:
                self._run_id = None
            return ProcessStopResult(tuple(stopped), tuple(residue))

    def residue(self) -> tuple[str, ...]:
        with self._lock:
            return tuple(
                process_id
                for process_id, (_, process) in self._processes.items()
                if _group_exists(process.pid)
            )

    def _start_one(self, run_root: Path, spec: ProcessSpec) -> None:
        stdout_path = run_root / f"{spec.process_id}.stdout.log"
        stderr_path = run_root / f"{spec.process_id}.stderr.log"
        environment = dict(self.base_environment)
        environment.update(spec.environment)
        with stdout_path.open("xb") as stdout, stderr_path.open("xb") as stderr:
            process = subprocess.Popen(
                spec.argv,
                cwd=run_root,
                env=environment,
                stdin=subprocess.PIPE if spec.keep_stdin_open else subprocess.DEVNULL,
                stdout=stdout,
                stderr=stderr,
                start_new_session=True,
                close_fds=True,
            )
        self._processes[spec.process_id] = (spec, process)
        return_code = process.poll()
        if return_code is not None:
            phase = "before readiness" if spec.ready_log_patterns else "during startup"
            raise ProcessSupervisorError(
                f"provider {spec.process_id} exited {phase} with {return_code}"
            )
        try:
            process_group_id = os.getpgid(process.pid)
        except ProcessLookupError as exc:
            return_code = process.poll()
            phase = "before readiness" if spec.ready_log_patterns else "during startup"
            raise ProcessSupervisorError(
                f"provider {spec.process_id} exited {phase} with {return_code}"
            ) from exc
        if process_group_id != process.pid:
            raise ProcessSupervisorError(
                f"provider {spec.process_id} did not acquire a dedicated process group"
            )
        if spec.ready_log_patterns:
            self._wait_until_ready(
                process,
                spec,
                stdout_path=stdout_path,
                stderr_path=stderr_path,
            )
        elif spec.startup_grace_seconds:
            time.sleep(spec.startup_grace_seconds)
        return_code = process.poll()
        if return_code is not None:
            raise ProcessSupervisorError(
                f"provider {spec.process_id} exited during startup with {return_code}"
            )

    def _wait_until_ready(
        self,
        process: subprocess.Popen[bytes],
        spec: ProcessSpec,
        *,
        stdout_path: Path,
        stderr_path: Path,
    ) -> None:
        pending = {pattern: pattern.encode("utf-8") for pattern in spec.ready_log_patterns}
        maximum_tail_bytes = max(len(pattern) for pattern in pending.values()) - 1
        deadline = time.monotonic() + spec.health_timeout_seconds
        with (
            stdout_path.open("rb", buffering=0) as stdout,
            stderr_path.open("rb", buffering=0) as stderr,
        ):
            streams = (stdout, stderr)
            tails = [b"", b""]
            while pending:
                return_code = process.poll()
                if return_code is not None:
                    raise ProcessSupervisorError(
                        f"provider {spec.process_id} exited before readiness with {return_code}; "
                        f"missing markers: {', '.join(pending)}"
                    )
                for index, stream in enumerate(streams):
                    content = tails[index] + stream.read()
                    if content:
                        pending = {
                            pattern: encoded
                            for pattern, encoded in pending.items()
                            if encoded not in content
                        }
                        tails[index] = content[-maximum_tail_bytes:] if maximum_tail_bytes else b""
                    if not pending:
                        return
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    raise ProcessSupervisorError(
                        f"provider {spec.process_id} did not reach readiness within "
                        f"{spec.health_timeout_seconds:g} seconds; "
                        f"missing markers: {', '.join(pending)}"
                    )
                time.sleep(min(spec.health_poll_interval_seconds, remaining))

    def _stop_group(
        self,
        process: subprocess.Popen[bytes],
        pgid: int,
        spec: ProcessSpec,
    ) -> bool:
        group_signal_supported = True
        for sent_signal, timeout in (
            (signal.SIGINT, spec.interrupt_timeout_seconds),
            (signal.SIGTERM, spec.terminate_timeout_seconds),
            (signal.SIGKILL, 1.0),
        ):
            if process.poll() is not None and not _group_exists(pgid):
                process.wait(timeout=0)
                _close_stdin(process)
                return True
            try:
                if group_signal_supported:
                    os.killpg(pgid, sent_signal)
                else:
                    process.send_signal(sent_signal)
            except ProcessLookupError:
                _close_stdin(process)
                return True
            except PermissionError:
                # Some macOS application sandboxes deny killpg for an owned
                # child session. The reviewed WSL target retains group signals;
                # local development falls back to the group leader and does
                # not claim target residue acceptance from this path.
                group_signal_supported = False
                process.send_signal(sent_signal)
            deadline = time.monotonic() + timeout
            while time.monotonic() < deadline:
                if process.poll() is not None:
                    process.wait(timeout=0)
                    if not group_signal_supported or not _group_exists(pgid):
                        _close_stdin(process)
                        return True
                time.sleep(0.01)
        stopped = process.poll() is not None and (
            not group_signal_supported or not _group_exists(pgid)
        )
        if stopped:
            _close_stdin(process)
        return stopped

    def _persist_registry(
        self,
        run_root: Path,
        records: tuple[OwnedProcess, ...],
        *,
        terminal: bool,
        residue: tuple[str, ...] = (),
    ) -> None:
        write_json_atomic(
            run_root / "process-registry.json",
            {
                "schema_version": "missioncore.simulation-process-registry/v1",
                "run_id": run_root.name,
                "terminal": terminal,
                "processes": [record.to_dict() for record in records],
                "residue_process_ids": list(residue),
            },
        )


def _group_exists(pgid: int) -> bool:
    try:
        os.killpg(pgid, 0)
    except ProcessLookupError:
        return False
    except PermissionError:
        return True
    return True


def _close_stdin(process: subprocess.Popen[bytes]) -> None:
    if process.stdin is not None and not process.stdin.closed:
        process.stdin.close()
