"""Cleanup for the Firecracker backend.

Reaps *orphans* only — resources with no live VM behind them:

  * orphan run dirs: a per-bottle run dir (holding the ~1G rootfs.ext4)
    whose firecracker process has exited. These leak when a launch is
    hard-killed before its teardown runs (host OOM/crash, a cancelled CI
    job, `kill -9`); the clean-exit path already removes its own dir in
    launch.py.
  * orphan VM pids: a firecracker process whose run dir is already gone
    — a VMM left lingering after its dir was removed.

A run dir with a *live* firecracker process is a running bottle and is
left strictly alone: it is neither killed nor removed. (The backend's
`enumerate_active` registry is still a stub — #354 — so a live process
is the only reliable "this bottle is in use" signal we have. Once the
registry lands, registry-orphaned-but-running VMs can be reaped too.)

TAP slots free themselves (the flock drops when the launcher exits), so
there is nothing to reclaim there.
"""

from __future__ import annotations

import os
import shutil
import signal
import subprocess
from pathlib import Path

from ...log import info
from . import util
from .bottle_cleanup_plan import FirecrackerBottleCleanupPlan


def _run_root() -> Path:
    return util.cache_dir() / "run"


def _run_dir_of(cmd: str, run_root: Path) -> Path | None:
    """The bottle run dir a firecracker cmdline belongs to, or None.

    A bottle VM is launched with `--config-file <run_root>/<slug>/config.json`,
    so the run dir is the config file's parent when it sits directly under
    the run root. Anything else (a builder VM, the infra VM elsewhere) is
    not ours to reap here.
    """
    toks = cmd.split()
    for i, tok in enumerate(toks):
        if tok == "--config-file" and i + 1 < len(toks):
            parent = Path(toks[i + 1]).parent
            if parent.parent == run_root:
                return parent
    return None


def _scan_processes(run_root: Path) -> tuple[set[str], list[int]]:
    """Inspect running firecracker VMs under ``run_root``.

    Returns ``(live_run_dirs, orphan_pids)``:
      * ``live_run_dirs`` — run dirs backed by a running VM (never reaped);
      * ``orphan_pids`` — firecracker pids whose run dir no longer exists
        (a lingering VMM to kill).
    """
    result = subprocess.run(
        ["pgrep", "-a", "firecracker"],
        capture_output=True, text=True, check=False,
    )
    if result.returncode != 0:
        return set(), []
    live: set[str] = set()
    orphan_pids: list[int] = []
    for line in result.stdout.splitlines():
        parts = line.split(None, 1)
        if len(parts) != 2:
            continue
        try:
            pid = int(parts[0])
        except ValueError:
            continue
        run_dir = _run_dir_of(parts[1], run_root)
        if run_dir is None:
            continue
        if run_dir.is_dir():
            live.add(str(run_dir))
        else:
            orphan_pids.append(pid)
    return live, orphan_pids


def _orphan_run_dirs(run_root: Path, live: set[str]) -> list[str]:
    """Run dirs with no live VM behind them — the leaked ones to remove."""
    if not run_root.is_dir():
        return []
    return sorted(
        str(p) for p in run_root.iterdir()
        if p.is_dir() and str(p) not in live
    )


def prepare_cleanup() -> FirecrackerBottleCleanupPlan:
    run_root = _run_root()
    live, orphan_pids = _scan_processes(run_root)
    return FirecrackerBottleCleanupPlan(
        vm_pids=tuple(orphan_pids),
        run_dirs=tuple(_orphan_run_dirs(run_root, live)),
    )


def cleanup(plan: FirecrackerBottleCleanupPlan) -> None:
    for pid in plan.vm_pids:
        info(f"kill firecracker VM pid {pid}")
        try:
            os.kill(pid, signal.SIGTERM)
        except ProcessLookupError:
            pass
    for path in plan.run_dirs:
        info(f"rm -rf {path}")
        shutil.rmtree(path, ignore_errors=True)