import subprocess


def get_drive_health(device):
    try:
        result = subprocess.run(
            ["smartctl", "-H", device],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        output = result.stdout.decode()
        if "PASSED" in output:
            return 10
        elif "FAILED" in output:
            return 0
        return 5
    except Exception as e:
        print(f"Error checking SMART status for {device}: {e}")
        return 5


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


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 list_scored_devices():
    """Return [{name, size, score}, ...] for all physical disks, highest score first."""
    devices = []
    try:
        result = subprocess.run(
            ["lsblk", "-dn", "-o", "NAME,SIZE"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=True,
        )
        for line in result.stdout.strip().split("\n"):
            if not line:
                continue
            parts = line.split()
            if len(parts) < 2:
                continue
            name, size = parts[0], parts[1]
            device = f"/dev/{name}"
            devices.append(
                {
                    "name": device,
                    "size": size,
                    "score": score_device(device, parse_size_gb(size)),
                }
            )
    except subprocess.CalledProcessError as e:
        print(f"Error listing devices: {e}")

    devices.sort(key=lambda d: d["score"], reverse=True)
    return devices


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()