"""Firecracker process lifecycle and PID identity handling."""

from __future__ import annotations

import ctypes
import errno
import math
import os
import select
import signal
import stat
import subprocess
import time
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import Protocol

from ..config import Settings
from ..domain import VmRecord
from ..errors import FirecrackerError


class _PopenLike(Protocol):
    pid: int

    def poll(self) -> int | None: ...


@dataclass(frozen=True, slots=True)
class ProcessInfo:
    pid: int
    start_time: str
    boot_id: str | None = None


def _libc_pidfd(name: str, *arguments: int | None) -> int:
    # Some Linux Python builds omit pidfd bindings even when libc supports them.
    try:
        function = getattr(ctypes.CDLL(None, use_errno=True), name)
    except (AttributeError, OSError) as error:
        raise OSError(errno.ENOSYS, f"{name} is unavailable") from error
    function.argtypes = (
        [ctypes.c_int, ctypes.c_uint]
        if name == "pidfd_open"
        else [ctypes.c_int, ctypes.c_int, ctypes.c_void_p, ctypes.c_uint]
    )
    function.restype = ctypes.c_int
    result = function(*arguments)
    if result == -1:
        error_number = ctypes.get_errno()
        raise OSError(error_number, os.strerror(error_number))
    return result


def _open_pidfd(pid: int) -> int:
    native = getattr(os, "pidfd_open", None)
    if native is not None:
        return native(pid, 0)
    return _libc_pidfd("pidfd_open", pid, 0)


def _signal_pidfd(pidfd: int, signal_number: int) -> None:
    native = getattr(signal, "pidfd_send_signal", None)
    if native is not None:
        native(pidfd, signal_number)
    else:
        _libc_pidfd("pidfd_send_signal", pidfd, signal_number, None, 0)


def _pidfd_exited(pidfd: int) -> bool:
    poller = select.poll()
    poller.register(pidfd, select.POLLIN)
    events = poller.poll(0)
    if any(event & (select.POLLERR | select.POLLNVAL) for _, event in events):
        raise OSError(errno.EIO, "could not observe the Firecracker pidfd")
    return any(event & (select.POLLIN | select.POLLHUP) for _, event in events)


