#!/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

    # Fixed width for consistent box drawing (fits device names like "nvme0n1")
    local drive_width=10

    # Main chassis section
    printf "┌────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┐\n"
    printf "│  %-126s  │\n" "$hostname - Sliger CX4712 (10x 3.5\" Hot-swap)"
    printf "│                                                                                                                                │\n"

    # Show storage controllers
    printf "│  Storage Controllers:                                                                                                          │\n"
    while IFS= read -r ctrl; do
        [[ -n "$ctrl" ]] && printf "│  %-126s│\n" "$ctrl"
    done < <(get_storage_controllers)
    printf "│                                                                                                                                │\n"

    # M.2 NVMe slot if present
    if [[ -n "${DRIVE_MAP[m2-1]}" ]]; then
        printf "│    M.2 NVMe: %-10s                                                                                                          │\n" "${DRIVE_MAP[m2-1]}"
        printf "│                                                                                                                                │\n"
    fi

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

    # Bay top borders
    printf "│    "
    for bay in {1..10}; do
        printf "┌──────────┐ "
    done
    printf "  │\n"

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

    # Bay bottom borders
    printf "│    "
    for bay in {1..10}; do
        printf "└──────────┘ "
    done
    printf "  │\n"

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

generate_micro_layout() {
    local hostname=$1
    build_drive_map

    # Check for eMMC storage
    local emmc_device=""
    if [[ -b /dev/mmcblk0 ]]; then
        emmc_device="mmcblk0"
    fi

    printf "┌─────────────────────────────────────────────────────────────┐\n"
    printf "│  %-57s  │\n" "$hostname - Micro SBC"
    printf "│                                                             │\n"
    printf "│  Storage Controllers:                                       │\n"
    while IFS= read -r ctrl; do
        [[ -n "$ctrl" ]] && printf "│  %-57s│\n" "$ctrl"
    done < <(get_storage_controllers)
    printf "│                                                             │\n"

    # Show eMMC if present
    if [[ -n "$emmc_device" ]]; then
        local emmc_size=$(lsblk -d -n -o SIZE "/dev/$emmc_device" 2>/dev/null | xargs)
        printf "│    ┌─────────────────────────────────────────────────────┐ │\n"
        printf "│    │  Onboard eMMC: %-10s (%s)                  │ │\n" "$emmc_device" "$emmc_size"
        printf "│    └─────────────────────────────────────────────────────┘ │\n"
        printf "│                                                             │\n"
    fi

    printf "│    SATA Ports (rear):                                       │\n"
    printf "│    ┌──────────────┐  ┌──────────────┐                        │\n"
    printf "│    │ 1: %-9s │  │ 2: %-9s │                        │\n" "${DRIVE_MAP[1]:-EMPTY}" "${DRIVE_MAP[2]:-EMPTY}"
    printf "│    └──────────────┘  └──────────────┘                        │\n"
    printf "└─────────────────────────────────────────────────────────────┘\n"
}

