#!/bin/bash

get_device_info() {
    local pci_addr=$1
    local info=$(lspci -s "$pci_addr")
    echo "$info"
}

get_drive_details() {
    local device=$1
    local size=$(lsblk -d -o NAME,SIZE | grep "$device" | awk '{print $2}')
    echo "$size"
}

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" | cut -d: -f2 | xargs)
    
    echo "$type|$temp°C|$health|$model"
}

get_drives_info() {
    local path="/dev/disk/by-path"
    for drive in "$path"/*; do
        if [ -L "$drive" ]; then
            echo "$(basename "$drive") $(readlink -f "$drive")"
        fi
    done
}

declare -A DRIVE_MAPPINGS=(
    ["medium2"]="
        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 11
        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
    "
)

build_drive_map() {
    local host=$(hostname)
    declare -A drive_map
    
    echo "DEBUG: Current host: $host"
    echo "DEBUG: Mapping found: ${DRIVE_MAPPINGS[$host]}"
    
    local mapping=${DRIVE_MAPPINGS[$host]}
    
    if [[ -n "$mapping" ]]; then
        while read -r path slot; do
            [[ -z "$path" || -z "$slot" ]] && continue
            
            echo "DEBUG: Checking path: $path for slot: $slot"
            
            if [[ -L "/dev/disk/by-path/$path" ]]; then
                local drive=$(readlink -f "/dev/disk/by-path/$path" | sed 's/.*\///')
                drive_map[$slot]=$drive
                echo "DEBUG: Mapped slot $slot to drive $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]}
        echo "DEBUG: Final mapping - slot $key: ${drive_map[$key]}"
    done
}

# Define the ASCII art maps
large1=''' 
┌─────────────────────────────────────────────────────────────┐
│                           large1                            │
│                                                             │
│ ┌──────────────────────────────────────────────┐            │
│ │               Motherboard                    │            │
│ │                                              │            │
│ │                                  ┌──┐┌──┐    │            │
│ │                                  │M1││M2│    │            │
│ │                                  └──┘└──┘    │            │
│ └──────────────────────────────────────────────┘            │
│                                                             │
│ ┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ │
│ │                 │ │                 │ │                 │ │
│ │       1         │ │        2        │ │        3        │ │
│ │                 │ │                 │ │                 │ │
│ └─────────────────┘ └─────────────────┘ └─────────────────┘ │
│ ┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ │
│ │                 │ │                 │ │                 │ │
│ │       4         │ │        5        │ │        6        │ │
│ │                 │ │                 │ │                 │ │
│ └─────────────────┘ └─────────────────┘ └─────────────────┘ │
│ ┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐ │
│ │                 │ │                 │ │                 │ │
│ │       7         │ │        8        │ │        9        │ │
│ │                 │ │                 │ │                 │ │
│ └─────────────────┘ └─────────────────┘ └─────────────────┘ │
└─────────────────────────────────────────────────────────────┘
'''

compute-storage-gpu-01=''' 
┌─────────────────────────────────────────────────────────────┐
│                                                             │
│ ┌────────────┐ ┌────────────┐ ┌────────────┐ ┌────────────┐ │
│ │     1      │ │     2      │ │      3     │ │     4      │ │
│ └────────────┘ └────────────┘ └────────────┘ └────────────┘ │
│                                                             │
│ ┌────────────┐ ┌────────────┐ ┌────────────┐ ┌────────────┐ │
│ │     5      │ │     6      │ │      7     │ │     8      │ │
│ └────────────┘ └────────────┘ └────────────┘ └────────────┘ │
│                                                             │
│                                                             │
│                                                             │
│                                                 ┌─────────┐ │
│                           compute-storage-gpu-01               │    9    │ │
│                                                 └─────────┘ │
│                                                 ┌─────────┐ │
│                                                 │   10    │ │
│                                                 └─────────┘ │
│                                                             │
└─────────────────────────────────────────────────────────────┘
'''

generate_medium2_layout() {
    build_drive_map
    
    # Calculate max width needed for drive names
    max_width=0
    for bay in {1..10} "11" "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
    printf "\n    External USB [0b:00.0]\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"
    
    # Main chassis section
    printf "┌──────────────────────────────────────────────────────────────┐\n"
    printf "│                     B650D4U3-2Q/BCM                          │\n"
    printf "│                                                              │\n"
    printf "│    NVMe [0d:00.0] Bay 11                                    │\n"
    printf "│    ┌%s┐                                             │\n" "$h_line"
    printf "│    │ %-${max_width}s │                                             │\n" "${DRIVE_MAP[11]:-EMPTY}"
    printf "│    └%s┘                                             │\n" "$h_line"
    printf "│                                                              │\n"
    printf "│    Front Hot-swap Bays [0c:00.0]                            │\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"
}

microGeneric=''' 
 ┌─┐  ┌─┐ 
┌└─┘──└─┘┐
│ 1    2 │
│        │
│        │
│        │
│        │
│        │
│        │
│        │
└────────┘
'''

# Get the hostname
HOSTNAME=$(hostname)

# ASCII art based on hostname
case "$HOSTNAME" in
    "large1")
        echo -e "$large1"
        ;;
    "compute-storage-gpu-01")
        echo -e "$compute-storage-gpu-01"
        ;;
    "medium2")
        generate_medium2_layout
        ;;
    "micro1" | "monitor-02")
        echo -e "$microGeneric"
        ;;
    *)
        echo -e "No ASCII map defined for this hostname."
        ;;
esac

map_drives_to_layout() {
    local server_type=$1
    case $server_type in
        "large1")
            for i in {1..9}; do
                local drive_info=$(get_drive_info_for_position $i)
                echo "Position $i: $drive_info"
            done
            ;;
    esac
}

DRIVE_PATHS=$(get_drives_info | awk '{print $1, $2}')

# Initialize array for "not found" messages
not_found=()


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

# For SATA drives
lsblk -d -o NAME | grep -v "nvme" | grep -v "rbd" | while read device; do
    size=$(get_drive_details "$device")
    smart_info=$(get_drive_smart_info "$device")
    IFS='|' read -r type temp health model <<< "$smart_info"
    printf "%-15s %-10s %-8s %-8s %-20s %-30s\n" "/dev/$device" "$size" "$type" "$temp" "$health" "$model"
done

# For NVMe drives
if [ -n "$nvme_drives" ]; then
    while read -r line; do
        device=$(echo "$line" | awk '{print $1}' | sed 's/.*\///')
        size=$(echo "$line" | awk '{print $6}')
        smart_info=$(get_drive_smart_info "$device")
        IFS='|' read -r type temp health model <<< "$smart_info"
        printf "%-15s %-10s %-8s %-8s %-20s %-30s\n" "/dev/$device" "$size" "$type" "$temp" "$health" "$model"
    done <<< "$nvme_drives"
fi

# Show NVMe Drives only if present
nvme_drives=$(sudo nvme list | grep "^/dev")
if [ -n "$nvme_drives" ]; then
    echo -e "\n=== NVMe Drives ===\n"
    printf "%-15s %-10s %-10s %-20s\n" "DEVICE" "SIZE" "TYPE" "MODEL"
    echo "------------------------------------------------------------"
    echo "$nvme_drives" | awk '{printf "%-15s %-10s %-10s %-20s\n", $1, $6, "NVMe", $3}'
else
    not_found+=("NVMe drives")
fi

# Show MMC Drives only if present
mmc_output=$(lsblk -o NAME,SIZE,TYPE,MOUNTPOINT | grep "mmcblk" | sort)
if [ -n "$mmc_output" ]; then
    echo -e "\n=== MMC Drives ===\n"
    printf "%-15s %-10s %-10s %-20s\n" "DEVICE" "SIZE" "TYPE" "MOUNTPOINT"
    echo "------------------------------------------------------------"
    echo "$mmc_output"
fi

# Show SATA Drives only if present
sata_output=$(lsblk -d -o NAME,SIZE,TYPE,MOUNTPOINT | grep "disk" | grep -v "nvme" | grep -v "rbd" | sort | column -t)s
if [ -n "$sata_output" ]; then
    echo -e "\n=== SATA Drives ===\n"
    printf "%-15s %-10s %-10s %-20s\n" "DEVICE" "SIZE" "TYPE" "MOUNTPOINT"
    echo "------------------------------------------------------------"
    echo "$sata_output"
fi

# Show Ceph RBD Devices only if present
rbd_output=$(lsblk -o NAME,SIZE,TYPE,MOUNTPOINT | grep "rbd" | sort -V)
if [ -n "$rbd_output" ]; then
    echo -e "\n=== Ceph RBD Devices ===\n"
    printf "%-15s %-10s %-10s %-20s\n" "DEVICE" "SIZE" "TYPE" "MOUNTPOINT"
    echo "------------------------------------------------------------"
    echo "$rbd_output"
else
    not_found+=("RBD devices")
fi

# Check RAID
if ! [ -f /proc/mdstat ] || ! grep -q "active" /proc/mdstat; then
    not_found+=("Software RAID")
fi

# Check ZFS
if ! command -v zpool >/dev/null 2>&1 || [ -z "$(sudo zpool status 2>/dev/null)" ]; then
    not_found+=("ZFS pools")
fi

# Display consolidated "not found" messages at the end
if [ ${#not_found[@]} -gt 0 ]; then
    echo -e "\n=== Not Found ===\n"
    printf "%s\n" "${not_found[@]}"
fi