prep_disk.py

#!/usr/bin/env python3

import importlib.util
import json
import math
import os
import re
import shutil
import subprocess
import sys
import time
from pathlib import Path

# Label rendering defaults / fallbacks
DEFAULT_LABEL_SHORT_MM = 41
DEFAULT_LABEL_LONG_MM = 89
LABEL_DPI = 300
MM_PER_INCH = 25.4
PT_PER_INCH = 72.0

LABEL_MARGIN_X = 40
LABEL_MARGIN_Y = 40
LABEL_GAP = 25
TARGET_TEXT_WIDTH_MM = 70

FONT_CANDIDATES = [
    "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
    "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
    "/usr/share/fonts/truetype/liberation2/LiberationSans-Bold.ttf",
    "/usr/share/fonts/truetype/liberation2/LiberationSans-Regular.ttf",
    "/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf",
    "/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf",
]


def run(cmd, check=True, capture_output=True, text=True):
    return subprocess.run(
        cmd,
        check=check,
        capture_output=capture_output,
        text=text
    )


def require_root():
    if os.geteuid() != 0:
        print("This script must be run as root.")
        print("Run it with: sudo ./prep_disk.py")
        sys.exit(1)


def python_module_exists(module_name):
    return importlib.util.find_spec(module_name) is not None


def check_prerequisites():
    required_commands = {
        "lsblk": "util-linux",
        "findmnt": "util-linux",
        "smartctl": "smartmontools",
        "parted": "parted",
        "partprobe": "parted",
        "mkfs.exfat": "exfatprogs",
        "wipefs": "util-linux",
        "dd": "coreutils",
        "udevadm": "udev",
        "blockdev": "util-linux",
        "umount": "util-linux",
        "mount": "util-linux",
        "cp": "coreutils",
        "lp": "cups-client",
        "lpstat": "cups-client",
        "lpoptions": "cups-client",
    }

    required_python_modules = {
        "PIL": "python3-pil",
    }

    missing_packages = set()

    print("\nChecking required tools...\n")

    for cmd, pkg in required_commands.items():
        if shutil.which(cmd) is None:
            print(f"Missing: {cmd} (package: {pkg})")
            missing_packages.add(pkg)

    for module_name, pkg in required_python_modules.items():
        if not python_module_exists(module_name):
            print(f"Missing Python module: {module_name} (package: {pkg})")
            missing_packages.add(pkg)

    if not missing_packages:
        print("All required tools are installed.\n")
        return

    print("\nThe following packages need to be installed:")
    for pkg in sorted(missing_packages):
        print(f"  {pkg}")

    choice = input("\nInstall them now? [Y/n]: ").strip().lower()

    if choice in ("", "y", "yes"):
        print("\nInstalling required packages...\n")

        update_result = subprocess.run(["apt", "update"])
        if update_result.returncode != 0:
            print("apt update failed. Please install the missing packages manually.")
            sys.exit(1)

        install_cmd = ["apt", "install", "-y"] + sorted(missing_packages)
        install_result = subprocess.run(install_cmd)

        if install_result.returncode != 0:
            print("Package installation failed. Please install manually and run again.")
            sys.exit(1)

        print("\nDependencies installed successfully.\n")
    else:
        print("\nPlease install the required packages and run again.")
        print(f"Example: sudo apt install {' '.join(sorted(missing_packages))}")
        sys.exit(1)


def sanitize_filename(value):
    value = (value or "").strip()
    if not value:
        value = "Unknown"

    value = re.sub(r'[\\/*?:"<>|]', '-', value)
    value = re.sub(r"\s+", " ", value).strip(" .")
    return value if value else "Unknown"


def sanitize_folder_name(value):
    value = (value or "").strip()
    if not value:
        value = "UnknownSerial"
    value = re.sub(r'[\\/*?:"<>|]', '-', value)
    value = re.sub(r"\s+", " ", value).strip(" .")
    return value if value else "UnknownSerial"


def sanitize_exfat_label(value):
    value = (value or "").strip()
    if not value:
        value = "NO_SERIAL"

    value = re.sub(r"[^A-Za-z0-9 _-]", "", value)
    value = re.sub(r"\s+", " ", value).strip()

    if not value:
        value = "NO_SERIAL"

    return value[:15]


def sanitize_label_text(value, fallback):
    value = (value or "").strip()
    value = re.sub(r"[\r\n\t]+", " ", value)
    value = re.sub(r"\s+", " ", value).strip()
    return value if value else fallback


def human_size(num_bytes):
    size = float(num_bytes)
    units = ["B", "KB", "MB", "GB", "TB", "PB"]

    for unit in units:
        if size < 1024 or unit == units[-1]:
            return f"{size:.2f} {unit}"
        size /= 1024


def get_root_parent_disk():
    try:
        root_source = run(["findmnt", "-n", "-o", "SOURCE", "/"]).stdout.strip()

        if not root_source.startswith("/dev/"):
            return None

        base = os.path.basename(root_source)

        nvme_match = re.match(r"^(nvme\d+n\d+)p\d+$", base)
        if nvme_match:
            return f"/dev/{nvme_match.group(1)}"

        mmc_match = re.match(r"^(mmcblk\d+)p\d+$", base)
        if mmc_match:
            return f"/dev/{mmc_match.group(1)}"

        sd_match = re.match(r"^([a-zA-Z]+)\d+$", base)
        if sd_match and not base.startswith(("nvme", "mmcblk", "loop")):
            parent = re.sub(r"\d+$", "", base)
            return f"/dev/{parent}"

        return None
    except Exception:
        return None


def get_root_pkname():
    try:
        root_source = run(["findmnt", "-n", "-o", "SOURCE", "/"]).stdout.strip()
        if not root_source.startswith("/dev/"):
            return None

        result = run(["lsblk", "-no", "PKNAME", root_source], check=False)
        pkname = result.stdout.strip()
        return pkname if pkname else None
    except Exception:
        return None


def get_disk_size_bytes(dev_name):
    sys_size = Path(f"/sys/class/block/{dev_name}/size")
    if sys_size.exists():
        try:
            sectors = int(sys_size.read_text().strip())
            return sectors * 512
        except Exception:
            pass

    try:
        result = run(["blockdev", "--getsize64", f"/dev/{dev_name}"], check=False)
        if result.returncode == 0 and result.stdout.strip().isdigit():
            return int(result.stdout.strip())
    except Exception:
        pass

    return 0