generate_large1_layout() {
    local hostname=$1
    build_drive_map

    # large1 has 3 stacks of 5 bays at front (15 total) + 2 M.2 slots
    # Physical bay mapping TBD - current mapping is by controller order
    printf "┌─────────────────────────────────────────────────────────────────────────┐\n"
    printf "│  %-69s  │\n" "$hostname - Rosewill RSV-L4500U (15x 3.5\" Bays)"
    printf "│                                                                         │\n"
    printf "│  Storage Controllers:                                                   │\n"
    while IFS= read -r ctrl; do
        [[ -n "$ctrl" ]] && printf "│  %-69s│\n" "$ctrl"
    done < <(get_storage_controllers)
    printf "│                                                                         │\n"
    printf "│  M.2 NVMe:  M1: %-10s   M2: %-10s                            │\n" "${DRIVE_MAP[m2-1]:-EMPTY}" "${DRIVE_MAP[m2-2]:-EMPTY}"
    printf "│                                                                         │\n"
    printf "│  Front Bays (3 stacks x 5 rows):  [Bay mapping TBD]                     │\n"
    printf "│     Stack A       Stack B       Stack C                                 │\n"
    printf "│  ┌──────────┐  ┌──────────┐  ┌──────────┐                               │\n"
    printf "│  │1:%-8s│  │2:%-8s│  │3:%-8s│                               │\n" "${DRIVE_MAP[1]:-EMPTY}" "${DRIVE_MAP[2]:-EMPTY}" "${DRIVE_MAP[3]:-EMPTY}"
    printf "│  ├──────────┤  ├──────────┤  ├──────────┤                               │\n"
    printf "│  │4:%-8s│  │5:%-8s│  │6:%-8s│                               │\n" "${DRIVE_MAP[4]:-EMPTY}" "${DRIVE_MAP[5]:-EMPTY}" "${DRIVE_MAP[6]:-EMPTY}"
    printf "│  ├──────────┤  ├──────────┤  ├──────────┤                               │\n"
    printf "│  │7:%-8s│  │8:%-8s│  │9:%-8s│                               │\n" "${DRIVE_MAP[7]:-EMPTY}" "${DRIVE_MAP[8]:-EMPTY}" "${DRIVE_MAP[9]:-EMPTY}"
    printf "│  ├──────────┤  ├──────────┤  ├──────────┤                               │\n"
    printf "│  │10:%-7s│  │11:%-7s│  │12:%-7s│                               │\n" "${DRIVE_MAP[10]:-EMPTY}" "${DRIVE_MAP[11]:-EMPTY}" "${DRIVE_MAP[12]:-EMPTY}"
    printf "│  ├──────────┤  ├──────────┤  ├──────────┤                               │\n"
    printf "│  │13:%-7s│  │14:%-7s│  │15:%-7s│                               │\n" "${DRIVE_MAP[13]:-EMPTY}" "${DRIVE_MAP[14]:-EMPTY}" "${DRIVE_MAP[15]:-EMPTY}"
    printf "│  └──────────┘  └──────────┘  └──────────┘                               │\n"
    printf "└─────────────────────────────────────────────────────────────────────────┘\n"
}

