#!/usr/bin/env python3

"""Read storage layout, usage, and drive health for Panama's Storage page.

Two boundaries, deliberately separate:

  snapshot   topology, usage, health, removable media. Cheap -- lsblk and a
             single udisks call -- so the page can open with it.
  scan       what is actually filling the drive. Expensive: measuring a folder
             means walking it, and this machine has a 1.2 TiB Steam library.
             The page asks for this on demand and remembers the answer.

  unmount PATH / eject PATH   removable media only, by explicit request.

Deliberately absent: partitioning and formatting. A settings pane is the wrong
place to hand someone a way to erase a disk in two clicks; GNOME Disks is one
button away on the page for that.
"""

from __future__ import annotations

import json
import os
from pathlib import Path
import re
import shutil
import subprocess
import sys
import time

# Bounded so a pathological tree cannot hang the page. A folder that cannot be
# measured in this long is reported as unmeasured rather than silently omitted,
# because a missing row reads as "this folder is small".
SCAN_TIMEOUT_SECONDS = 90

# Where "what is filling my drive" actually gets answered. Ordered so the most
# likely culprits are measured first; the walk stops when the budget runs out.
SCAN_TARGETS = [
    ("Steam library", "~/.local/share/Steam"),
    ("Documents", "~/Documents"),
    ("Downloads", "~/Downloads"),
    ("Videos", "~/Videos"),
    ("Pictures", "~/Pictures"),
    ("Music", "~/Music"),
    ("Flatpak applications", "~/.var/app"),
    ("Caches", "~/.cache"),
    ("Trash", "~/.local/share/Trash"),
]


class BoundaryError(RuntimeError):
    """A user-visible validation or command failure."""


def run(command: list[str], timeout: float = 15.0) -> str:
    try:
        completed = subprocess.run(command, check=False, capture_output=True,
                                   text=True, timeout=timeout)
    except (OSError, subprocess.TimeoutExpired) as error:
        raise BoundaryError(f"{command[0]} did not answer") from error
    if completed.returncode != 0:
        raise BoundaryError(completed.stderr.strip() or f"{command[0]} failed")
    return completed.stdout


def lsblk() -> list[dict]:
    """The block device tree.

    PANAMA_DISKS_LSBLK substitutes a recorded tree for the real one. It exists
    so the refusal rules below can be tested against a removable drive that is
    not plugged in -- the alternative is a test that runs unmount against
    whatever is actually mounted, which is not a test anyone should write.
    """
    fixture = os.environ.get("PANAMA_DISKS_LSBLK")
    if fixture:
        try:
            return json.loads(Path(fixture).read_text(encoding="utf-8"))["blockdevices"]
        except (OSError, json.JSONDecodeError, KeyError) as error:
            raise BoundaryError("could not read the recorded block device layout") from error
    fields = ("NAME,PATH,SIZE,TYPE,FSTYPE,MOUNTPOINTS,MODEL,SERIAL,RM,HOTPLUG,"
              "ROTA,FSSIZE,FSUSED,FSAVAIL")
    try:
        return json.loads(run(["lsblk", "-J", "-b", "-o", fields]))["blockdevices"]
    except (json.JSONDecodeError, KeyError) as error:
        raise BoundaryError("could not read the block device layout") from error


def udisks_objects() -> dict:
    """Every udisks object in one call, or nothing if udisks is not running.

    Health is a nice-to-have: a machine without udisks still gets its layout and
    usage, which is most of the page. Failing the whole snapshot because a
    temperature is unavailable would be the wrong trade.
    """
    if not shutil.which("busctl"):
        return {}
    try:
        raw = run(["busctl", "--system", "--json=short", "call",
                   "org.freedesktop.UDisks2", "/org/freedesktop/UDisks2",
                   "org.freedesktop.DBus.ObjectManager", "GetManagedObjects"])
        payload = json.loads(raw)
    except (BoundaryError, json.JSONDecodeError):
        return {}
    # busctl wraps the reply as {"type": "...", "data": [ {path: {iface: {prop: {"type","data"}}}} ]}
    data = payload.get("data") or []
    return data[0] if data else {}


def unwrap(value):
    """busctl encodes every value as {"type": t, "data": v}."""
    if isinstance(value, dict) and "data" in value and "type" in value:
        return value["data"]
    return value


