import json
import re
import subprocess
from pathlib import Path

# Smallest disk we offer as an install target. The base system plus Docker
# needs ~6 GB; the rest is headroom for images and app data.
MIN_DISK_GB = 16

# MBR (what archinstall lays out on BIOS boxes) can't address past 2 TiB.
MBR_LIMIT_GB = 2048

# Block devices lsblk reports as TYPE=disk that are never install targets:
# compressed RAM swap and the eMMC hardware boot / RPMB areas (4 MB, usually
# read-only) that show up next to mmcblk0 on every eMMC mini-PC and laptop.
_PSEUDO_DISK_RE = re.compile(r"^(zram\d+|mmcblk\d+boot\d+|mmcblk\d+rpmb)$")

_SUBPROCESS_TIMEOUT = 15


def _run(cmd, timeout=_SUBPROCESS_TIMEOUT):
    """subprocess.run with a timeout; None when the tool is missing or hangs."""
    try:
        return subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
    except (FileNotFoundError, subprocess.TimeoutExpired):
        return None


def _parent_disk(device):
    result = _run(["lsblk", "-no", "PKNAME", device])
    if result is None or result.returncode != 0:
        return None
    lines = [line.strip() for line in result.stdout.splitlines() if line.strip()]
    return lines[0] if lines else None


def _boot_device_from_cmdline(cmdline):
    """Resolve the boot partition from archiso's kernel args, or None."""
    args = dict(arg.split("=", 1) for arg in cmdline.split() if "=" in arg)
    if args.get("archisodevice", "").startswith("/dev/"):
        return args["archisodevice"]
    if args.get("archisosearchuuid"):
        cmd = ["blkid", "-U", args["archisosearchuuid"]]
    elif args.get("archisolabel"):
        cmd = ["blkid", "-L", args["archisolabel"]]
    else:
        return None
    result = _run(cmd)
    if result is None or result.returncode != 0:
        return None
    return result.stdout.strip() or None


def _boot_disk_name():
    """Return the parent disk name of the live-ISO boot media (e.g. "sdb"), or None.

    On bare metal booted from USB this is the stick we booted from — we want
    to filter it out so the user can't accidentally pick it as the install
    target. Primary source is the mount behind /run/archiso/bootmnt; when
    that's gone (copytoram unmounts it) fall back to the UUID/label archiso
    put on the kernel command line.
    """
    partition = None
    result = _run(["findmnt", "-no", "SOURCE", "/run/archiso/bootmnt"])
    if result is not None and result.returncode == 0:
        partition = result.stdout.strip() or None
    if partition is None:
        try:
            partition = _boot_device_from_cmdline(Path("/proc/cmdline").read_text())
        except OSError:
            partition = None
    if partition is None:
        return None
    # A stick written with dd has the ISO filesystem on the whole disk, where
    # PKNAME is empty — the device itself is then the disk.
    return _parent_disk(partition) or Path(partition).name


def _smart_status(device):
    result = _run(["smartctl", "-H", device])
    if result is None:
        return "unknown"
    if "PASSED" in result.stdout:
        return "passed"
    if "FAILED" in result.stdout:
        return "failed"
    return "unknown"


_HEALTH_SCORE = {"passed": 10, "failed": 0, "unknown": 5}
_HEALTH_LABEL = {
    "passed": "Healthy",
    "failed": "SMART warning",
    "unknown": "Status unknown",
}

_TYPE_SCORE = {"NVMe": 15, "SSD": 10, "eMMC": 8, "HDD": 5, "USB": 0}


def get_drive_health(device):
    return _HEALTH_SCORE[_smart_status(device)]


def classify_drive(name, tran="", rota=None, removable=False):
    """Human label for a disk: NVMe, SSD, eMMC, HDD or USB.

    Uses lsblk's transport and rotational flags — the kernel name alone
    can't tell an SSD from a spinning disk (both are sdX).
    """
    tran = (tran or "").lower()
    if tran == "usb" or (removable and not name.startswith("mmcblk")):
        return "USB"
    if tran == "nvme" or name.startswith("nvme"):
        return "NVMe"
    if tran == "mmc" or name.startswith("mmcblk"):
        return "eMMC"
    if rota is False:
        return "SSD"
    return "HDD"


