#!/bin/sh

# Script to display zpool layout with physical disk sizes
# Usage: ./zpool_disk_sizes.sh <poolname>

if [ -z "$1" ]; then
    echo "Usage: $0 <poolname>"
    exit 1
fi

POOL="$1"

# Check if pool exists
if ! zpool status "$POOL" >/dev/null 2>&1; then
    echo "Error: Pool '$POOL' not found"
    exit 1
fi

# Create temporary files for mappings
GPTID_MAP=$(mktemp)
DISK_SIZE_MAP=$(mktemp)

# Cleanup on exit
trap "rm -f $GPTID_MAP $DISK_SIZE_MAP" EXIT

# Build gptid -> device mapping from glabel status
glabel status | grep "gptid/" | awk '{print $1, $3}' | sed 's/p[12]$//' > "$GPTID_MAP"

# Build device -> size/model/serial mapping from geom disk list
# Convert bytes to base-10 TB (manufacturer size)
geom disk list | awk '
    /^Geom name:/ {
        dev = $3
        size = ""
        model = ""
        serial = ""
    }
    /Mediasize:/ {
        # Extract bytes (first number)
        bytes = $2
        if (bytes > 0) {
            # Convert to base-10 TB (1TB = 1000^4 bytes)
            tb = bytes / 1000000000000
            # Round to whole number for typical drives
            if (tb >= 1) {
                rounded = int(tb + 0.5)
                size = rounded "TB"
            } else {
                # For smaller drives, show in GB
                gb = bytes / 1000000000
                rounded = int(gb + 0.5)
                size = rounded "GB"
            }
        }
    }
    /descr:/ {
        # Extract everything after "descr: "
        sub(/^[[:space:]]*descr:[[:space:]]*/, "")
        model = $0
    }
    /ident:/ {
        # Extract everything after "ident: "
        sub(/^[[:space:]]*ident:[[:space:]]*/, "")
        serial = $0
    }
    # After each disk block, output the mapping
    /^$/ {
        if (dev != "" && size != "") {
            # Use pipe character as delimiter since model names may have spaces
            print dev "|" size "|" model "|" serial
            dev = ""
            size = ""
            model = ""
            serial = ""
        }
    }
    # Handle end of file
    END {
        if (dev != "" && size != "") {
            print dev "|" size "|" model "|" serial
        }
    }
' > "$DISK_SIZE_MAP"

# Function to get device from gptid
get_device() {
    local gptid="$1"
    grep "^$gptid " "$GPTID_MAP" | awk '{print $2}' | head -1
}

# Function to get size/model/serial from device
get_disk_info() {
    local device="$1"
    grep "^$device|" "$DISK_SIZE_MAP" | head -1
}

# Parse zpool status and annotate with sizes
echo "Pool: $POOL"
echo "===================================================================================================================================="

zpool status "$POOL" | awk -v gptid_map="$GPTID_MAP" -v disk_map="$DISK_SIZE_MAP" '
BEGIN {
    # Load gptid -> device mapping
    while ((getline < gptid_map) > 0) {
        gptid_to_dev[$1] = $2
    }
    close(gptid_map)

    # Load device -> size/model/serial mapping (pipe-delimited)
    while ((getline < disk_map) > 0) {
        split($0, parts, "|")
        dev = parts[1]
        dev_to_size[dev] = parts[2]
        dev_to_model[dev] = parts[3]
        dev_to_serial[dev] = parts[4]
    }
    close(disk_map)
}

# Skip until we hit the config section
/^config:/ { in_config = 1; next }

# Print header as-is
in_config && /NAME.*STATE.*READ.*WRITE.*CKSUM/ {
    printf "%-48s %-9s %-6s %-25s %-18s %s\n", "NAME", "DEVICE", "SIZE", "MODEL", "SERIAL", "STATE"
    printf "%s\n", "------------------------------------------------------------------------------------------------------------------------------------"
    next
}

# Process device lines
in_config && /gptid\// {
    # Extract the gptid and state
    match($0, /^[[:space:]]*/)
    indent = RLENGTH
    gptid = $1
    state = $2

    # Look up device and size
    device = gptid_to_dev[gptid]
    if (device == "") device = "???"

    size = dev_to_size[device]
    if (size == "") size = "???"

    model = dev_to_model[device]
    if (model == "") model = "???"

    serial = dev_to_serial[device]
    if (serial == "") serial = "???"

    # Print with alignment
    spaces = ""
    for (i = 0; i < indent; i++) spaces = spaces " "

    printf "%s%-43s %-9s %-6s %-25s %-18s %s\n", spaces, gptid, device, size, model, serial, state
    next
}

# Print vdev lines (raidz2-0, etc.) and pool name
in_config && ($1 ~ /^(raidz|mirror|'$POOL')/ || $1 == "'$POOL'") {
    match($0, /^[[:space:]]*/)
    indent = RLENGTH
    name = $1
    state = $2

    spaces = ""
    for (i = 0; i < indent; i++) spaces = spaces " "

    printf "%s%-43s %-9s %-6s %-25s %-18s %s\n", spaces, name, "", "", "", "", state
    next
}

# Stop at errors line
in_config && /^errors:/ { exit }
'

echo ""
