"""Buildbox disk admission at the 15 GiB red floor."""

from __future__ import annotations

import os
import stat
from dataclasses import dataclass
from pathlib import Path
from typing import Callable

RED_FLOOR_BYTES = 15 * 1024**3
FreeBytesFn = Callable[[str], int | None]


@dataclass(frozen=True)
class AdmissionVerdict:
    admit: bool
    free_bytes: int | None
    reason: str
    floor_bytes: int = RED_FLOOR_BYTES
    paths: tuple[str, ...] = ()
    failing_mount: str | None = None


def _at_least_red_floor(value: object) -> int:
    try:
        parsed = int(value)  # type: ignore[arg-type]
    except (TypeError, ValueError):
        return RED_FLOOR_BYTES
    return max(parsed, RED_FLOOR_BYTES)


def resolve_floor_bytes(env: os._Environ[str] | None = None) -> int:
    env = os.environ if env is None else env
    raw_bytes = env.get("BUILDBOX_DISK_MIN_FREE_BYTES")
    if raw_bytes is not None:
        try:
            return _at_least_red_floor(int(raw_bytes))
        except ValueError:
            pass
    raw_gib = env.get("BUILDBOX_DISK_MIN_FREE_GIB")
    if raw_gib is not None:
        try:
            return _at_least_red_floor(int(raw_gib) * 1024**3)
        except ValueError:
            pass
    return RED_FLOOR_BYTES


def resolve_probe_path(destination: str) -> str | None:
    path = Path(destination)
    if path.exists():
        return str(path.resolve())
    parent = path.parent
    while True:
        if parent.exists():
            return str(parent.resolve())
        if parent == parent.parent:
            return None
        parent = parent.parent


def unique_filesystem_probe_paths(destinations: list[str]) -> tuple[list[str], bool]:
    seen: dict[int, str] = {}
    for destination in destinations:
        probe = resolve_probe_path(destination)
        if probe is None:
            return [], False
        try:
            dev = os.stat(probe).st_dev
        except OSError:
            return [], False
        seen.setdefault(dev, probe)
    return list(seen.values()), True


def disk_free_bytes(target: str) -> int | None:
    try:
        st = os.statvfs(target)
        return int(st.f_bavail) * int(st.f_frsize)
    except OSError:
        return None


def admit_new_work(
  *,
  destinations: list[str],
  floor_bytes: int | None = None,
  free_bytes_fn: FreeBytesFn | None = None,
  env: os._Environ[str] | None = None,
) -> AdmissionVerdict:
    floor = resolve_floor_bytes(env) if floor_bytes is None else _at_least_red_floor(floor_bytes)
    probe = free_bytes_fn or disk_free_bytes
    paths, known = unique_filesystem_probe_paths(destinations)
    if not known:
        return AdmissionVerdict(
            admit=False,
            free_bytes=None,
            reason="unknown",
            floor_bytes=floor,
            paths=tuple(paths),
        )
    minimum: int | None = None
    for path in paths:
        free = probe(path)
        if free is None:
            return AdmissionVerdict(
                admit=False,
                free_bytes=None,
                reason="unknown",
                floor_bytes=floor,
                paths=tuple(paths),
                failing_mount=path,
            )
        minimum = free if minimum is None else min(minimum, free)
        if free < floor:
            return AdmissionVerdict(
                admit=False,
                free_bytes=free,
                reason="below-floor",
                floor_bytes=floor,
                paths=tuple(paths),
                failing_mount=path,
            )
    return AdmissionVerdict(
        admit=True,
        free_bytes=minimum,
        reason="ok",
        floor_bytes=floor,
        paths=tuple(paths),
    )


def format_capacity_refusal(host: str, verdict: AdmissionVerdict) -> str:
    required = verdict.floor_bytes
    if verdict.reason == "unknown":
        mount = verdict.failing_mount or (verdict.paths[0] if verdict.paths else "unknown")
        return (
            f"{host}: disk admission refused: mount={mount} free=unknown "
            f"required={required} bytes"
        )
    mount = verdict.failing_mount or (verdict.paths[0] if verdict.paths else "unknown")
    return (
        f"{host}: disk admission refused: mount={mount} free={verdict.free_bytes} "
        f"required={required} bytes"
    )


def remote_probe_script(destinations: list[str]) -> str:
  quoted = " ".join(f"'{d.replace(chr(39), chr(39) + chr(92) + chr(39) + chr(39))}'" for d in destinations)
  return (
      "set -euo pipefail\n"
      ". /usr/local/lib/buildbox/disk-admission.sh\n"
      f'buildbox_disk_admit "$(hostname -s 2>/dev/null || hostname)" {quoted}\n'
  )