def get_drive_type_score(type_label):
    return _TYPE_SCORE.get(type_label, 5)


def get_size_score(size_gb):
    if size_gb is None:
        return 5
    if size_gb < 128:
        return 5
    if size_gb < 512:
        return 7
    return 10


def format_size(size_bytes):
    gb = size_bytes / 1024**3
    if gb >= 1024:
        return f"{gb / 1024:.1f} TB".replace(".0 ", " ")
    if gb >= 10:
        return f"{gb:.0f} GB"
    return f"{gb:.1f} GB"


def _flag(value):
    # util-linux ≥2.37 emits JSON booleans; older builds emit "0"/"1".
    if isinstance(value, bool):
        return value
    return str(value).strip() in ("1", "true")


def parse_lsblk_json(output, boot_disk=None, uefi=True):
    """Parse `lsblk -J -d -b -o NAME,SIZE,TYPE,RO,RM,TRAN,ROTA` into device dicts.

    Keeps only real, writable, non-empty TYPE=disk devices (drops the live
    ISO's loop/rom, zram, eMMC boot areas, empty card-reader slots) and the
    USB stick we booted from. Disks too small to install on stay in the list
    with `usable=False` and a reason so the user sees why they're greyed out.
    Sorted best-first; unusable disks last.
    """
    try:
        entries = json.loads(output or "{}").get("blockdevices", [])
    except json.JSONDecodeError:
        return []
    devices = []
    for e in entries:
        name = e.get("name") or ""
        if e.get("type") != "disk" or not name:
            continue
        if _PSEUDO_DISK_RE.match(name) or (boot_disk and name == boot_disk):
            continue
        try:
            size_bytes = int(e.get("size") or 0)
        except (TypeError, ValueError):
            size_bytes = 0
        if size_bytes <= 0 or _flag(e.get("ro", False)):
            continue
        size_gb = size_bytes / 1024**3
        type_label = classify_drive(
            name,
            tran=e.get("tran") or "",
            rota=_flag(e.get("rota")) if e.get("rota") is not None else None,
            removable=_flag(e.get("rm", False)),
        )
        device = f"/dev/{name}"
        usable = size_gb >= MIN_DISK_GB
        status = _smart_status(device) if usable else "unknown"
        warning = ""
        if not uefi and size_gb > MBR_LIMIT_GB:
            warning = (
                "This computer starts in legacy (BIOS) mode, which can only use "
                "the first 2 TB of this disk. Switch the firmware to UEFI to use all of it."
            )
        elif type_label == "USB":
            warning = (
                "This is a USB drive — it works, but an internal disk is faster and more reliable."
            )
        devices.append(
            {
                "name": device,
                "size": format_size(size_bytes),
                "size_gb": size_gb,
                "type_label": type_label,
                "health_label": _HEALTH_LABEL[status],
                "score": get_drive_type_score(type_label)
                + _HEALTH_SCORE[status]
                + get_size_score(size_gb),
                "usable": usable,
                "reason": "" if usable else f"Too small — Furtka needs at least {MIN_DISK_GB} GB",
                "warning": warning,
            }
        )
    devices.sort(key=lambda d: (d["usable"], d["score"]), reverse=True)
    return devices


def list_scored_devices():
    """Return [{name, size, score, usable, ...}, ...] for all physical disks, best first."""
    result = _run(["lsblk", "-J", "-d", "-b", "-o", "NAME,SIZE,TYPE,RO,RM,TRAN,ROTA"])
    if result is None or result.returncode != 0:
        print(f"Error listing devices: {result.stderr if result else 'lsblk unavailable'}")
        return []
    return parse_lsblk_json(
        result.stdout,
        boot_disk=_boot_disk_name(),
        uefi=Path("/sys/firmware/efi").exists(),
    )


def main():
    devices = list_scored_devices()
    if not devices:
        print("No storage devices found.")
        return
    print(f"\n{'Device':<20} {'Size':<10} {'Type':<6} {'Score'}")
    print("-" * 46)
    for d in devices:
        note = "" if d["usable"] else f"  ({d['reason']})"
        print(f"{d['name']:<20} {d['size']:<10} {d['type_label']:<6} {d['score']}{note}")


if __name__ == "__main__":
    main()