import subprocess


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

    On a normal box `/run/archiso/bootmnt` does not exist and we return None,
    leaving the device list untouched. 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.
    """
    try:
        result = subprocess.run(
            ["findmnt", "-no", "SOURCE", "/run/archiso/bootmnt"],
            capture_output=True,
            text=True,
        )
    except FileNotFoundError:
        return None
    if result.returncode != 0:
        return None
    partition = result.stdout.strip()
    if not partition:
        return None
    try:
        parent = subprocess.run(
            ["lsblk", "-no", "PKNAME", partition],
            capture_output=True,
            text=True,
        )
    except FileNotFoundError:
        return None
    if parent.returncode != 0:
        return None
    name = parent.stdout.strip().splitlines()[0] if parent.stdout.strip() else ""
    return name or None


def _smart_status(device):
    try:
        result = subprocess.run(
            ["smartctl", "-H", device],
            capture_output=True,
        )
        output = result.stdout.decode()
        if "PASSED" in output:
            return "passed"
        elif "FAILED" in output:
            return "failed"
        return "unknown"
    except Exception as e:
        print(f"Error checking SMART status for {device}: {e}")
        return "unknown"


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


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


def get_drive_type_score(device):
    name = device.lower()
    if "nvme" in name:
        return 15
    if "ssd" in name:
        return 10
    return 5


def get_drive_type_label(device):
    name = device.lower()
    if "nvme" in name:
        return "NVMe"
    if "ssd" in name:
        return "SSD"
    return "HDD"


def parse_size_gb(size_str):
    size_str = size_str.strip().upper().replace(",", ".")
    if not size_str:
        return None
    if size_str.endswith("T"):
        return float(size_str[:-1]) * 1024
    if size_str.endswith("G"):
        return float(size_str[:-1])
    if size_str.endswith("M"):
        return float(size_str[:-1]) / 1024
    return None


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 score_device(device, size_gb):
    return get_drive_type_score(device) + get_drive_health(device) + get_size_score(size_gb)


def parse_lsblk_output(output, boot_disk=None):
    """Parse `lsblk -dn -o NAME,SIZE,TYPE` output into scored device dicts.

    Keeps only TYPE=disk so the live ISO's own squashfs (loop) and the boot
    CD-ROM (rom) don't show up as install targets. If `boot_disk` is given,
    that disk is also dropped — it's the USB stick the live ISO booted from
    on bare metal, where it appears as TYPE=disk and would otherwise be a
    valid-looking install target.
    """
    devices = []
    for line in output.strip().split("\n"):
        if not line:
            continue
        parts = line.split()
        if len(parts) < 3:
            continue
        name, size, dev_type = parts[0], parts[1], parts[2]
        if dev_type != "disk":
            continue
        if boot_disk and name == boot_disk:
            continue
        device = f"/dev/{name}"
        size_gb = parse_size_gb(size)
        status = _smart_status(device)
        score = get_drive_type_score(device) + _HEALTH_SCORE[status] + get_size_score(size_gb)
        devices.append(
            {
                "name": device,
                "size": size,
                "type_label": get_drive_type_label(device),
                "health_label": _HEALTH_LABEL[status],
                "score": score,
            }
        )
    devices.sort(key=lambda d: d["score"], reverse=True)
    return devices


def list_scored_devices():
    """Return [{name, size, score}, ...] for all physical disks, highest score first."""
    try:
        result = subprocess.run(
            ["lsblk", "-dn", "-o", "NAME,SIZE,TYPE"],
            capture_output=True,
            text=True,
            check=True,
        )
    except subprocess.CalledProcessError as e:
        print(f"Error listing devices: {e}")
        return []
    return parse_lsblk_output(result.stdout, boot_disk=_boot_disk_name())


def main():
    devices = list_scored_devices()
    if not devices:
        print("No storage devices found.")
        return
    print(f"\n{'Device':<20} {'Size':<10} {'Score'}")
    print("-" * 40)
    for d in devices:
        print(f"{d['name']:<20} {d['size']:<10} {d['score']}")
    print(f"\nBest drive for boot: {devices[0]['name']} (score {devices[0]['score']})")


if __name__ == "__main__":
    main()