#!/bin/bash

#==============================================================================
# Drive Atlas - Server Drive Mapping Tool
# Maps physical drive bays to logical device names using PCI paths
#==============================================================================

#------------------------------------------------------------------------------
# Chassis Type Definitions
# These define the physical layout and display formatting for each chassis type
#------------------------------------------------------------------------------

generate_10bay_layout() {
    local hostname=$1
    build_drive_map

    # Calculate max width needed for drive names
    max_width=0
    for bay in {1..10} "m2-1" "usb1" "usb2"; do
        drive_text="${DRIVE_MAP[$bay]:-EMPTY}"
        text_len=$((${#bay} + 1 + ${#drive_text}))
        [[ $text_len -gt $max_width ]] && max_width=$text_len
    done

    # Add padding for box borders
    box_width=$((max_width + 4))

    # Create box drawing elements
    h_line=$(printf '%*s' "$box_width" '' | tr ' ' '─')

    # USB Section (if applicable)
    if [[ -n "${DRIVE_MAP[usb1]}" || -n "${DRIVE_MAP[usb2]}" ]]; then
        printf "\n    External USB\n"
        printf "    ┌%s┐ ┌%s┐\n" "$h_line" "$h_line"
        printf "    │ %-${max_width}s │ │ %-${max_width}s │\n" "${DRIVE_MAP[usb1]:-EMPTY}" "${DRIVE_MAP[usb2]:-EMPTY}"
        printf "    └%s┘ └%s┘\n\n" "$h_line" "$h_line"
    fi

    # Main chassis section
    printf "┌──────────────────────────────────────────────────────────────┐\n"
    printf "│  %-58s  │\n" "$hostname"
    printf "│  %-58s  │\n" "10-Bay Hot-swap Chassis"
    printf "│                                                              │\n"

    # M.2 NVMe slot if present
    if [[ -n "${DRIVE_MAP[m2-1]}" ]]; then
        printf "│    M.2 NVMe Slot                                             │\n"
        printf "│    ┌%s┐                                             │\n" "$h_line"
        printf "│    │ %-${max_width}s │                                             │\n" "${DRIVE_MAP[m2-1]:-EMPTY}"
        printf "│    └%s┘                                             │\n" "$h_line"
        printf "│                                                              │\n"
    fi

    printf "│    Front Hot-swap Bays                                       │\n"

    # Create bay rows
    printf "│ "
    for bay in {1..10}; do
        printf "┌%s┐" "$h_line"
    done
    printf " │\n│ "

    for bay in {1..10}; do
        printf "│%-2d:%-${max_width}s │" "$bay" "${DRIVE_MAP[$bay]:-EMPTY}"
    done
    printf " │\n│ "

    for bay in {1..10}; do
        printf "└%s┘" "$h_line"
    done
    printf " │\n"

    printf "└──────────────────────────────────────────────────────────────┘\n"
}

generate_large1_layout() {
    build_drive_map

    cat << 'EOF'
┌─────────────────────────────────────────────────────────────┐
│                           large1                            │
│                      Unique 3x5 Grid Chassis                │
│                                                             │
│ ┌──────────────────────────────────────────────┐            │
│ │               Motherboard                    │            │
│ │                                              │            │
│ │                                  ┌──┐┌──┐    │            │
│ │                                  │M1││M2│    │            │
│ │                                  └──┘└──┘    │            │
│ └──────────────────────────────────────────────┘            │
│                                                             │
│ ┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ │
│ │                 │ │                 │ │                 │ │
│ │       1         │ │        2        │ │        3        │ │
│ │                 │ │                 │ │                 │ │
│ └─────────────────┘ └─────────────────┘ └─────────────────┘ │
│ ┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ │
│ │                 │ │                 │ │                 │ │
│ │       4         │ │        5        │ │        6        │ │
│ │                 │ │                 │ │                 │ │
│ └─────────────────┘ └─────────────────┘ └─────────────────┘ │
│ ┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ │
│ │                 │ │                 │ │                 │ │
│ │       7         │ │        8        │ │        9        │ │
│ │                 │ │                 │ │                 │ │
│ └─────────────────┘ └─────────────────┘ └─────────────────┘ │
└─────────────────────────────────────────────────────────────┘
EOF
}

#------------------------------------------------------------------------------
# Server-Specific Drive Mappings
# Maps PCI paths to physical bay numbers for each server
# Format: "pci-path bay-number"
#------------------------------------------------------------------------------

declare -A SERVER_MAPPINGS=(
    # compute-storage-01 (formerly medium2)
    # Motherboard: B650D4U3-2Q/BCM
    # Controller at 0c:00.0 for hot-swap bays
    # Controller at 0d:00.0 for M.2 NVMe
    ["compute-storage-01"]="
        pci-0000:0c:00.0-ata-3 5
        pci-0000:0c:00.0-ata-4 6
        pci-0000:0c:00.0-ata-1 3
        pci-0000:0c:00.0-ata-2 4
        pci-0000:0d:00.0-nvme-1 m2-1
        pci-0000:0b:00.0-usb-0:3:1.0-scsi-0:0:0:0 usb1
        pci-0000:0b:00.0-usb-0:4:1.0-scsi-0:0:0:0 usb2
    "

    # storage-01
    # Different motherboard, no HBA currently
    # TODO: Map actual PCI paths after running diagnose-drives.sh
    ["storage-01"]="
    "

    # large1
    # Unique chassis - 1/1 configuration
    # TODO: Map actual PCI paths after running diagnose-drives.sh
    ["large1"]="
    "
)

declare -A CHASSIS_TYPES=(
    ["compute-storage-01"]="10bay"
    ["compute-storage-gpu-01"]="10bay"
    ["storage-01"]="10bay"
    ["large1"]="large1"
    ["micro1"]="micro"
    ["monitor-02"]="micro"
)

#------------------------------------------------------------------------------
# Core Functions
#------------------------------------------------------------------------------

build_drive_map() {
    local host=$(hostname)
    declare -A drive_map

    local mapping=${SERVER_MAPPINGS[$host]}

    if [[ -n "$mapping" ]]; then
        while read -r path slot; do
            [[ -z "$path" || -z "$slot" ]] && continue

            if [[ -L "/dev/disk/by-path/$path" ]]; then
                local drive=$(readlink -f "/dev/disk/by-path/$path" | sed 's/.*\///')
                drive_map[$slot]=$drive
            fi
        done <<< "$mapping"
    fi

    # Make drive_map available globally
    declare -g -A DRIVE_MAP=()
    for key in "${!drive_map[@]}"; do
        DRIVE_MAP[$key]=${drive_map[$key]}
    done
}

get_drive_smart_info() {
    local device=$1
    local smart_info=$(sudo smartctl -A -i -H /dev/$device 2>/dev/null)
    local temp=$(echo "$smart_info" | grep "Temperature" | awk '{print $10}' | head -1)
    local type=$(echo "$smart_info" | grep "Rotation Rate" | grep -q "Solid State" && echo "SSD" || echo "HDD")
    local health=$(echo "$smart_info" | grep "SMART overall-health" | grep -q "PASSED" && echo "✓" || echo "✗")
    local model=$(echo "$smart_info" | grep "Device Model\|Model Number" | cut -d: -f2 | xargs)

    echo "$type|$temp°C|$health|$model"
}

#------------------------------------------------------------------------------
# Main Display Logic
#------------------------------------------------------------------------------

HOSTNAME=$(hostname)
CHASSIS_TYPE=${CHASSIS_TYPES[$HOSTNAME]:-"unknown"}

# Display chassis layout
case "$CHASSIS_TYPE" in
    "10bay")
        generate_10bay_layout "$HOSTNAME"
        ;;
    "large1")
        generate_large1_layout
        ;;
    "micro")
        echo "Micro server layout not yet implemented"
        ;;
    *)
        echo "┌─────────────────────────────────────────────────────────┐"
        echo "│  Unknown server: $HOSTNAME"
        echo "│  No chassis mapping defined yet"
        echo "│  Run diagnose-drives.sh to gather PCI path information"
        echo "└─────────────────────────────────────────────────────────┘"
        ;;