def get_disks():
    result = run([
        "lsblk",
        "-J",
        "-d",
        "-b",
        "-o",
        "NAME,SIZE,MODEL,SERIAL,VENDOR,TYPE,TRAN,RM,RO,HOTPLUG"
    ])

    data = json.loads(result.stdout)
    disks = []

    root_disk_by_path = get_root_parent_disk()
    root_pkname = get_root_pkname()

    for dev in data.get("blockdevices", []):
        if dev.get("type") != "disk":
            continue

        path = f"/dev/{dev['name']}"
        is_root_disk = (path == root_disk_by_path) or (dev["name"] == root_pkname)

        size_raw = dev.get("size", 0)
        try:
            size_bytes = int(size_raw)
        except Exception:
            size_bytes = 0

        if size_bytes <= 0:
            size_bytes = get_disk_size_bytes(dev["name"])

        disks.append({
            "path": path,
            "name": dev["name"],
            "size_bytes": size_bytes,
            "size_human": human_size(size_bytes) if size_bytes > 0 else "Unknown",
            "model": (dev.get("model") or "").strip(),
            "serial": (dev.get("serial") or "").strip(),
            "vendor": (dev.get("vendor") or "").strip(),
            "tran": (dev.get("tran") or "").strip(),
            "rm": str(dev.get("rm", "")),
            "ro": str(dev.get("ro", "")),
            "hotplug": str(dev.get("hotplug", "")),
            "is_root_disk": is_root_disk
        })

    return disks


def get_partition_path(disk_path):
    base = os.path.basename(disk_path)
    if re.search(r"^(nvme\d+n\d+|mmcblk\d+)$", base):
        return disk_path + "p1"
    return disk_path + "1"


#def collect_smart_report(disk_path):
#    result = run(["smartctl", "-x", disk_path], check=False)
#    output = result.stdout or ""
#    if result.stderr:
#        output += "\n" + result.stderr
#    return output

