#!/bin/bash

# Script to check BlueStore compression settings for OSDs in a specified tree
# Usage: ./check_osd_compression.sh [--root root_name] [--hosts host1,host2,host3]

set -e

# Parse command line arguments
FILTER_HOSTS=""
ROOT_NAME=""  # No default, require explicit selection
while [[ $# -gt 0 ]]; do
    case $1 in
        --hosts)
            FILTER_HOSTS="$2"
            shift 2
            ;;
        --root)
            ROOT_NAME="$2"
            shift 2
            ;;
        -h|--help)
            echo "Usage: $0 [options]"
            echo
            echo "This script checks BlueStore compression settings for all active OSDs"
            echo "in a specified Ceph tree and displays a formatted report."
            echo
            echo "Options:"
            echo "  --root root_name             Specify which root tree to examine (REQUIRED)"
            echo "  --hosts host1,host2,host3    Only check OSDs on specified hosts (comma-separated)"
            echo "  -h, --help                   Show this help message"
            echo
            echo "Examples:"
            echo "  $0 --root platter           # Check all hosts in platter tree"
            echo "  $0 --root default           # Check all hosts in default tree"
            echo "  $0 --root rbd               # Check all hosts in default tree (using alias)"
            echo "  $0 --root platter --hosts csn01,csn06      # Check only csn01 and csn06 in platter tree"
            echo "  $0 --root rbd --hosts csn12                # Check only csn12 in default tree"
            echo
            echo "To see available trees, run: ceph osd tree | grep root"
            echo
            echo "Requirements:"
            echo "  - Ceph client tools must be installed"
            echo "  - Valid Ceph configuration (ceph.conf and keyring)"
            echo "  - Network connectivity to Ceph cluster"
            echo
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            echo "Use --help for usage information"
            exit 1
            ;;
    esac
done

# If no root specified, show help
if [[ -z "$ROOT_NAME" ]]; then
    echo "Error: --root parameter is required"
    echo
    echo "Usage: $0 --root <tree_name> [--hosts host1,host2]"
    echo
    echo "Available trees in your cluster:"
    if command -v ceph &> /dev/null && ceph status &> /dev/null 2>&1; then
        ceph osd tree | grep "root" | awk '{print "  - " $4}'
        echo "  - rbd (alias for default)"
    else
        echo "  (Cannot connect to Ceph cluster to show available trees)"
        echo "  - rbd (alias for default)"
    fi
    echo
    echo "Use --help for full usage information"
    exit 1
fi

# Handle "rbd" alias for "default"
if [[ "$ROOT_NAME" == "rbd" ]]; then
    ROOT_NAME="default"
    echo "Note: Using 'rbd' as alias for 'default' tree"
fi

# Convert comma-separated hosts to array immediately after parsing
declare -a TARGET_HOSTS
if [[ -n "$FILTER_HOSTS" ]]; then
    IFS=',' read -ra TARGET_HOSTS <<< "$FILTER_HOSTS"
    # Trim whitespace from each host
    for i in "${!TARGET_HOSTS[@]}"; do
        TARGET_HOSTS[$i]=$(echo "${TARGET_HOSTS[$i]}" | xargs)
    done
fi
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Function to print colored output
print_status() {
    local color=$1
    local message=$2
    echo -e "${color}${message}${NC}"
}