#------------------------------------------------------------------------------
# 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 with AMD SATA controller
    # HBA: LSI SAS3008 at 01:00.0 (mini-SAS HD ports)
    # Cable mapping from user notes:
    #   - Mobo SATA: top-right=bay1, bottom-right=bay2, bottom-left=bay3, top-left=bay4
    #   - HBA bottom mini-SAS: bays 5,6,7,8
    #   - HBA top mini-SAS: bays 9,10
    ["compute-storage-01"]="
        pci-0000:0d:00.0-ata-2 1
        pci-0000:0d:00.0-ata-1 2
        pci-0000:0d:00.0-ata-3 3
        pci-0000:0d:00.0-ata-4 4
        pci-0000:01:00.0-sas-phy6-lun-0 5
        pci-0000:01:00.0-sas-phy7-lun-0 6
        pci-0000:01:00.0-sas-phy5-lun-0 7
        pci-0000:01:00.0-sas-phy2-lun-0 8
        pci-0000:01:00.0-sas-phy4-lun-0 9
        pci-0000:01:00.0-sas-phy3-lun-0 10
        pci-0000:0e:00.0-nvme-1 m2-1
    "

    # compute-storage-gpu-01
    # Motherboard: ASUS PRIME B550-PLUS with AMD SATA controller at 02:00.1
    # 5 SATA ports + 1 M.2 NVMe slot
    # sdf is USB/card reader - not mapped
    ["compute-storage-gpu-01"]="
        pci-0000:02:00.1-ata-1 1
        pci-0000:02:00.1-ata-2 2
        pci-0000:02:00.1-ata-3 3
        pci-0000:02:00.1-ata-4 4
        pci-0000:02:00.1-ata-5 5
        pci-0000:0c:00.0-nvme-1 m2-1
    "

    # storage-01
    # Motherboard: ASRock A320M-HDV R4.0 with AMD SATA controller at 02:00.1
    # 4 SATA ports used (ata-1, ata-2, ata-5, ata-6) - ata-3/4 empty
    ["storage-01"]="
        pci-0000:02:00.1-ata-1 1
        pci-0000:02:00.1-ata-2 2
        pci-0000:02:00.1-ata-5 3
        pci-0000:02:00.1-ata-6 4
    "

    # large1
    # Custom tower with multiple controllers:
    # - HBA: LSI SAS2008 at 10:00.0 (7 drives)
    # - AMD SATA at 16:00.1 (3 drives)
    # - ASMedia SATA at 25:00.0 (2 drives)
    # - 2x NVMe slots
    ["large1"]="
        pci-0000:10:00.0-sas-phy0-lun-0 1
        pci-0000:10:00.0-sas-phy1-lun-0 2
        pci-0000:10:00.0-sas-phy3-lun-0 3
        pci-0000:10:00.0-sas-phy4-lun-0 4
        pci-0000:10:00.0-sas-phy5-lun-0 5
        pci-0000:10:00.0-sas-phy6-lun-0 6
        pci-0000:10:00.0-sas-phy7-lun-0 7
        pci-0000:16:00.1-ata-3 8
        pci-0000:16:00.1-ata-7 9
        pci-0000:16:00.1-ata-8 10
        pci-0000:25:00.0-ata-1 11
        pci-0000:25:00.0-ata-2 12
        pci-0000:2a:00.0-nvme-1 m2-1
        pci-0000:26:00.0-nvme-1 m2-2
    "

    # micro1
    # ZimaBoard 832 - Single board computer
    # 2 SATA ports on rear (currently unused)
    # Boot from onboard eMMC (mmcblk0)
    # SATA controller at 00:12.0
    ["micro1"]="
    "

    # monitor-02
    # ZimaBoard 832 - Single board computer
    # 2 SATA ports on rear (currently unused)
    # Boot from onboard eMMC (mmcblk0)
    # SATA controller would be at a specific PCI address when drives connected
    ["monitor-02"]="
    "
)

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

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

get_storage_controllers() {
    # Returns a formatted list of storage controllers (HBAs, SATA, NVMe)
    lspci 2>/dev/null | grep -iE "SAS|SATA|RAID|Mass storage|NVMe" | while read -r line; do
        pci_addr=$(echo "$line" | awk '{print $1}')
        # Get short description (strip PCI address)
        desc=$(echo "$line" | sed 's/^[0-9a-f:.]\+ //')
        echo "    $pci_addr: $desc"
    done
}

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)
    local serial=$(echo "$smart_info" | grep "Serial Number" | awk '{print $3}')

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

#------------------------------------------------------------------------------
# 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 "$HOSTNAME"
        ;;
    "micro")
        generate_micro_layout "$HOSTNAME"
        ;;
    *)
        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 (by Bay Position) ==="
printf "%-5s %-15s %-10s %-8s %-8s %-8s %-30s %-20s %-12s %-10s %-10s\n" "BAY" "DEVICE" "SIZE" "TYPE" "TEMP" "HEALTH" "MODEL" "SERIAL" "CEPH OSD" "STATUS" "USAGE"
echo "----------------------------------------------------------------------------------------------------------------------------------------------------"

# Build reverse map: device -> bay
declare -A DEVICE_TO_BAY
for bay in "${!DRIVE_MAP[@]}"; do
    device="${DRIVE_MAP[$bay]}"
    if [[ -n "$device" && "$device" != "EMPTY" ]]; then
        DEVICE_TO_BAY[$device]=$bay
    fi
done

