#!/usr/bin/env python3
import os
import sys
import re

# Physical mapping updated to 1-based indexing: (Row 1-6, Column 1-4)
BAY_MAP = {
    "0000:01": {
        "phy0": (1, 1), "phy1": (1, 2), "phy2": (1, 3), "phy3": (1, 4),
        "phy4": (2, 1), "phy5": (2, 2), "phy6": (2, 3), "phy7": (2, 4),
    },
    "0000:02": {
        "phy0": (3, 4), "phy1": (3, 3), "phy2": (3, 2), "phy3": (3, 1),
        "phy4": (4, 4), "phy5": (4, 3), "phy6": (4, 2), "phy7": (4, 1),
    },
    "0000:04": {
        "phy0": (5, 4), "phy1": (5, 3), "phy2": (5, 2), "phy3": (5, 1),
        "phy4": (6, 4), "phy5": (6, 3), "phy6": (6, 2), "phy7": (6, 1),
    }
}

def main():
    export_mapping = (len(sys.argv) > 1 and sys.argv[1] == "--mapping")

    # Initialize a 6x4 empty grid (1-indexed rows 1-6, columns 1-4)
    # We use size 7x5 internally so we can just use the indices 1 to 6 and 1 to 4 directly
    grid = [["[ EMPTY ]" for _ in range(5)] for _ in range(7)]
    path_dir = "/dev/disk/by-path"
    
    if not os.path.exists(path_dir):
        if not export_mapping:
            print(f"Error: {path_dir} does not exist.")
        return

    sas_regex = re.compile(r"pci-(0000:\d{2}):\d{2}\.\d-sas-(phy\d)-lun-0")
    mappings = {}

    for filename in os.listdir(path_dir):
        match = sas_regex.match(filename)
        if match:
            pci_addr = match.group(1)   
            phy_id = match.group(2)     
            
            if pci_addr in BAY_MAP and phy_id in BAY_MAP[pci_addr]:
                row, col = BAY_MAP[pci_addr][phy_id]
                full_path = os.path.join(path_dir, filename)
                try:
                    target = os.readlink(full_path)
                    dev_name = os.path.basename(target) 
                    grid[row][col] = f"[ /dev/{dev_name} ]"
                    mappings[dev_name] = f"R{row}-C{col}"
                except OSError:
                    grid[row][col] = "[ ERROR ]"

    # Bash Mode: Output "sda:R1-C3" format
    if export_mapping:
        for dev, bay in mappings.items():
            print(f"{dev}:{bay}")
        return

    # Visual Mode: Print your layout chart for the operators
    print("\n" + "="*53)
    print("                DRIVE WIPING BAY LAYOUT                ")
    print("="*53)
    
    # Iterate from Row 1 to 6
    for row_idx in range(1, 7):
        # Slice columns 1 through 4
        row_cells = grid[row_idx][1:5]
        row_str = " | ".join(f"{cell:<11}" for cell in row_cells)
        print(f"Row {row_idx}:  | {row_str} |")
        
        # Add visual separation between shelf units
        if row_idx in [2, 4]:
            print("-" * 53)
            
    print("="*53 + "\n")

if __name__ == "__main__":
    main()