def drive_health(objects: dict) -> dict[str, dict]:
    """Health keyed by drive serial, which is the only stable join to lsblk."""
    health: dict[str, dict] = {}
    for path, interfaces in objects.items():
        if "/drives/" not in path:
            continue
        drive = {key: unwrap(value)
                 for key, value in (interfaces.get("org.freedesktop.UDisks2.Drive") or {}).items()}
        serial = str(drive.get("Serial") or "")
        if not serial:
            continue

        entry = {
            "ejectable": bool(drive.get("Ejectable")),
            "removable": bool(drive.get("Removable")),
            "temperatureC": None,
            "powerOnHours": None,
            "healthy": None,
            "warnings": [],
            "selfTest": "",
        }

        nvme = {key: unwrap(value) for key, value in
                (interfaces.get("org.freedesktop.UDisks2.NVMe.Controller") or {}).items()}
        ata = {key: unwrap(value) for key, value in
               (interfaces.get("org.freedesktop.UDisks2.Drive.Ata") or {}).items()}

        if nvme:
            kelvin = nvme.get("SmartTemperature")
            if isinstance(kelvin, (int, float)) and kelvin > 0:
                entry["temperatureC"] = round(kelvin - 273.15, 1)
            hours = nvme.get("SmartPowerOnHours")
            if isinstance(hours, (int, float)):
                entry["powerOnHours"] = int(hours)
            warnings = nvme.get("SmartCriticalWarning") or []
            entry["warnings"] = [str(item) for item in warnings]
            entry["healthy"] = not entry["warnings"]
            entry["selfTest"] = str(nvme.get("SmartSelftestStatus") or "")
        elif ata:
            kelvin = ata.get("SmartTemperature")
            if isinstance(kelvin, (int, float)) and kelvin > 0:
                entry["temperatureC"] = round(kelvin - 273.15, 1)
            seconds = ata.get("SmartPowerOnSeconds")
            if isinstance(seconds, (int, float)) and seconds > 0:
                entry["powerOnHours"] = int(seconds // 3600)
            failing = ata.get("SmartFailing")
            if isinstance(failing, bool):
                entry["healthy"] = not failing
                entry["warnings"] = ["failing"] if failing else []
            entry["selfTest"] = str(ata.get("SmartSelftestStatus") or "")

        health[serial] = entry
    return health


def walk(node: dict, depth: int = 0):
    yield node, depth
    for child in node.get("children") or []:
        yield from walk(child, depth + 1)


def mountpoints(node: dict) -> list[str]:
    return [point for point in (node.get("mountpoints") or []) if point and point != "[SWAP]"]


def snapshot() -> dict:
    tree = lsblk()
    health = drive_health(udisks_objects())

    drives = []
    swap = []
    filesystems: dict[str, dict] = {}

    for root in tree:
        if root.get("type") != "disk":
            continue

        # zram is compressed swap in RAM. It is a block device and it is not
        # storage; listing it as a drive would be actively misleading about how
        # much of this machine is disk.
        if str(root.get("name", "")).startswith("zram"):
            swap.append({
                "name": root.get("name"),
                "sizeBytes": root.get("size") or 0,
                "kind": "zram",
            })
            continue

        serial = str(root.get("serial") or "")
        drive_health_entry = health.get(serial, {})
        partitions = []
        encrypted = False

        for node, depth in walk(root):
            if depth == 0:
                continue
            if node.get("fstype") == "crypto_LUKS":
                encrypted = True
            if node.get("type") in ("part", "crypt"):
                partitions.append({
                    "name": node.get("name"),
                    "path": node.get("path"),
                    "sizeBytes": node.get("size") or 0,
                    "fstype": node.get("fstype") or "",
                    "type": node.get("type"),
                    "mountpoints": mountpoints(node),
                })
            for point in mountpoints(node):
                if node.get("fssize") is None:
                    continue
                # Keyed by device: btrfs subvolumes mounted at / and /home are
                # ONE filesystem with one pool of free space. Reporting them as
                # two independent bars, which is what df does, doubles the free
                # space on screen.
                key = str(node.get("path"))
                entry = filesystems.setdefault(key, {
                    "device": node.get("path"),
                    "fstype": node.get("fstype") or "",
                    "sizeBytes": int(node.get("fssize") or 0),
                    "usedBytes": int(node.get("fsused") or 0),
                    "availBytes": int(node.get("fsavail") or 0),
                    "mountpoints": [],
                    "encrypted": node.get("type") == "crypt",
                })
                entry["mountpoints"].append(point)
                entry["mountpoints"].sort(key=lambda mount: (mount != "/", mount))

        drives.append({
            "name": root.get("name"),
            "path": root.get("path"),
            "model": (root.get("model") or "").strip() or str(root.get("name")),
            "serial": serial,
            "sizeBytes": root.get("size") or 0,
            "rotational": bool(root.get("rota")),
            "removable": bool(root.get("rm")) or bool(root.get("hotplug"))
                         or bool(drive_health_entry.get("removable")),
            "ejectable": bool(drive_health_entry.get("ejectable")),
            "encrypted": encrypted,
            "partitions": partitions,
            "temperatureC": drive_health_entry.get("temperatureC"),
            "powerOnHours": drive_health_entry.get("powerOnHours"),
            "healthy": drive_health_entry.get("healthy"),
            "warnings": drive_health_entry.get("warnings", []),
            "selfTest": drive_health_entry.get("selfTest", ""),
        })

    ordered = sorted(filesystems.values(),
                     key=lambda entry: (0 if "/" in entry["mountpoints"] else 1,
                                        entry["mountpoints"][0] if entry["mountpoints"] else ""))
    return {"drives": drives, "filesystems": ordered, "swap": swap}


def measure(path: Path, budget: float) -> tuple[int | None, float]:
    """Bytes used by a folder, and what is left of the time budget."""
    if not path.is_dir():
        return None, budget
    started = time.monotonic()
    try:
        completed = subprocess.run(["du", "-sxb", str(path)], check=False,
                                   capture_output=True, text=True, timeout=budget)
    except (OSError, subprocess.TimeoutExpired):
        # A folder that ran out of time has consumed the whole budget by
        # definition; the caller stops rather than starting another walk.
        return None, 0.0
    left = max(budget - (time.monotonic() - started), 0.0)
    if completed.returncode != 0:
        return None, left
    match = re.match(r"^(\d+)", completed.stdout)
    return (int(match.group(1)) if match else None), left


def container_reclaimable() -> dict | None:
    """Container images nothing is using. Absent when podman is not installed."""
    if not shutil.which("podman"):
        return None
    try:
        raw = run(["podman", "system", "df", "--format", "json"], timeout=20.0)
        entries = json.loads(raw)
    except (BoundaryError, json.JSONDecodeError):
        return None
    for entry in entries if isinstance(entries, list) else []:
        if not str(entry.get("Type", "")).lower().startswith("image"):
            continue
        # Raw* are integers; Size and Reclaimable are display strings like
        # "10.8GB" and "5.31GB (49%)". Reading the display strings is how this
        # first crashed, so only the raw fields are trusted, and an old podman
        # that lacks them reports nothing rather than a wrong number.
        total = entry.get("RawSize")
        reclaimable = entry.get("RawReclaimable")
        if not isinstance(total, int) or not isinstance(reclaimable, int):
            return None
        return {"totalBytes": total, "reclaimableBytes": reclaimable}
    return None


def scan() -> dict:
    folders = []
    remaining = float(SCAN_TIMEOUT_SECONDS)
    truncated = False
    for label, target in SCAN_TARGETS:
        path = Path(os.path.expanduser(target))
        if remaining <= 1.0:
            truncated = True
            break
        size, remaining = measure(path, remaining)
        if size is None:
            continue
        folders.append({"label": label, "path": str(path), "bytes": size})
    folders.sort(key=lambda item: item["bytes"], reverse=True)
    return {
        "folders": folders,
        "truncated": truncated,
        "containers": container_reclaimable(),
    }


def removable_device(path: str) -> dict:
    """Resolve a device path, refusing anything that is not removable.

    Unmounting the root filesystem is not a feature. The check is on the device
    rather than the button because the caller is a UI that can be wrong.
    """
    if not re.fullmatch(r"/dev/[A-Za-z0-9/_-]+", path or ""):
        raise BoundaryError("That is not a device path.")
    for root in lsblk():
        for node, _ in walk(root):
            if node.get("path") != path:
                continue
            top = root
            if not (bool(top.get("rm")) or bool(top.get("hotplug"))):
                raise BoundaryError("That drive is not removable.")
            return node
    raise BoundaryError("No such device.")


def main(arguments: list[str]) -> int:
    try:
        if arguments == ["snapshot"]:
            print(json.dumps(snapshot(), separators=(",", ":")))
        elif arguments == ["scan"]:
            print(json.dumps(scan(), separators=(",", ":")))
        elif len(arguments) == 2 and arguments[0] in ("unmount", "eject"):
            removable_device(arguments[1])
            action = "unmount" if arguments[0] == "unmount" else "power-off"
            flag = "-b" if arguments[0] == "unmount" else "-b"
            run(["udisksctl", action, flag, arguments[1]], timeout=30.0)
        else:
            raise BoundaryError(
                "Usage: panama-disks snapshot | scan | unmount DEVICE | eject DEVICE")
    except BoundaryError as error:
        print(str(error), file=sys.stderr)
        return 2
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))
