import psutil
import subprocess

def get_drive_health(device):
    """
    Get the health of a storage device using smartctl.
    Returns a score based on the drive's SMART health status.
    """
    try:
        result = subprocess.run(['smartctl', '-H', device], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        output = result.stdout.decode()
        if "PASSED" in output:
            return 10  # Healthy drive
        elif "FAILED" in output:
            return 0  # Failed drive
        else:
            return 5  # Unknown or problematic drive
    except Exception as e:
        print(f"Error checking SMART status for {device}: {e}")
        return 5  # Default score for uncheckable devices

def get_drive_type(device):
    """
    Determine if a device is an SSD or HDD based on its device type.
    """
    if 'NVME' in device:
        return 15  # SSDs are optimal
    elif 'SSD' in device:
        return 10  # SSDs are optimal
    else:
        return 5  # HDDs are less optimal for boot drives

def get_drive_size(device):
    """
    Get size of a block device using lsblk (works for disks).
    Always returns an integer score.
    """
    try:
        result = subprocess.run(
            ["lsblk", "-dn", "-o", "SIZE", device],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=True
        )

        size_str = result.stdout.strip().upper()

        if not size_str:
            return 5  # fallback

        size_str = size_str.replace(",", ".")

        # Convert to GB
        if size_str.endswith("T"):
            size_gb = float(size_str[:-1]) * 1024
        elif size_str.endswith("G"):
            size_gb = float(size_str[:-1])
        elif size_str.endswith("M"):
            size_gb = float(size_str[:-1]) / 1024
        else:
            return 5  # unknown format

        if size_gb < 128:
            return 5
        elif size_gb < 512:
            return 7
        else:
            return 10

    except Exception as e:
        print(f"Error getting size for {device}: {e}")
        return 5  # ALWAYS return something

def get_device_score(device):
    """
    Calculate a suitability score for each drive based on:
    - Type (SSD/HDD)
    - Health (SMART status)
    - Size (GB)
    """
    score = 0

    # Type
    score += get_drive_type(device)
    
    # Health
    score += get_drive_health(device)

    # Size
    score += get_drive_size(device)

    return score


def list_storage_devices():
    """
    List all physical storage devices (e.g. /dev/sda, /dev/nvme0n1)
    and return them with their computed scores.

    Compatible with the existing scoring pipeline.
    """
    devices = []

    try:
        result = subprocess.run(
            ["lsblk", "-dn", "-o", "NAME"],  # -d = disks only, no partitions
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            check=True
        )

        for line in result.stdout.strip().split("\n"):
            if not line:
                continue

            device = f"/dev/{line.strip()}"
            score = get_device_score(device)  # <-- reuses your existing logic
            devices.append((device, score))

    except subprocess.CalledProcessError as e:
        print(f"Error listing devices: {e}")

    return devices

def main():
    devices = list_storage_devices()
    print(devices)
    
    if not devices:
        print("No storage devices found.")
        return

    print(f"\n{'Device':<20} {'Score'}")
    print("-" * 30)

    for device, score in devices:
        print(f"{device:<20} {score}")

    # Find the highest scoring drive
    best_device = max(devices, key=lambda x: x[1])
    print(f"\nBest drive for boot: {best_device[0]} with score: {best_device[1]}")

if __name__ == "__main__":
    main()