class FirecrackerProcessManager:
    """Start, observe, and terminate detached Firecracker processes."""

    def __init__(
        self,
        settings: Settings,
        *,
        popen: Callable[..., _PopenLike] = subprocess.Popen,
        sleep: Callable[[float], None] = time.sleep,
        monotonic: Callable[[], float] = time.monotonic,
    ) -> None:
        self._settings = settings
        self._popen = popen
        self._sleep = sleep
        self._monotonic = monotonic
        self._children: dict[int, _PopenLike] = {}

    def start(
        self,
        vm: VmRecord,
        *,
        on_spawn: Callable[[ProcessInfo], None] | None = None,
    ) -> ProcessInfo:
        """Publish identity before readiness; never replace an ambiguous process."""

        observation = self.observe(vm)
        if observation != "gone":
            raise FirecrackerError(
                f"refusing to start Firecracker: existing process is {observation}"
            )
        boot_id = self._boot_id()
        if boot_id is None:
            raise FirecrackerError("could not establish the current host boot identity")

        # Check support before creating a child that we might be unable to stop safely.
        try:
            pidfd = _open_pidfd(os.getpid())
            try:
                _signal_pidfd(pidfd, 0)
            finally:
                os.close(pidfd)
        except (OSError, NotImplementedError) as error:
            raise FirecrackerError(
                f"safe pidfd support is required to start Firecracker: {error}"
            ) from error

        socket_path = Path(vm.socket)
        self._remove_runtime_sockets(vm)

        log_path = Path(vm.log)
        process: _PopenLike | None = None
        try:
            log_path.parent.mkdir(parents=True, exist_ok=True)
            with log_path.open("ab", buffering=0) as log_file:
                previous_identity = (
                    vm.pid, vm.process_start_time, vm.boot_id, vm.operation, vm.spawn_pending,
                )
                vm.pid = None
                vm.process_start_time = None
                vm.boot_id = boot_id
                vm.operation = "spawning"
                vm.spawn_pending = True
                try:
                    process = self._popen(
                        [str(self._settings.firecracker_binary), "--api-sock", str(socket_path)],
                        stdout=log_file,
                        stderr=log_file,
                        start_new_session=True,
                    )
                except OSError as error:
                    # Popen cleans up an exec/fork failure, but an interrupt is uncertain.
                    if not isinstance(error, InterruptedError):
                        (
                            vm.pid, vm.process_start_time, vm.boot_id, vm.operation, vm.spawn_pending,
                        ) = previous_identity
                    raise
                vm.pid = process.pid
                self._children[process.pid] = process
                vm.process_start_time = self.process_start_time(process.pid)
                if vm.process_start_time is None:
                    raise FirecrackerError(
                        f"could not establish a safe process identity for Firecracker process {process.pid}"
                    )
                info = ProcessInfo(process.pid, vm.process_start_time, boot_id)
                if on_spawn is not None:
                    on_spawn(info)
                else:
                    vm.spawn_pending = False
            self._wait_for_socket(process, socket_path, log_path)
            return info
        except BaseException as error:
            if process is not None:
                try:
                    self.terminate(vm)
                except BaseException as cleanup_error:
                    raise FirecrackerError(
                        f"Firecracker startup failed ({type(error).__name__}: {error}); "
                        f"cleanup incomplete for PID {vm.pid}: {cleanup_error}"
                    ) from error
            if isinstance(error, OSError):
                raise FirecrackerError(f"could not start Firecracker: {error}") from error
            raise

    def adopt_legacy_identity(self, vm: VmRecord) -> bool:
        """Validate a v1 process; the caller must lock and persist the adopted boot ID."""

        pid, start_time = vm.pid, vm.process_start_time
        if (
            vm.boot_id is not None
            or type(pid) is not int
            or not 0 < pid < 2**31
            or not isinstance(start_time, str)
            or not start_time.isascii()
            or not start_time.isdigit()
        ):
            return False

        try:
            pidfd = _open_pidfd(pid)
            try:
                process_stat = self._process_stat(pid)
                if (
                    process_stat is None
                    or process_stat[1] != start_time
                    or process_stat[0] in {"Z", "X", "x"}
                ):
                    return False

                binary_path = self._settings.firecracker_binary
                socket_path = self._settings.vm_dir(vm.id) / "firecracker.sock"
                if (
                    not binary_path.is_absolute()
                    or ".." in binary_path.parts
                    or not vm.id
                    or vm.id in {".", ".."}
                    or Path(vm.id).name != vm.id
                    or vm.socket != str(socket_path)
                ):
                    return False

                # Sticky ancestors are safe only because every existing entry is root-owned.
                for directory in set(binary_path.parents) | set(socket_path.parents):
                    metadata = directory.lstat()
                    if (
                        not stat.S_ISDIR(metadata.st_mode)
                        or metadata.st_uid != 0
                        or (directory == socket_path.parent and metadata.st_mode & 0o077)
                        or (metadata.st_mode & 0o022 and not metadata.st_mode & stat.S_ISVTX)
                    ):
                        return False
                binary_stat = binary_path.lstat()
                socket_stat = socket_path.lstat()
                if (
                    not stat.S_ISREG(binary_stat.st_mode)
                    or binary_stat.st_uid != 0
                    or binary_stat.st_mode & 0o022
                    or not stat.S_ISSOCK(socket_stat.st_mode)
                    or socket_stat.st_uid != 0
                ):
                    return False

                proc_path = Path(f"/proc/{pid}")
                executable = proc_path / "exe"
                if executable.readlink() != binary_path or not executable.samefile(binary_path):
                    return False
                expected_cmdline = b"\0".join(
                    (os.fsencode(binary_path), b"--api-sock", os.fsencode(socket_path), b"")
                )
                if (proc_path / "cmdline").read_bytes() != expected_cmdline:
                    return False
                boot_id = self._boot_id()
                if boot_id is None or _pidfd_exited(pidfd):
                    return False
            finally:
                os.close(pidfd)
        except (OSError, ValueError, NotImplementedError):
            return False

        if (vm.pid, vm.process_start_time, vm.boot_id) != (pid, start_time, None):
            return False
        vm.boot_id = boot_id
        return True

    def observe(self, vm: VmRecord) -> str:
        """Distinguish proven exit from missing or unreadable process identity."""

        if vm.pid is not None and vm.pid <= 0:
            return "unknown"
        try:
            self._reap_child(vm.pid)
        except OSError:
            return "unknown"
        boot_id = self._boot_id()
        if vm.boot_id and boot_id and vm.boot_id != boot_id:
            return "gone"
        if vm.pid is None:
            if vm.spawn_pending or vm.operation == "spawning" or vm.process_start_time is not None:
                return "unknown"
            for socket_path in (Path(vm.socket), self._settings.guest_socket_path(vm.id)):
                try:
                    socket_path.lstat()
                except FileNotFoundError:
                    continue
                except OSError:
                    return "unknown"
                # An unowned socket may belong to a child whose identity was never saved.
                return "unknown"
            return "gone"

        stat = self._process_stat(vm.pid)
        if stat is None:
            return "unknown" if self._pid_exists(vm.pid) else "gone"
        state, start_time = stat
        has_start_time = (
            isinstance(vm.process_start_time, str)
            and vm.process_start_time.isascii()
            and vm.process_start_time.isdigit()
        )
        if has_start_time and vm.process_start_time != start_time:
            return "gone"
        if state in {"Z", "X", "x"}:
            # A zombie thread-group leader can still have live worker threads.
            try:
                pidfd = _open_pidfd(vm.pid)
                try:
                    return "gone" if _pidfd_exited(pidfd) else "unknown"
                finally:
                    os.close(pidfd)
            except ProcessLookupError:
                return "gone"
            except (OSError, NotImplementedError):
                return "unknown"
        if not has_start_time or not vm.boot_id or not boot_id:
            return "unknown"
        return "alive"

    def is_alive(self, vm: VmRecord) -> bool:
        """Conservative compatibility query: unknown processes may still be alive."""

        return self.observe(vm) != "gone"

    def terminate(self, vm: VmRecord, *, timeout: float | None = None) -> None:
        """Wait for exit before unlinking; timeout controls the SIGTERM grace period."""

        if vm.pid is not None and vm.pid <= 0:
            raise FirecrackerError(f"invalid persisted Firecracker PID: {vm.pid}")
        grace_timeout = self._settings.terminate_timeout_s if timeout is None else timeout
        if not math.isfinite(grace_timeout) or grace_timeout < 0:
            raise FirecrackerError("Firecracker termination timeout must be finite and non-negative")
        observation = self.observe(vm)
        if observation == "unknown":
            raise FirecrackerError(
                f"refusing to stop Firecracker PID {vm.pid}: process identity is ambiguous"
            )
        if observation == "gone":
            self._remove_runtime_sockets(vm)
            return

        try:
            pidfd = _open_pidfd(vm.pid)
        except ProcessLookupError:
            self._reap_child(vm.pid)
            self._remove_runtime_sockets(vm)
            return
        except (OSError, NotImplementedError) as error:
            raise FirecrackerError(
                f"could not safely open a pidfd for Firecracker process {vm.pid}: {error}"
            ) from error

        try:
            if not _pidfd_exited(pidfd):
                # Open the handle BEFORE rechecking identity; never signal by bare PID.
                observation = self.observe(vm)
                if observation == "unknown":
                    raise FirecrackerError(
                        f"refusing to stop Firecracker PID {vm.pid}: process identity became ambiguous"
                    )
                if observation == "alive":
                    try:
                        _signal_pidfd(pidfd, signal.SIGTERM)
                    except ProcessLookupError:
                        pass
                    if not self._wait_for_exit(pidfd, grace_timeout):
                        try:
                            _signal_pidfd(pidfd, signal.SIGKILL)
                        except ProcessLookupError:
                            pass
                        # Even an immediate forced stop needs time to confirm kernel exit.
                        if not self._wait_for_exit(pidfd, max(1.0, grace_timeout)):
                            raise FirecrackerError(
                                f"could not confirm exit of Firecracker process {vm.pid} after SIGKILL; "
                                "process identity and socket retained"
                            )
            self._reap_child(vm.pid)
            self._remove_runtime_sockets(vm)
        except (OSError, NotImplementedError) as error:
            raise FirecrackerError(
                f"could not safely stop Firecracker process {vm.pid}: {error}"
            ) from error
        finally:
            os.close(pidfd)

    def process_start_time(self, pid: int) -> str | None:
        """Read Linux proc start time so a reused PID is not mistaken for a VM."""

        stat = self._process_stat(pid)
        return stat[1] if stat is not None else None

    def host_boot_id(self) -> str | None:
        """Return the current host boot identity, or None when it is unreadable."""

        return self._boot_id()

    @staticmethod
    def _boot_id() -> str | None:
        try:
            return Path("/proc/sys/kernel/random/boot_id").read_text(encoding="ascii").strip() or None
        except (OSError, UnicodeError):
            return None

    @staticmethod
    def _process_stat(pid: int) -> tuple[str, str] | None:
        if pid <= 0:
            return None
        try:
            stat = Path(f"/proc/{pid}/stat").read_text(encoding="utf-8", errors="replace")
        except (OSError, UnicodeError):
            return None
        opening_parenthesis = stat.find("(")
        closing_parenthesis = stat.rfind(")")
        if (
            opening_parenthesis < 0
            or closing_parenthesis < opening_parenthesis
            or stat[:opening_parenthesis].strip() != str(pid)
        ):
            return None
        fields = stat[closing_parenthesis + 1 :].split()
        if (
            len(fields) < 20
            or fields[0] not in {"R", "S", "D", "Z", "T", "t", "X", "x", "K", "W", "P", "I"}
            or not fields[19].isascii()
            or not fields[19].isdigit()
        ):
            return None
        return fields[0], fields[19]

    def _wait_for_socket(
        self,
        process: _PopenLike,
        socket_path: Path,
        log_path: Path,
    ) -> None:
        deadline = self._monotonic() + self._settings.api_socket_timeout_s
        while True:
            if process.poll() is not None:
                raise FirecrackerError(f"Firecracker exited early; see {log_path}")
            if socket_path.is_socket():
                return
            remaining = deadline - self._monotonic()
            if remaining <= 0:
                raise FirecrackerError("timed out waiting for Firecracker API socket")
            self._sleep(min(0.05, remaining))

    def _wait_for_exit(self, pidfd: int, timeout: float) -> bool:
        deadline = self._monotonic() + timeout
        while True:
            if _pidfd_exited(pidfd):
                return True
            remaining = deadline - self._monotonic()
            if remaining <= 0:
                return False
            self._sleep(min(0.05, remaining))

    def _reap_child(self, pid: int | None) -> None:
        if pid is None:
            return
        process = self._children.get(pid)
        if process is not None and process.poll() is not None:
            self._children.pop(pid, None)

    def _remove_runtime_sockets(self, vm: VmRecord) -> None:
        self._remove_socket(self._settings.guest_socket_path(vm.id))
        self._remove_socket(Path(vm.socket))

    @staticmethod
    def _remove_socket(socket_path: Path) -> None:
        try:
            socket_path.unlink()
        except FileNotFoundError:
            pass
        except OSError as error:
            raise FirecrackerError(
                f"could not remove Firecracker socket {socket_path}: {error}"
            ) from error

    @staticmethod
    def _pid_exists(pid: int) -> bool:
        try:
            os.kill(pid, 0)
        except ProcessLookupError:
            return False
        except (OSError, OverflowError):
            return True
        return True