# Function to check if a host should be processed
should_process_host() {
    local host=$1
    
    # If no filter specified, process all hosts
    if [[ ${#TARGET_HOSTS[@]} -eq 0 ]]; then
        return 0
    fi
    
    # Check if host is in the target list
    for target_host in "${TARGET_HOSTS[@]}"; do
        if [[ "$host" == "$target_host" ]]; then
            return 0
        fi
    done
    
    return 1
}
# Function to get OSD list with hosts from specified tree
get_tree_osds_with_hosts() {
    local root_name=$1
    # Parse the osd tree output to extract host and OSD relationships ONLY from specified tree
    ceph osd tree | awk -v root="$root_name" '
    BEGIN { in_target_tree = 0 }
    $0 ~ ("root " root) { in_target_tree = 1; next }
    /root/ && !($0 ~ ("root " root)) { in_target_tree = 0; next }
    in_target_tree && /host/ { 
        current_host = $4 
    }
    in_target_tree && /osd\.[0-9]+/ && !/down|destroyed/ { 
        gsub(/osd\./, "", $4)
        print current_host "," $4 
    }'
}

# Function to check compression settings for a single OSD
check_osd_compression() {
    local osd_id=$1
    
    # Try to get actual running configuration from the OSD daemon first
    local config_output
    if config_output=$(ceph tell osd.${osd_id} config show 2>/dev/null); then
        # Parse compression mode from JSON output
        local comp_mode=$(echo "$config_output" | grep '"bluestore_compression_mode"' | cut -d'"' -f4)
        local comp_algo=$(echo "$config_output" | grep '"bluestore_compression_algorithm"' | cut -d'"' -f4)
        
        # Fallback if parsing failed
        if [[ -z "$comp_mode" ]]; then
            comp_mode="unknown"
        fi
        if [[ -z "$comp_algo" ]]; then
            comp_algo="unknown"
        fi
        
        echo "$osd_id,$comp_mode,$comp_algo,running"
    else
        # Fallback to config get if tell fails - this shows configured values, not running values
        local comp_mode=$(ceph config get osd.${osd_id} bluestore_compression_mode 2>/dev/null || echo "none")
        local comp_algo=$(ceph config get osd.${osd_id} bluestore_compression_algorithm 2>/dev/null || echo "snappy")
        
        # If we get empty values, try global defaults
        if [[ -z "$comp_mode" || "$comp_mode" == "" ]]; then
            comp_mode=$(ceph config get global bluestore_compression_mode 2>/dev/null || echo "none")
        fi
        if [[ -z "$comp_algo" || "$comp_algo" == "" ]]; then
            comp_algo=$(ceph config get global bluestore_compression_algorithm 2>/dev/null || echo "snappy")
        fi
        
        echo "$osd_id,$comp_mode,$comp_algo,config"
    fi
}

# Main execution
main() {
    print_status $BLUE "=== Ceph OSD BlueStore Compression Status Report ==="
    print_status $BLUE "Checking OSDs in the '$ROOT_NAME' tree..."
    echo
    
    # Check if ceph command is available
    if ! command -v ceph &> /dev/null; then
        print_status $RED "Error: 'ceph' command not found. Please ensure Ceph is installed and configured."
        exit 1
    fi
    
    # Check cluster connectivity
    if ! ceph status &> /dev/null; then
        print_status $RED "Error: Cannot connect to Ceph cluster. Please check your configuration."
        exit 1
    fi
    
    # Get list of OSDs with their hosts from specified tree
    if [[ ${#TARGET_HOSTS[@]} -gt 0 ]]; then
        print_status $YELLOW "Gathering OSD list for specified hosts: ${TARGET_HOSTS[*]} (from '$ROOT_NAME' tree)"
    else
        print_status $YELLOW "Gathering OSD list with hosts from '$ROOT_NAME' tree..."
    fi
    
    # Create associative array to group OSDs by host
    declare -A host_osds
    while IFS=',' read -r host osd_id; do
        if [[ -n "$host" && -n "$osd_id" ]]; then
            # Only include hosts we want to process
            if should_process_host "$host"; then
                host_osds["$host"]+="$osd_id "
            fi
        fi
    done < <(get_tree_osds_with_hosts "$ROOT_NAME")
    
    # Count total OSDs
    local total_osds=0
    for host in "${!host_osds[@]}"; do
        local osd_count=$(echo ${host_osds[$host]} | wc -w)
        total_osds=$((total_osds + osd_count))
    done
    
    if [[ $total_osds -eq 0 ]]; then
        if [[ ${#TARGET_HOSTS[@]} -gt 0 ]]; then
            print_status $RED "No active OSDs found on specified hosts: ${TARGET_HOSTS[*]} in '$ROOT_NAME' tree"
            print_status $YELLOW "Available hosts in '$ROOT_NAME' tree:"
            get_tree_osds_with_hosts "$ROOT_NAME" | cut -d',' -f1 | sort -u | sed 's/^/  - /'
        else
            print_status $RED "No active OSDs found in '$ROOT_NAME' tree"
            print_status $YELLOW "Available trees:"
            ceph osd tree | grep "root" | awk '{print "  - " $4}'
        fi
        exit 1
    fi
    
    if [[ ${#TARGET_HOSTS[@]} -gt 0 ]]; then
        print_status $GREEN "Found $total_osds active OSDs across ${#host_osds[@]} specified hosts in '$ROOT_NAME' tree"
    else
        print_status $GREEN "Found $total_osds active OSDs across ${#host_osds[@]} hosts in '$ROOT_NAME' tree"
    fi
    echo
    
    # Check compression settings for each host and their OSDs
    local config_only_osds=()
    local total_osds_checked=0
    
    # Sort hosts alphabetically for consistent output
    local sorted_hosts=($(printf '%s\n' "${!host_osds[@]}" | sort))
    
    for host in "${sorted_hosts[@]}"; do
        local osd_list=(${host_osds[$host]})
        
        # Print host header
        echo
        print_status $BLUE "=== Host: $host (${#osd_list[@]} OSDs) ==="
        printf "%-8s %-20s %-25s %-12s %s\n" "OSD ID" "Compression Mode" "Compression Algorithm" "Source" "Status"
        printf "%-8s %-20s %-25s %-12s %s\n" "------" "----------------" "---------------------" "------" "------"
        
        # Process each OSD on this host
        for osd_id in "${osd_list[@]}"; do
            result=$(check_osd_compression $osd_id)
            IFS=',' read -r osd comp_mode comp_algo source <<< "$result"
            
            # Track OSDs where we only got config data
            if [[ "$source" == "config" ]]; then
                config_only_osds+=($osd_id)
            fi
            
            # Color code based on compression status and source
            local status_color=$GREEN
            local status="Active"
            
            if [[ "$comp_mode" == "none" ]]; then
                status_color=$YELLOW
                status="No Compression"
            elif [[ "$source" == "config" ]]; then
                status_color=$BLUE
                status="Config Only"
            elif [[ "$comp_mode" == "aggressive" || "$comp_mode" == "force" ]]; then
                status_color=$GREEN
                status="High Compression"
            elif [[ "$comp_mode" == "passive" ]]; then
                status_color=$GREEN  
                status="Passive Compression"
            fi
            
            printf "${status_color}%-8s %-20s %-25s %-12s %s${NC}\n" \
                "osd.$osd" "$comp_mode" "$comp_algo" "$source" "$status"
            
            total_osds_checked=$((total_osds_checked + 1))
        done
    done
    
    echo
    
    # Summary
    print_status $BLUE "=== Summary ==="
    local config_only_count=${#config_only_osds[@]}
    local running_count=$((total_osds_checked - config_only_count))
    
    if [[ ${#TARGET_HOSTS[@]} -gt 0 ]]; then
        echo "Filtered hosts: ${TARGET_HOSTS[*]}"
        echo "Processed hosts: ${#host_osds[@]}"
    else
        echo "Total hosts: ${#host_osds[@]}"
    fi
    echo "Total OSDs checked: $total_osds_checked"
    echo "OSDs with running config: $running_count"
    
    if [[ $config_only_count -gt 0 ]]; then
        print_status $YELLOW "OSDs with config-only data: $config_only_count"
        echo "  (These OSDs were unreachable, showing configured values only)"
        for config_osd in "${config_only_osds[@]}"; do
            echo "  - osd.$config_osd"
        done
    fi
    
    echo
    print_status $GREEN "Compression status check completed!"
}

# Run main function
main "$@"