#!/usr/bin/env python3
import sys
import os
from typing import NamedTuple
import subprocess


def printc(text, colour):
    """Print text in color."""
    colours = {
        "red": "\033[91m",
        "green": "\033[92m",
        "yellow": "\033[93m",
        "blue": "\033[94m",
        "purple": "\033[95m",
        "cyan": "\033[96m",
        "white": "\033[97m",
        "gray": "\033[90m",
        "reset": "\033[0m",
    }
    print(colours[colour] + text + colours["reset"])


DISK_PATHS = [
    "/dev/",
    "/dev/disk/by-id/",
    "/dev/disk/by-vdev/",
    "/dev/disk/by-path/",
    "/dev/disk/by-uuid/",
    "/dev/disk/by-label/",
]

PREFERED_PATH = "/dev/disk/by-vdev/"


def find_disk_path(disk_identifier):
    for path in DISK_PATHS:
        if os.path.exists(path + disk_identifier):
            return path + disk_identifier


def normalize_path(path):
    """Resolve the symbolic link of a path."""
    return os.path.realpath(path)


def find_other_links(normal_path):
    """Find other links to the same disk."""
    for path in DISK_PATHS:
        # loop through all the paths in the directory checking if they are links to the same disk
        for file in os.listdir(path):
            if os.path.islink(path + file):
                if os.path.realpath(path + file) == normal_path:
                    yield path + file


class DiskInfo(NamedTuple):
    name: str
    model: str
    capacity: str
    serial: str
    state: str
    wwn: str


def get_disk_info(disk_path):
    """Get disk information."""
    get_info = lambda x: subprocess.check_output(x, shell=True).decode().strip()
    return DiskInfo(
        name=get_info(f"lsblk -ndo NAME {disk_path}"),
        model=get_info(f"lsblk -ndo MODEL {disk_path}"),
        capacity=get_info(f"lsblk -ndo SIZE {disk_path}"),
        serial=get_info(f"lsblk -ndo SERIAL {disk_path}"),
        state=get_info(f"lsblk -ndo STATE {disk_path}"),
        wwn=get_info(f"lsblk -ndo WWN {disk_path}"),
    )


def choose_preferred_path(paths):
    for path in paths:
        if path.startswith(PREFERED_PATH):
            return path


arg = sys.argv[1]

# only the file part of the path is needed
if "/" in arg:
    arg = os.path.basename(arg)

disk_path = find_disk_path(arg)

if not disk_path:
    printc("Disk not found", "red")
    sys.exit(1)

normal_path = normalize_path(disk_path)
other_links = list(find_other_links(normal_path))
preferred_path = choose_preferred_path(other_links)
disk_info = get_disk_info(disk_path)

if preferred_path:
    printc(f"--- DISK {preferred_path.removeprefix(PREFERED_PATH)} ---", "green")
else:
    printc(f"DISK {disk_path}", "yellow")
print(f"Device:    {normal_path}")
print(f"Name:      {disk_info.name}")
print(f"Model:     {disk_info.model}")
print(f"Capacity:  {disk_info.capacity}")
print(f"Serial:    {disk_info.serial}")
print(f"State:     {disk_info.state}")
print(f"WWN:       {disk_info.wwn}")

for link in other_links:
    printc(f"Alias:     {link}", colour="gray")