# Sort drives by bay position
for bay in $(printf '%s\n' "${!DRIVE_MAP[@]}" | grep -E '^[0-9]+$' | sort -n); do
    device="${DRIVE_MAP[$bay]}"
    if [[ -n "$device" && "$device" != "EMPTY" && -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 serial <<< "$smart_info"

        # Check for Ceph OSD
        osd_id=$(ceph-volume lvm list 2>/dev/null | grep -B 20 "/dev/$device" | grep "osd id" | awk '{print "osd."$3}' | head -1)

        # Get Ceph status if OSD exists
        ceph_status="-"
        if [[ -n "$osd_id" ]]; then
            # Get in/out and up/down status from ceph osd tree
            osd_num=$(echo "$osd_id" | sed 's/osd\.//')
            # Parse ceph osd tree output - column 5 is STATUS (up/down), column 6 is REWEIGHT (1.0 = in, 0 = out)
            tree_line=$(ceph osd tree 2>/dev/null | grep -E "^\s*${osd_num}\s+" | grep "osd.${osd_num}")
            up_status=$(echo "$tree_line" | awk '{print $5}')
            reweight=$(echo "$tree_line" | awk '{print $6}')

            # Default to unknown if we can't parse
            [[ -z "$up_status" ]] && up_status="unknown"
            [[ -z "$reweight" ]] && reweight="0"

            # Determine in/out based on reweight (1.0 = in, 0 = out)
            if (( $(echo "$reweight > 0" | bc -l 2>/dev/null || echo 0) )); then
                in_status="in"
            else
                in_status="out"
            fi

            ceph_status="${up_status}/${in_status}"
        else
            osd_id="-"
        fi

        # Check if boot drive
        usage="-"
        if mount | grep -q "^/dev/${device}"; then
            mount_point=$(mount | grep "^/dev/${device}" | awk '{print $3}' | head -1)
            if [[ "$mount_point" == "/" ]]; then
                usage="BOOT"
            else
                usage="$mount_point"
            fi
        fi

        printf "%-5s %-15s %-10s %-8s %-8s %-8s %-30s %-20s %-12s %-10s %-10s\n" "$bay" "/dev/$device" "$size" "$type" "$temp" "$health" "$model" "$serial" "$osd_id" "$ceph_status" "$usage"
    fi
done

# NVMe drives
nvme_devices=$(lsblk -d -n -o NAME,SIZE | grep "^nvme" 2>/dev/null)
if [ -n "$nvme_devices" ]; then
    echo -e "\n=== NVMe Drives ==="
    printf "%-15s %-10s %-10s %-40s %-25s\n" "DEVICE" "SIZE" "TYPE" "MODEL" "SERIAL"
    echo "------------------------------------------------------------------------------------------------------"
    echo "$nvme_devices" | while read -r name size; do
        device="/dev/$name"
        # Get model and serial from smartctl for accuracy
        smart_info=$(sudo smartctl -i "$device" 2>/dev/null)
        model=$(echo "$smart_info" | grep "Model Number" | cut -d: -f2 | xargs)
        serial=$(echo "$smart_info" | grep "Serial Number" | cut -d: -f2 | xargs)
        [[ -z "$model" ]] && model="-"
        [[ -z "$serial" ]] && serial="-"
        printf "%-15s %-10s %-10s %-40s %-25s\n" "$device" "$size" "NVMe" "$model" "$serial"
    done
fi

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

# Ceph RBD Devices
rbd_devices=$(lsblk -d -n -o NAME,SIZE,TYPE 2>/dev/null | grep "rbd" | sort -V)
if [ -n "$rbd_devices" ]; then
    echo -e "\n=== Ceph RBD Devices ==="
    printf "%-15s %-10s %-10s %-30s\n" "DEVICE" "SIZE" "TYPE" "MOUNTPOINT"
    echo "------------------------------------------------------------"
    echo "$rbd_devices" | while read -r name size type; do
        # Get mountpoint if any
        mountpoint=$(lsblk -n -o MOUNTPOINT "/dev/$name" 2>/dev/null | head -1)
        [[ -z "$mountpoint" ]] && mountpoint="-"
        printf "%-15s %-10s %-10s %-30s\n" "/dev/$name" "$size" "$type" "$mountpoint"
    done
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