def score_smart_output(text):
    score = 0
    if not text:
        return score

    patterns = [
        r"SMART overall-health self-assessment test result",
        r"ID#\s+ATTRIBUTE_NAME",
        r"Power_On_Hours",
        r"Reallocated_Sector_Ct",
        r"Current_Pending_Sector",
        r"Offline_Uncorrectable",
        r"Power_Cycle_Count",
        r"Temperature_Celsius",
        r"Current Drive Temperature",
        r"Serial Number:",
        r"Model Family:",
        r"Device Model:",
        r"User Capacity:",
        r"Vendor Specific SMART Attributes with Thresholds:",
    ]

    for pattern in patterns:
        if re.search(pattern, text, re.IGNORECASE | re.MULTILINE):
            score += 10

    # Prefer outputs that are longer, up to a point
    score += min(len(text) // 200, 50)

    # Penalize obvious failures a bit
    failure_patterns = [
        r"Unknown USB bridge",
        r"Please specify device type",
        r"SMART support is:\s+Unavailable",
        r"Read Device Identity failed",
        r"A mandatory SMART command failed",
    ]

    for pattern in failure_patterns:
        if re.search(pattern, text, re.IGNORECASE):
            score -= 15

    return score


def collect_smart_report(disk_path):
    probe_cmds = [
        ["smartctl", "-x", disk_path],
        ["smartctl", "-x", "-d", "sat", disk_path],
        ["smartctl", "-x", "-d", "sat,12", disk_path],
        ["smartctl", "-x", "-d", "sat,16", disk_path],
        ["smartctl", "-x", "-d", "scsi", disk_path],
        ["smartctl", "-x", "-d", "usbjmicron", disk_path],
        ["smartctl", "-x", "-d", "usbprolific", disk_path],
        ["smartctl", "-x", "-d", "usbsunplus", disk_path],
    ]

    best_output = ""
    best_score = -9999
    best_cmd = None

    for cmd in probe_cmds:
        result = run(cmd, check=False)
        output = result.stdout or ""
        if result.stderr:
            output += "\n" + result.stderr

        score = score_smart_output(output)

        if score > best_score:
            best_score = score
            best_output = output
            best_cmd = cmd

    header = []
    header.append("SMART COLLECTION METHOD")
    header.append("=" * 72)
    header.append(f"Command used: {' '.join(best_cmd) if best_cmd else 'Unknown'}")
    header.append(f"Score: {best_score}")
    header.append("")

    return "\n".join(header) + best_output


def get_identity_for_filename(disk):
    vendor = sanitize_filename(disk["vendor"] or "UnknownMake")
    model = sanitize_filename(disk["model"] or "UnknownModel")
    serial = sanitize_filename(disk["serial"] or "UnknownSerial")
    return vendor, model, serial


def get_report_dir(script_dir, disk):
    serial = sanitize_folder_name(disk.get("serial") or "UnknownSerial")
    report_dir = Path(script_dir) / serial
    report_dir.mkdir(parents=True, exist_ok=True)
    return report_dir


def extract_first_match(patterns, text, flags=re.MULTILINE):
    for pattern in patterns:
        match = re.search(pattern, text, flags | re.IGNORECASE)
        if match:
            return match.group(1).strip()
    return "Unknown"


def extract_temperature_celsius(smart_text):
    patterns = [
        r"Current Drive Temperature:\s+(\d+)\s*C",
        r"Temperature:\s+(\d+)\s+Celsius",
        r"Temperature Sensor \d+:\s+(\d+)\s+Celsius",
        r"^\s*194\s+Temperature_Celsius\s+\S+\s+\S+\s+\S+\s+\S+\s+\S+\s+\S+\s+(\d+)\s*$",
        r"^\s*190\s+Airflow_Temperature_Cel\s+\S+\s+\S+\s+\S+\s+\S+\s+\S+\s+\S+\s+(\d+)\s*$",
        r"^\s*190\s+Temperature_Internal\s+\S+\s+\S+\s+\S+\s+\S+\s+\S+\s+\S+\s+(\d+)\s*$",
    ]
    value = extract_first_match(patterns, smart_text)
    if value.isdigit():
        return int(value)
    return None


def extract_total_host_writes(smart_text):
    matches = [
        re.search(r"Data Units Written:\s*([\d,]+)", smart_text),
        re.search(r"Host Writes:\s*([\d,]+)", smart_text),
        re.search(r"Total Host Writes:\s*([\d,]+)\s*GB", smart_text),
    ]

    for m in matches:
        if m:
            return m.group(1).replace(",", "")

    return None


def parse_smart_summary(smart_text, disk):
    summary = {
        "Make": disk.get("vendor") or "Unknown",
        "Model": disk.get("model") or "Unknown",
        "Serial Number": disk.get("serial") or "Unknown",
        "Capacity": disk.get("size_human") or "Unknown",
        "Interface": disk.get("tran") or "Unknown",
        "SMART Overall Health": "Unknown",
        "Power-On Hours": "Unknown",
        "Power Cycle Count": "Unknown",
        "Reallocated Sectors": "Unknown",
        "Current Pending Sectors": "Unknown",
        "Offline Uncorrectable": "Unknown",
        "Temperature": "Unknown",
        "Total Host Writes": "Unknown",
    }

    health = extract_first_match([
        r"SMART overall-health self-assessment test result:\s*(.+)",
        r"SMART Health Status:\s*(.+)",
        r"SMART overall-health self-assessment test result\s*:\s*(.+)",
        r"SMART overall-health self-assessment test result =\s*(.+)",
    ], smart_text)

    poh = extract_first_match([
        r"^\s*9\s+Power_On_Hours\s+\S+\s+\S+\s+\S+\s+\S+\s+\S+\s+\S+\s+(\d+)\s*$",
        r"^\s*9\s+Power_On_Hours\s+\S+.*?(\d+)\s*$",
        r"Power on hours:\s*(.+)",
        r"Accumulated power on time, hours:\s*(.+)"
    ], smart_text)

    pcc = extract_first_match([
        r"^\s*12\s+Power_Cycle_Count\s+\S+\s+\S+\s+\S+\s+\S+\s+\S+\s+\S+\s+(\d+)\s*$",
        r"^\s*12\s+Power_Cycle_Count.*?(\d+)\s*$",
        r"Power cycle count:\s*(.+)",
        r"start[- ]stop count:\s*(.+)",
    ], smart_text)

    realloc = extract_first_match([
        r"^\s*5\s+Reallocated_Sector_Ct\s+\S+\s+\S+\s+\S+\s+\S+\s+\S+\s+\S+\s+(\d+)\s*$",
        r"^\s*5\s+Reallocated_Sector_Ct.*?(\d+)\s*$",
        r"Reallocated sector count:\s*(.+)",
        r"Elements in grown defect list:\s*(.+)",
    ], smart_text)

    pending = extract_first_match([
        r"^\s*197\s+Current_Pending_Sector\s+\S+\s+\S+\s+\S+\s+\S+\s+\S+\s+\S+\s+(\d+)\s*$",
        r"^\s*197\s+Current_Pending_Sector.*?(\d+)\s*$",
        r"Current pending sector count:\s*(.+)",
    ], smart_text)

    offline_unc = extract_first_match([
        r"^\s*198\s+Offline_Uncorrectable\s+\S+\s+\S+\s+\S+\s+\S+\s+\S+\s+\S+\s+(\d+)\s*$",
        r"^\s*198\s+Offline_Uncorrectable.*?(\d+)\s*$",
        r"Offline uncorrectable sector count:\s*(.+)",
    ], smart_text)

    temp_c = extract_temperature_celsius(smart_text)
    host_writes = extract_total_host_writes(smart_text)

    summary["SMART Overall Health"] = health
    summary["Power-On Hours"] = poh
    summary["Power Cycle Count"] = pcc
    summary["Reallocated Sectors"] = realloc
    summary["Current Pending Sectors"] = pending
    summary["Offline Uncorrectable"] = offline_unc
    summary["Temperature"] = f"{temp_c} C" if temp_c is not None else "Unknown"
    summary["Total Host Writes"] = host_writes if host_writes else "Unknown"

    return summary


def parse_ata_smart_attributes(smart_text):
    attributes = []
    lines = smart_text.splitlines()

    table_started = False

    for line in lines:
        stripped = line.rstrip()

        if re.search(r"ID#\s+ATTRIBUTE_NAME\s+FLAG\s+VALUE\s+WORST\s+THRESH\s+TYPE\s+UPDATED\s+WHEN_FAILED\s+RAW_VALUE", stripped):
            table_started = True
            continue

        if table_started:
            if not stripped.strip():
                break

            m = re.match(
                r"^\s*(\d+)\s+([A-Za-z0-9_\-]+)\s+([0-9A-Fa-fx]+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(.+?)\s*$",
                stripped
            )
            if m:
                attributes.append({
                    "id": m.group(1),
                    "name": m.group(2).replace("_", " "),
                    "current": m.group(4),
                    "worst": m.group(5),
                    "threshold": m.group(6),
                    "raw": m.group(10).strip(),
                })
            else:
                # If we hit something clearly no longer table-ish, bail out.
                if not re.match(r"^\s*\d+\s+", stripped):
                    break

    return attributes


def format_summary_block(summary):
    lines = []
    lines.append("LISTING SUMMARY")
    lines.append("=" * 72)

    ordered_keys = [
        "Make",
        "Model",
        "Serial Number",
        "Capacity",
        "Interface",
        "SMART Overall Health",
        "Temperature",
        "Power-On Hours",
        "Power Cycle Count",
        "Total Host Writes",
        "Reallocated Sectors",
        "Current Pending Sectors",
        "Offline Uncorrectable",
    ]

    for key in ordered_keys:
        lines.append(f"{key}: {summary.get(key, 'Unknown')}")

    lines.append("")
    lines.append("Suggested listing note:")
    lines.append(
        f"{summary['Model']} hard drive, {summary['Capacity']}, "
        f"tested and formatted exFAT. "
        f"SMART health: {summary['SMART Overall Health']}. "
        f"Temperature: {summary['Temperature']}. "
        f"Power-on hours: {summary['Power-On Hours']}. "
        f"Reallocated sectors: {summary['Reallocated Sectors']}. "
        f"Pending sectors: {summary['Current Pending Sectors']}."
    )
    return "\n".join(lines)


def write_report(script_dir, disk, report_text, summary_text):
    vendor, model, serial = get_identity_for_filename(disk)
    report_dir = get_report_dir(script_dir, disk)
    filename = f"{vendor} - {model} - {serial}.txt"
    path = report_dir / filename

    with open(path, "w", encoding="utf-8") as f:
        f.write("Drive Report\n")
        f.write("=" * 72 + "\n")
        f.write(f"Generated: {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"Disk Path: {disk['path']}\n")
        f.write(f"Vendor: {disk['vendor'] or 'Unknown'}\n")
        f.write(f"Model: {disk['model'] or 'Unknown'}\n")
        f.write(f"Serial: {disk['serial'] or 'Unknown'}\n")
        f.write(f"Size: {disk['size_human']}\n")
        f.write(f"Transport: {disk['tran'] or 'Unknown'}\n")
        f.write("\nSMART / DRIVE DATA\n")
        f.write("=" * 72 + "\n")
        f.write(report_text.rstrip() + "\n\n")
        f.write(summary_text.rstrip() + "\n")

    return path


def draw_vertical_gradient(img, top_rgb, bottom_rgb):
    from PIL import ImageDraw
    draw = ImageDraw.Draw(img)
    width, height = img.size
    for y in range(height):
        r = int(top_rgb[0] + (bottom_rgb[0] - top_rgb[0]) * (y / max(1, height - 1)))
        g = int(top_rgb[1] + (bottom_rgb[1] - top_rgb[1]) * (y / max(1, height - 1)))
        b = int(top_rgb[2] + (bottom_rgb[2] - top_rgb[2]) * (y / max(1, height - 1)))
        draw.line((0, y, width, y), fill=(r, g, b))


def create_smart_graphic_png(script_dir, disk, summary, smart_text, tested_timestamp):
    from PIL import Image, ImageDraw, ImageFont

    vendor, model, serial = get_identity_for_filename(disk)
    report_dir = get_report_dir(script_dir, disk)
    filename = f"{vendor} - {model} - {serial}.png"
    path = report_dir / filename

    width = 1700
    margin = 30
    panel_gap = 18
    left_col_w = 260
    top_panel_h = 300

    attrs = parse_ata_smart_attributes(smart_text)

    font_path_bold = get_font_path()
    font_path_regular = None
    for p in FONT_CANDIDATES:
        if Path(p).exists() and "Bold" not in Path(p).name:
            font_path_regular = p
            break
    if font_path_regular is None:
        font_path_regular = font_path_bold

    title_font = ImageFont.truetype(font_path_bold, 40)
    small_title_font = ImageFont.truetype(font_path_bold, 24)
    big_box_font = ImageFont.truetype(font_path_bold, 38)
    medium_font = ImageFont.truetype(font_path_regular, 28)
    medium_bold_font = ImageFont.truetype(font_path_bold, 28)
    small_font = ImageFont.truetype(font_path_regular, 24)
    table_font = ImageFont.truetype(font_path_regular, 24)
    table_bold_font = ImageFont.truetype(font_path_bold, 24)

    table_row_h = 34
    table_header_h = 42
    table_h = table_header_h + (len(attrs) * table_row_h if attrs else 2 * table_row_h)
    footer_h = 55
    height = margin + top_panel_h + panel_gap + table_h + panel_gap + footer_h + margin

    img = Image.new("RGB", (width, height), (236, 242, 250))
    draw = ImageDraw.Draw(img)

    draw_vertical_gradient(img, (239, 245, 252), (219, 231, 244))

    def rounded_box(x1, y1, x2, y2, fill, outline, radius=22, width_px=2):
        draw.rounded_rectangle((x1, y1, x2, y2), radius=radius, fill=fill, outline=outline, width=width_px)

    def gradient_badge(x1, y1, x2, y2, top_rgb, bottom_rgb, outline=(112, 143, 205), radius=26):
        badge_w = x2 - x1
        badge_h = y2 - y1
        badge = Image.new("RGB", (badge_w, badge_h), (255, 255, 255))
        draw_vertical_gradient(badge, top_rgb, bottom_rgb)
        mask = Image.new("L", (badge_w, badge_h), 0)
        from PIL import ImageDraw as ID
        md = ID.Draw(mask)
        md.rounded_rectangle((0, 0, badge_w - 1, badge_h - 1), radius=radius, fill=255)
        img.paste(badge, (x1, y1), mask)
        draw.rounded_rectangle((x1, y1, x2, y2), radius=radius, outline=outline, width=2)

    header_x1 = margin
    header_y1 = margin
    header_x2 = width - margin
    header_y2 = margin + 95
    rounded_box(header_x1, header_y1, header_x2, header_y2, (255, 255, 255), (170, 186, 210), 18, 2)

    model_header = disk.get("model") or "Unknown Model"
    size_header = disk.get("size_human") or "Unknown Size"
    title_text = f"{model_header} | {size_header} | Tested: {tested_timestamp}"
    draw.text((header_x1 + 25, header_y1 + 20), title_text, font=title_font, fill=(20, 30, 50))

    left_x1 = margin
    left_y1 = header_y2 + panel_gap
    left_x2 = left_x1 + left_col_w
    left_y2 = left_y1 + top_panel_h
    rounded_box(left_x1, left_y1, left_x2, left_y2, (255, 255, 255), (170, 186, 210), 20, 2)

    draw.text((left_x1 + 18, left_y1 + 16), "Health Status", font=small_title_font, fill=(40, 55, 80))

    health_text = summary.get("SMART Overall Health", "Unknown")
    health_clean = health_text.upper()
    if "PASS" in health_clean or "GOOD" in health_clean or "OK" in health_clean:
        health_fill_top = (202, 227, 255)
        health_fill_bottom = (83, 173, 231)
        health_outline = (94, 131, 210)
        health_primary = "GOOD"
    else:
        health_fill_top = (255, 227, 193)
        health_fill_bottom = (255, 171, 77)
        health_outline = (196, 124, 55)
        health_primary = health_text[:14] if health_text != "Unknown" else "UNKNOWN"

    gradient_badge(left_x1 + 18, left_y1 + 56, left_x2 - 18, left_y1 + 150,
                   health_fill_top, health_fill_bottom, health_outline, 18)

    health_primary_box = draw.textbbox((0, 0), health_primary, font=big_box_font)
    hpw = health_primary_box[2] - health_primary_box[0]
    draw.text((left_x1 + ((left_col_w - hpw) // 2), left_y1 + 76), health_primary,
              font=big_box_font, fill=(10, 20, 40))

    detail_text = health_text if health_text != health_primary else ""
    if detail_text:
        detail_box = draw.textbbox((0, 0), detail_text, font=small_font)
        dtw = detail_box[2] - detail_box[0]
        draw.text((left_x1 + ((left_col_w - dtw) // 2), left_y1 + 118),
                  detail_text, font=small_font, fill=(10, 20, 40))

    draw.text((left_x1 + 18, left_y1 + 182), "Temperature", font=small_title_font, fill=(40, 55, 80))

    temp_text = summary.get("Temperature", "Unknown")
    gradient_badge(left_x1 + 18, left_y1 + 220, left_x2 - 18, left_y1 + 286,
                   (200, 216, 255), (95, 196, 239), (94, 131, 210), 30)
    temp_box = draw.textbbox((0, 0), temp_text, font=big_box_font)
    tw = temp_box[2] - temp_box[0]
    draw.text((left_x1 + ((left_col_w - tw) // 2), left_y1 + 232),
              temp_text, font=big_box_font, fill=(10, 20, 40))

    right_x1 = left_x2 + panel_gap
    right_y1 = left_y1
    right_x2 = width - margin
    right_y2 = left_y2
    rounded_box(right_x1, right_y1, right_x2, right_y2, (255, 255, 255), (170, 186, 210), 20, 2)

    info_fields = [
        ("Firmware", extract_first_match([r"Firmware Version:\s*(.+)", r"Firmware:\s*(.+)"], smart_text), None),
        ("Serial Number", summary["Serial Number"], None),
        ("Interface", extract_first_match([r"Transport protocol:\s*(.+)", r"Interface:\s*(.+)"], smart_text), None),
        ("Capacity", summary["Capacity"], None),
        ("Power On Count", summary["Power Cycle Count"], None),
        ("Power On Hours", summary["Power-On Hours"], None),
        ("Host Writes", summary["Total Host Writes"], "GB" if summary["Total Host Writes"] not in ("Unknown", "") else None),
        ("Transport", disk.get("tran") or "Unknown", None),
    ]

    info_start_y = right_y1 + 22
    label_x = right_x1 + 24
    value_x = right_x1 + 260
    row_h = 30
    row_gap = 10
    box_w = right_x2 - value_x - 20

    for idx, (label, value, suffix) in enumerate(info_fields):
        y = info_start_y + idx * (row_h + row_gap)
        if y + row_h > right_y2 - 18:
            break

        draw.text((label_x, y + 2), label, font=medium_font, fill=(30, 45, 70))
        draw.rounded_rectangle((value_x, y, value_x + box_w, y + row_h + 6),
                               radius=8, fill=(246, 248, 252), outline=(180, 190, 206), width=1)

        value_display = value if value not in ("Unknown", "", None) else "----"
        if suffix and value_display != "----":
            value_display = f"{value_display} {suffix}"
        draw.text((value_x + 10, y + 3), value_display, font=medium_font, fill=(20, 30, 50))

    table_x1 = margin
    table_y1 = left_y2 + panel_gap
    table_x2 = width - margin
    table_y2 = table_y1 + table_h
    rounded_box(table_x1, table_y1, table_x2, table_y2, (255, 255, 255), (170, 186, 210), 20, 2)

    col_x = {
        "id": table_x1 + 24,
        "name": table_x1 + 90,
        "current": table_x1 + 760,
        "worst": table_x1 + 930,
        "threshold": table_x1 + 1090,
        "raw": table_x1 + 1260,
    }

    header_y = table_y1 + 14
    draw.text((col_x["id"], header_y), "ID", font=table_bold_font, fill=(25, 35, 60))
    draw.text((col_x["name"], header_y), "Attribute Name", font=table_bold_font, fill=(25, 35, 60))
    draw.text((col_x["current"], header_y), "Current", font=table_bold_font, fill=(25, 35, 60))
    draw.text((col_x["worst"], header_y), "Worst", font=table_bold_font, fill=(25, 35, 60))
    draw.text((col_x["threshold"], header_y), "Threshold", font=table_bold_font, fill=(25, 35, 60))
    draw.text((col_x["raw"], header_y), "Raw Value", font=table_bold_font, fill=(25, 35, 60))

    draw.line((table_x1 + 18, table_y1 + table_header_h, table_x2 - 18, table_y1 + table_header_h),
              fill=(170, 186, 210), width=2)

    if attrs:
        for i, attr in enumerate(attrs):
            row_y = table_y1 + table_header_h + (i * table_row_h)
            fill = (248, 250, 253) if i % 2 == 0 else (238, 243, 249)
            draw.rounded_rectangle((table_x1 + 14, row_y + 2, table_x2 - 14, row_y + table_row_h),
                                   radius=6, fill=fill)

            draw.ellipse((table_x1 + 22, row_y + 8, table_x1 + 38, row_y + 24),
                         fill=(96, 177, 235), outline=(88, 132, 199))

            draw.text((col_x["id"], row_y + 4), attr["id"], font=table_font, fill=(20, 30, 50))
            draw.text((col_x["name"], row_y + 4), attr["name"], font=table_font, fill=(20, 30, 50))
            draw.text((col_x["current"], row_y + 4), attr["current"], font=table_font, fill=(20, 30, 50))
            draw.text((col_x["worst"], row_y + 4), attr["worst"], font=table_font, fill=(20, 30, 50))
            draw.text((col_x["threshold"], row_y + 4), attr["threshold"], font=table_font, fill=(20, 30, 50))
            draw.text((col_x["raw"], row_y + 4), attr["raw"], font=table_font, fill=(20, 30, 50))
    else:
        msg = "SMART attribute table not available in a standard ATA format for this device."
        draw.text((table_x1 + 24, table_y1 + table_header_h + 18), msg, font=medium_font, fill=(50, 60, 80))
        draw.text((table_x1 + 24, table_y1 + table_header_h + 58),
                  "The text report still contains the full smartctl output.", font=small_font, fill=(70, 80, 100))

    footer_y = table_y2 + panel_gap
    rounded_box(margin, footer_y, width - margin, footer_y + footer_h, (255, 255, 255), (170, 186, 210), 16, 2)
    footer_text = f"Generated: {time.strftime('%Y-%m-%d %H:%M:%S')}    Source device: {disk['path']}"
    draw.text((margin + 20, footer_y + 12), footer_text, font=small_font, fill=(45, 55, 75))

    img.save(path, dpi=(300, 300))
    return path


def unmount_disk_partitions(disk_path):
    result = run(
        ["lsblk", "-J", "-o", "NAME,MOUNTPOINT", disk_path],
        check=False
    )

    if result.returncode != 0 or not result.stdout.strip():
        return

    data = json.loads(result.stdout)
    devices = data.get("blockdevices", [])

    def walk(children):
        mounts = []
        for child in children or []:
            name = child.get("name")
            mountpoint = child.get("mountpoint")
            if name and mountpoint:
                mounts.append((name, mountpoint))
            mounts.extend(walk(child.get("children", [])))
        return mounts

    for device in devices:
        for name, mountpoint in walk(device.get("children", [])):
            part_path = f"/dev/{name}"
            print(f"Unmounting {part_path} from {mountpoint} ...")
            run(["umount", part_path], check=False)

    run(["sync"], check=False)


def quick_wipe(disk_path):
    print(f"\nStarting quick wipe of {disk_path} ...")
    unmount_disk_partitions(disk_path)

    print("Removing known filesystem signatures...")
    run(["wipefs", "-a", disk_path], check=False)

    print("Blanking first 100 MiB...")
    start_cmd = [
        "dd",
        "if=/dev/zero",
        f"of={disk_path}",
        "bs=1M",
        "count=100",
        "conv=fsync",
        "status=progress"
    ]
    start_proc = subprocess.run(start_cmd)
    if start_proc.returncode != 0:
        raise RuntimeError("Quick wipe failed while blanking the start of the disk.")

    size_bytes = get_disk_size_bytes(os.path.basename(disk_path))
    mib = 1024 * 1024
    tail_size = 100 * mib

    if size_bytes > tail_size:
        seek_mib = (size_bytes - tail_size) // mib
        print("Blanking last 100 MiB...")
        end_cmd = [
            "dd",
            "if=/dev/zero",
            f"of={disk_path}",
            "bs=1M",
            f"seek={seek_mib}",
            "count=100",
            "conv=fsync",
            "status=progress"
        ]
        end_proc = subprocess.run(end_cmd)
        if end_proc.returncode != 0:
            raise RuntimeError("Quick wipe failed while blanking the end of the disk.")

    run(["sync"], check=False)
    print("Quick wipe complete.")


def full_zero_wipe(disk_path):
    print(f"\nStarting full zero wipe of {disk_path} ...")
    print("This can take a while on larger drives.\n")

    unmount_disk_partitions(disk_path)

    cmd = [
        "dd",
        "if=/dev/zero",
        f"of={disk_path}",
        "bs=16M",
        "conv=fsync",
        "status=progress"
    ]

    proc = subprocess.run(cmd)
    if proc.returncode != 0:
        raise RuntimeError("Full zero wipe failed.")

    run(["sync"], check=False)
    print("Full zero wipe complete.")


def wait_for_partition_node(part_path, timeout=15):
    for _ in range(timeout):
        if Path(part_path).exists():
            return True
        time.sleep(1)
    return False


def format_exfat(part_path, label):
    run(["udevadm", "settle"], check=False)
    run(["umount", part_path], check=False)

    print(f"Formatting {part_path} as exFAT ...")

    if label:
        print(f"Using volume label: {label}")
        result = run(["mkfs.exfat", "-L", label, part_path], check=False)
        if result.returncode == 0:
            return
        print("Formatting with label failed. Retrying without a label...")

    result = run(["mkfs.exfat", part_path], check=False)
    if result.returncode != 0:
        raise RuntimeError(f"mkfs.exfat failed for {part_path}.")


def repartition_and_format_exfat(disk_path, disk):
    print(f"\nCreating fresh MBR partition table on {disk_path} ...")
    run(["parted", "-s", disk_path, "mklabel", "msdos"])
    run(["parted", "-s", disk_path, "mkpart", "primary", "1MiB", "100%"])

    run(["partprobe", disk_path], check=False)
    run(["udevadm", "settle"], check=False)

    part_path = get_partition_path(disk_path)

    if not wait_for_partition_node(part_path, timeout=15):
        raise RuntimeError(f"Partition device was not created: {part_path}")

    label = sanitize_exfat_label(disk.get("serial") or "")
    format_exfat(part_path, label)

    return part_path, label


def mount_partition_temp(part_path):
    mount_dir = Path("/tmp") / f"prep_drive_mount_{os.getpid()}"
    mount_dir.mkdir(parents=True, exist_ok=True)

    run(["udevadm", "settle"], check=False)
    run(["mount", part_path, str(mount_dir)])

    return mount_dir


def copy_files_to_drive(part_path, files_to_copy):
    mount_dir = None
    copied = []

    try:
        mount_dir = mount_partition_temp(part_path)

        for src in files_to_copy:
            destination = mount_dir / Path(src).name
            shutil.copy2(src, destination)
            copied.append(destination)

        run(["sync"], check=False)
        return copied
    finally:
        if mount_dir is not None:
            run(["umount", str(mount_dir)], check=False)
            try:
                mount_dir.rmdir()
            except Exception:
                pass


def get_cups_printers():
    result = run(["lpstat", "-p"], check=False)
    if result.returncode != 0:
        return []

    printers = []
    for line in (result.stdout or "").splitlines():
        line = line.strip()
        m = re.match(r"^printer\s+(\S+)\s+", line)
        if m:
            printers.append(m.group(1))
    return printers


def choose_printer(printers):
    if not printers:
        print("No CUPS printers found.")
        return None

    print("\nAvailable printers:\n")
    for i, printer in enumerate(printers, start=1):
        print(f"{i}. {printer}")

    choice = input("\nChoose printer number [Enter to skip]: ").strip()
    if choice == "":
        return None

    if not choice.isdigit():
        print("Invalid printer choice.")
        return None

    idx = int(choice)
    if idx < 1 or idx > len(printers):
        print("Printer choice out of range.")
        return None

    return printers[idx - 1]


def get_printer_default_media(printer):
    media_value = None

    result = run(["lpoptions", "-p", printer], check=False)
    if result.returncode == 0 and result.stdout:
        m = re.search(r"(?:^|\s)media=([^\s]+)", result.stdout)
        if m:
            media_value = m.group(1).strip()

    if media_value:
        return media_value

    result = run(["lpoptions", "-p", printer, "-l"], check=False)
    if result.returncode == 0 and result.stdout:
        for line in result.stdout.splitlines():
            if line.lower().startswith("media/"):
                m = re.search(r"\*([^\s]+)", line)
                if m:
                    return m.group(1).strip()

    return None


def parse_media_dimensions_mm(media_name):
    if not media_name:
        return None

    m = re.search(r"w(\d+)h(\d+)", media_name)
    if not m:
        return None

    width_pt = int(m.group(1))
    height_pt = int(m.group(2))

    width_mm = width_pt / PT_PER_INCH * MM_PER_INCH
    height_mm = height_pt / PT_PER_INCH * MM_PER_INCH

    short_mm = min(width_mm, height_mm)
    long_mm = max(width_mm, height_mm)

    return short_mm, long_mm


def get_label_dimensions_from_media(media_name):
    dims = parse_media_dimensions_mm(media_name)
    if not dims:
        short_mm = DEFAULT_LABEL_SHORT_MM
        long_mm = DEFAULT_LABEL_LONG_MM
    else:
        short_mm, long_mm = dims

    short_px = round(short_mm / MM_PER_INCH * LABEL_DPI)
    long_px = round(long_mm / MM_PER_INCH * LABEL_DPI)
    target_text_width_px = round(min(70.0, max(20.0, long_mm - 12.0)) / MM_PER_INCH * LABEL_DPI)

    return {
        "short_mm": short_mm,
        "long_mm": long_mm,
        "short_px": short_px,
        "long_px": long_px,
        "target_text_width_px": target_text_width_px,
    }


def choose_label_plan():
    choice = input("\nDo you want to print a label after processing? [y/N]: ").strip().lower()
    if choice not in ("y", "yes"):
        return None

    printers = get_cups_printers()
    printer = choose_printer(printers)
    if not printer:
        print("Skipping label printing.")
        return None

    media = get_printer_default_media(printer)
    dims = get_label_dimensions_from_media(media)

    print("\nLabel plan:")
    print(f"  Printer       : {printer}")
    print(f"  CUPS media    : {media or 'Unknown / fallback used'}")
    print(f"  Label size    : {dims['short_mm']:.1f} x {dims['long_mm']:.1f} mm")
    print(f"  Render size   : {dims['long_px']} x {dims['short_px']} px")

    return {
        "printer": printer,
        "media": media,
        "dims": dims,
    }


def get_font_path(prefer_bold=True):
    for path in FONT_CANDIDATES:
        if Path(path).exists():
            if prefer_bold and "Bold" in Path(path).name:
                return path

    for path in FONT_CANDIDATES:
        if Path(path).exists():
            return path
    raise FileNotFoundError("No suitable TTF font found.")


def fit_label_fonts(draw, font_path, serial_text, model_text, max_width, max_height):
    from PIL import ImageFont

    if model_text:
        for serial_size in range(220, 20, -2):
            model_size = max(20, int(serial_size * 0.40))

            serial_font = ImageFont.truetype(font_path, serial_size)
            model_font = ImageFont.truetype(font_path, model_size)

            sb = draw.textbbox((0, 0), serial_text, font=serial_font)
            mb = draw.textbbox((0, 0), model_text, font=model_font)

            sw = sb[2] - sb[0]
            sh = sb[3] - sb[1]
            mw = mb[2] - mb[0]
            mh = mb[3] - mb[1]

            total_h = sh + LABEL_GAP + mh

            if max(sw, mw) <= max_width and total_h <= max_height:
                return serial_font, model_font
    else:
        for serial_size in range(260, 20, -2):
            serial_font = ImageFont.truetype(font_path, serial_size)
            sb = draw.textbbox((0, 0), serial_text, font=serial_font)

            sw = sb[2] - sb[0]
            sh = sb[3] - sb[1]

            if sw <= max_width and sh <= max_height:
                return serial_font, None

    raise RuntimeError("Could not fit the label text onto the selected label size.")


def create_label_image(serial_text, model_text, tested_text, label_plan):
    from PIL import Image, ImageDraw, ImageFont

    serial_text = sanitize_label_text(serial_text, "NO_SERIAL")
    model_text = sanitize_label_text(model_text, "")
    tested_text = sanitize_label_text(tested_text, "")

    dims = label_plan["dims"]
    label_long_px = dims["long_px"]
    label_short_px = dims["short_px"]
    target_text_width_px = dims["target_text_width_px"]

    font_path = get_font_path(prefer_bold=True)

    img = Image.new("L", (label_long_px, label_short_px), 255)
    draw = ImageDraw.Draw(img)

    max_width = min(target_text_width_px, label_long_px - (LABEL_MARGIN_X * 2))
    max_height = label_short_px - (LABEL_MARGIN_Y * 2)

    from PIL import ImageFont

    serial_font = None
    model_font = None
    tested_font = None

    for serial_size in range(220, 20, -2):
        model_size = max(18, int(serial_size * 0.38))
        tested_size = max(16, int(serial_size * 0.24))

        sf = ImageFont.truetype(font_path, serial_size)
        mf = ImageFont.truetype(font_path, model_size)
        tf = ImageFont.truetype(font_path, tested_size)

        sb = draw.textbbox((0, 0), serial_text, font=sf)
        mb = draw.textbbox((0, 0), model_text, font=mf) if model_text else (0, 0, 0, 0)
        tb = draw.textbbox((0, 0), tested_text, font=tf) if tested_text else (0, 0, 0, 0)

        sw = sb[2] - sb[0]
        sh = sb[3] - sb[1]
        mw = mb[2] - mb[0]
        mh = mb[3] - mb[1]
        tw = tb[2] - tb[0]
        th = tb[3] - tb[1]

        total_h = sh
        if model_text:
            total_h += LABEL_GAP + mh
        if tested_text:
            total_h += LABEL_GAP + th

        widest = max(sw, mw, tw)

        if widest <= max_width and total_h <= max_height:
            serial_font = sf
            model_font = mf if model_text else None
            tested_font = tf if tested_text else None
            break

    if serial_font is None:
        raise RuntimeError("Could not fit the label text onto the selected label size.")

    serial_box = draw.textbbox((0, 0), serial_text, font=serial_font)
    serial_w = serial_box[2] - serial_box[0]
    serial_h = serial_box[3] - serial_box[1]

    model_w = model_h = 0
    if model_text and model_font is not None:
        model_box = draw.textbbox((0, 0), model_text, font=model_font)
        model_w = model_box[2] - model_box[0]
        model_h = model_box[3] - model_box[1]

    tested_w = tested_h = 0
    if tested_text and tested_font is not None:
        tested_box = draw.textbbox((0, 0), tested_text, font=tested_font)
        tested_w = tested_box[2] - tested_box[0]
        tested_h = tested_box[3] - tested_box[1]

    total_h = serial_h
    if model_text:
        total_h += LABEL_GAP + model_h
    if tested_text:
        total_h += LABEL_GAP + tested_h

    y = (label_short_px - total_h) // 2

    serial_x = (label_long_px - serial_w) // 2
    draw.text((serial_x, y), serial_text, font=serial_font, fill=0)
    y += serial_h

    if model_text and model_font is not None:
        y += LABEL_GAP
        model_x = (label_long_px - model_w) // 2
        draw.text((model_x, y), model_text, font=model_font, fill=0)
        y += model_h

    if tested_text and tested_font is not None:
        y += LABEL_GAP
        tested_x = (label_long_px - tested_w) // 2
        draw.text((tested_x, y), tested_text, font=tested_font, fill=0)

    out_file = Path("/tmp") / f"dynamic_label_{os.getpid()}.png"
    img.save(out_file, dpi=(LABEL_DPI, LABEL_DPI))
    return out_file


def print_label(disk, label_plan, tested_timestamp):
    if not label_plan:
        return

    printer = label_plan["printer"]
    media = label_plan["media"]

    serial = disk.get("serial") or "NO_SERIAL"
    model = disk.get("model") or ""
    size = disk.get("size_human") or ""
    model_line = f"{model} | {size}" if model and size else (model or size or "")

    tested_short = tested_timestamp[:16]  # YYYY-MM-DD HH:MM
    tested_line = f"Tested: {tested_short} by _____"

    label_image = None
    try:
        label_image = create_label_image(serial, model_line, tested_line, label_plan)

        cmd = [
            "lp",
            "-d", printer,
            "-o", "scaling=100",
            "-o", "position=center",
            "-t", f"DriveLabel-{sanitize_filename(serial)}",
        ]

        if media:
            cmd.extend(["-o", f"media={media}"])

        cmd.append(str(label_image))

        result = run(cmd, check=False)

        if result.returncode == 0:
            print(f"Label sent to printer: {printer}")
            print(f"Label media: {media or 'printer default / fallback'}")
            print(f"Label text: {serial} / {model_line} / {tested_line}")
        else:
            print(f"Failed to print label on {printer}.")
            if result.stderr:
                print(result.stderr.strip())
    finally:
        if label_image is not None:
            try:
                label_image.unlink()
            except Exception:
                pass

def choose_disk(disks):
    print("\nAvailable disks:\n")

    for idx, disk in enumerate(disks, start=1):
        notes = []
        if disk["is_root_disk"]:
            notes.append("SYSTEM DISK - BLOCKED")
        if disk["ro"] == "1":
            notes.append("READ-ONLY")
        if disk["rm"] == "1":
            notes.append("REMOVABLE")

        note_text = f" [{' | '.join(notes)}]" if notes else ""

        print(
            f"{idx}. {disk['path']} | "
            f"Size: {disk['size_human']} | "
            f"Vendor: {disk['vendor'] or 'Unknown'} | "
            f"Model: {disk['model'] or 'Unknown'} | "
            f"Serial: {disk['serial'] or 'Unknown'} | "
            f"Bus: {disk['tran'] or 'Unknown'}"
            f"{note_text}"
        )

    choice = input("\nEnter the number of the disk you want to interrogate/wipe/format: ").strip()

    if not choice.isdigit():
        print("Invalid choice.")
        sys.exit(1)

    index = int(choice)
    if index < 1 or index > len(disks):
        print("Choice out of range.")
        sys.exit(1)

    disk = disks[index - 1]

    if disk["is_root_disk"]:
        print("Refusing to process the system disk.")
        sys.exit(1)

    if disk["ro"] == "1":
        print("Selected disk appears to be read-only. Aborting.")
        sys.exit(1)

    if disk["size_bytes"] >= 2 * 1024 * 1024 * 1024 * 1024:
        print(f"Warning: {disk['path']} is {disk['size_human']}, which is 2 TB or larger.")
        print("You said these are under 2 TB, so stopping to avoid surprises.")
        sys.exit(1)

    return disk


def choose_wipe_type():
    print("\nChoose wipe type:\n")
    print("1. Quick wipe")
    print("   Removes filesystem signatures and blanks the beginning/end of the disk.")
    print("   Fast, but old data may still be recoverable.\n")
    print("2. Full secure zero wipe")
    print("   Overwrites the entire disk with zeros.")
    print("   Much slower, but more thorough.\n")

    choice = input("Enter 1 or 2: ").strip()

    if choice == "1":
        return "quick"
    if choice == "2":
        return "full"

    print("Invalid wipe type selected.")
    sys.exit(1)


def main():
    require_root()
    check_prerequisites()

    disks = get_disks()
    if not disks:
        print("No disks found.")
        sys.exit(1)

    disk = choose_disk(disks)

    print("\nYou selected:")
    print(f"  Device : {disk['path']}")
    print(f"  Size   : {disk['size_human']}")
    print(f"  Vendor : {disk['vendor'] or 'Unknown'}")
    print(f"  Model  : {disk['model'] or 'Unknown'}")
    print(f"  Serial : {disk['serial'] or 'Unknown'}")

    label_plan = choose_label_plan()

    wipe_type = choose_wipe_type()

    print("\nSelected wipe type:")
    print(f"  {'Quick wipe' if wipe_type == 'quick' else 'Full secure zero wipe'}")

    confirm1 = input(f"\nType the device path exactly to continue ({disk['path']}): ").strip()
    if confirm1 != disk["path"]:
        print("Confirmation failed. Aborting.")
        sys.exit(1)

    confirm2 = input("Type WIPE to permanently continue: ").strip()
    if confirm2 != "WIPE":
        print("Final confirmation failed. Aborting.")
        sys.exit(1)

    script_dir = os.getcwd()
    tested_timestamp = time.strftime('%Y-%m-%d %H:%M:%S')

    print("\nCollecting SMART data...")
    smart_text = collect_smart_report(disk["path"])
    summary = parse_smart_summary(smart_text, disk)
    summary_text = format_summary_block(summary)

    print("\n----- LISTING SUMMARY -----")
    print(summary_text)
    print("---------------------------\n")

    print("Writing SMART report file...")
    report_path = write_report(script_dir, disk, smart_text, summary_text)

    print("Creating graphical SMART PNG...")
    png_path = create_smart_graphic_png(script_dir, disk, summary, smart_text, tested_timestamp)

    if wipe_type == "quick":
        quick_wipe(disk["path"])
    else:
        full_zero_wipe(disk["path"])

    part_path, label = repartition_and_format_exfat(disk["path"], disk)

    print("Copying report files to drive root...")
    copied_files = copy_files_to_drive(part_path, [report_path, png_path])

    print("\nDone.")
    print(f"Formatted partition: {part_path}")
    print(f"Volume label: {label}")
    print(f"SMART report saved to: {report_path}")
    print(f"SMART PNG saved to: {png_path}")

    if copied_files:
        for copied in copied_files:
            print(f"Copied to drive root: {copied.name}")

    print("\nListing summary for eBay:\n")
    print(summary_text)

    print_label(disk, label_plan, tested_timestamp)


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\nCancelled.")
        sys.exit(1)
    except Exception as e:
        print(f"\nError: {e}")
        sys.exit(1)