esac

#------------------------------------------------------------------------------
# Drive Details Section
#------------------------------------------------------------------------------

echo -e "\n=== Drive Details with SMART Status ==="
printf "%-15s %-10s %-8s %-8s %-8s %-30s\n" "DEVICE" "SIZE" "TYPE" "TEMP" "HEALTH" "MODEL"
echo "--------------------------------------------------------------------------------"

# SATA/SAS drives
lsblk -d -o NAME | grep -v "nvme" | grep -v "rbd" | grep -v "loop" | grep -v "NAME" | while read device; do
    if [ -b "/dev/$device" ]; then
        size=$(lsblk -d -n -o SIZE "/dev/$device" 2>/dev/null)
        smart_info=$(get_drive_smart_info "$device")
        IFS='|' read -r type temp health model <<< "$smart_info"
        printf "%-15s %-10s %-8s %-8s %-8s %-30s\n" "/dev/$device" "$size" "$type" "$temp" "$health" "$model"
    fi
done

# NVMe drives
if command -v nvme >/dev/null 2>&1; then
    nvme_drives=$(sudo nvme list 2>/dev/null | grep "^/dev")
    if [ -n "$nvme_drives" ]; then
        echo -e "\n=== NVMe Drives ==="
        printf "%-15s %-10s %-10s %-40s\n" "DEVICE" "SIZE" "TYPE" "MODEL"
        echo "--------------------------------------------------------------------------------"
        echo "$nvme_drives" | awk '{printf "%-15s %-10s %-10s %-40s\n", $1, $6, "NVMe", $3}'
    fi
fi

#------------------------------------------------------------------------------
# Optional sections
#------------------------------------------------------------------------------

# Ceph RBD Devices
rbd_output=$(lsblk -o NAME,SIZE,TYPE,MOUNTPOINT 2>/dev/null | grep "rbd" | sort -V)
if [ -n "$rbd_output" ]; then
    echo -e "\n=== Ceph RBD Devices ==="
    printf "%-15s %-10s %-10s %-20s\n" "DEVICE" "SIZE" "TYPE" "MOUNTPOINT"
    echo "------------------------------------------------------------"
    echo "$rbd_output"
fi

# Show mapping diagnostic info if DEBUG is set
if [[ -n "$DEBUG" ]]; then
    echo -e "\n=== DEBUG: Drive Mappings ==="
    for key in "${!DRIVE_MAP[@]}"; do
        echo "Bay $key: ${DRIVE_MAP[$key]}"
    done | sort -n
fi