#!/bin/bash

set -euo pipefail

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m'

print_color() {
    local color=$1
    shift
    echo -e "${color}$*${NC}"
}

usage() {
    echo "Usage: $0 [FILTER] [--set VALUE]"
    echo ""
    echo "Options:"
    echo "  --hosts HOST[,HOST2,...]  Filter to specific host(s)"
    echo "  --root ROOT               Filter to OSDs under root"
    echo "  --rack RACK               Filter to OSDs in rack"
    echo "  --tree TREE               Filter to OSDs in tree location"
    echo "  --set VALUE               Set osd_max_backfills to VALUE on filtered OSDs"
    echo "  -h, --help                Show this help message"
    echo ""
    echo "Examples:"
    echo "  $0 --hosts csn12"
    echo "  $0 --root ssd --set 3"
    echo "  $0 --root sata --set 1"
    echo "  $0 --rack rack1 --set 2"
    exit 1
}

check_ceph_cmd() {
    if ! command -v ceph &> /dev/null; then
        print_color $RED "Error: 'ceph' command not found"
        exit 1
    fi
}

get_osd_tree() {
    ceph osd tree --format=json 2>/dev/null
}

get_osds() {
    local filter_type="$1"
    local filter_name="$2"
    local osd_tree_json="$3"

    if [[ "$filter_type" == "none" ]]; then
        echo "$osd_tree_json" | jq -r '.nodes[] | select(.type == "osd") | .id'
    elif [[ "$filter_type" == "hosts" ]]; then
        local IFS=','
        local hosts=($filter_name)
        local all_osds=""

        for host in "${hosts[@]}"; do
            host=$(echo "$host" | xargs)

            local host_id=$(echo "$osd_tree_json" | jq -r ".nodes[] | select(.type == \"host\" and .name == \"$host\") | .id")
            if [[ -z "$host_id" || "$host_id" == "null" ]]; then
                print_color $RED "Error: Host '$host' not found"
                exit 1
            fi

            local host_osds=$(echo "$osd_tree_json" | jq -r ".nodes[] | select(.type == \"host\" and .name == \"$host\") | .children[]?")
            if [[ -n "$host_osds" ]]; then
                if [[ -n "$all_osds" ]]; then
                    all_osds="$all_osds"$'\n'"$host_osds"
                else
                    all_osds="$host_osds"
                fi
            fi
        done

        echo "$all_osds"
    elif [[ "$filter_type" == "root" || "$filter_type" == "rack" || "$filter_type" == "tree" ]]; then
        local tree_id
        if [[ "$filter_type" == "root" ]]; then
            tree_id=$(echo "$osd_tree_json" | jq -r ".nodes[] | select(.type == \"root\" and .name == \"$filter_name\") | .id")
        elif [[ "$filter_type" == "rack" ]]; then
            tree_id=$(echo "$osd_tree_json" | jq -r ".nodes[] | select(.type == \"rack\" and .name == \"$filter_name\") | .id")
        else
            tree_id=$(echo "$osd_tree_json" | jq -r ".nodes[] | select(.name == \"$filter_name\") | .id")
        fi

        if [[ -z "$tree_id" || "$tree_id" == "null" ]]; then
            print_color $RED "Error: Tree location '$filter_name' not found"
            exit 1
        fi

        find_osds_under_node() {
            local node_id=$1
            local json=$2

            echo "$json" | jq -r ".nodes[] | select(.id == $node_id) | .children[]?" | while read -r child_id; do
                if [[ -n "$child_id" && "$child_id" != "null" ]]; then
                    local child_type=$(echo "$json" | jq -r ".nodes[] | select(.id == $child_id) | .type")
                    if [[ "$child_type" == "osd" ]]; then
                        echo "$child_id"
                    else
                        find_osds_under_node "$child_id" "$json"
                    fi
                fi
            done
        }

        find_osds_under_node "$tree_id" "$osd_tree_json"
    fi
}

get_osd_host() {
    local osd_id=$1
    local osd_tree_json=$2

    local host_name=$(echo "$osd_tree_json" | jq -r ".nodes[] | select(.type == \"host\" and (.children[]? // empty) == $osd_id) | .name" 2>/dev/null || echo "")
    echo "${host_name:-unknown}"
}

get_max_backfills() {
    local osd_id=$1

    set +e
    local result=$(ceph config show osd.$osd_id | grep osd_max_backfills | awk '{print $2}' 2>/dev/null)
    local exit_code=$?
    set -e

    if [[ $exit_code -ne 0 || -z "$result" ]]; then
        result="1"
    fi

    echo "$result"
}

backfill_color() {
    local val=$1
    if [[ "$val" =~ ^[0-9]+$ ]]; then
        if (( val >= 3 )); then
            echo "$GREEN"
        elif (( val == 2 )); then
            echo "$YELLOW"
        else
            echo "$RED"
        fi
    else
        echo "$NC"
    fi
}

collect_osds() {
    local filter_type="$1"
    local filter_name="$2"

    local osd_tree_json
    if ! osd_tree_json=$(get_osd_tree); then
        print_color $RED "Error: Failed to get OSD tree"
        exit 1
    fi

    local osds
    if ! osds=$(get_osds "$filter_type" "$filter_name" "$osd_tree_json"); then
        print_color $RED "Error: Failed to get OSD list"
        exit 1
    fi

    if [[ -z "$osds" ]]; then
        print_color $YELLOW "No OSDs found matching the criteria"
        exit 0
    fi

    echo "$osd_tree_json"$'\x1f'"$osds"
}

print_filter_header() {
    local filter_type="$1"
    local filter_name="$2"

    if [[ "$filter_type" == "hosts" ]]; then
        if [[ "$filter_name" == *","* ]]; then
            print_color $BLUE "Hosts: $filter_name"
        else
            print_color $BLUE "Host: $filter_name"
        fi
    elif [[ "$filter_type" == "root" ]]; then
        print_color $BLUE "Root: $filter_name"
    elif [[ "$filter_type" == "rack" ]]; then
        print_color $BLUE "Rack: $filter_name"
    elif [[ "$filter_type" == "tree" ]]; then
        print_color $BLUE "Tree: $filter_name"
    else
        print_color $BLUE "All OSDs"
    fi
    echo
}

display_results() {
    local filter_type="$1"
    local filter_name="$2"

    print_color $CYAN "=== Ceph OSD Max Backfills ==="
    echo
    print_filter_header "$filter_type" "$filter_name"

    local combined
    combined=$(collect_osds "$filter_type" "$filter_name")
    local osd_tree_json="${combined%%$'\x1f'*}"
    local osds="${combined#*$'\x1f'}"

    printf "%-8s %-20s %-15s\n" "OSD" "Host" "Max Backfills"
    printf "%-8s %-20s %-15s\n" "---" "----" "-------------"

    local total_count=0
    declare -A backfill_summary

    while IFS= read -r osd_id; do
        [[ -z "$osd_id" || "$osd_id" == "null" ]] && continue

        local host=$(get_osd_host "$osd_id" "$osd_tree_json")
        local val=$(get_max_backfills "$osd_id")
        local color=$(backfill_color "$val")

        printf "${color}%-8s %-20s %-15s${NC}\n" "osd.$osd_id" "${host:-unknown}" "$val"

        : $((total_count++))
        if [[ -z "${backfill_summary["$val"]+x}" ]]; then
            backfill_summary["$val"]=0
        fi
        : $((backfill_summary["$val"]++))
    done <<< "$osds"

    echo
    print_color $CYAN "=== Summary ==="
    echo "Total OSDs checked: $total_count"
    echo
    print_color $BLUE "Max Backfills Distribution:"
    for val in $(printf '%s\n' "${!backfill_summary[@]}" | sort -n); do
        local count=${backfill_summary[$val]}
        local pct=$(( count * 100 / total_count ))
        printf "  %s: %d OSDs (%d%%)\n" "$val" "$count" "$pct"
    done
    echo
    print_color $BLUE "Legend:"
    print_color $GREEN  "  Green:  >= 3"
    print_color $YELLOW "  Yellow:  = 2"
    print_color $RED    "  Red:    <= 1"
}

apply_results() {
    local filter_type="$1"
    local filter_name="$2"
    local set_value="$3"

    print_color $CYAN "=== Ceph OSD Max Backfills Set ==="
    echo
    print_filter_header "$filter_type" "$filter_name"
    print_color $BLUE "Setting osd_max_backfills to $set_value"
    echo

    local combined
    combined=$(collect_osds "$filter_type" "$filter_name")
    local osd_tree_json="${combined%%$'\x1f'*}"
    local osds="${combined#*$'\x1f'}"

    printf "%-8s %-20s %-10s\n" "OSD" "Host" "Status"
    printf "%-8s %-20s %-10s\n" "---" "----" "------"

    local ok=0
    local fail=0

    while IFS= read -r osd_id; do
        [[ -z "$osd_id" || "$osd_id" == "null" ]] && continue

        local host=$(get_osd_host "$osd_id" "$osd_tree_json")

        set +e
        ceph tell osd.$osd_id injectargs --osd_max_backfills "$set_value" &>/dev/null
        local exit_code=$?
        set -e

        if [[ $exit_code -eq 0 ]]; then
            printf "${GREEN}%-8s %-20s %-10s${NC}\n" "osd.$osd_id" "${host:-unknown}" "OK"
            : $((ok++))
        else
            printf "${RED}%-8s %-20s %-10s${NC}\n" "osd.$osd_id" "${host:-unknown}" "FAIL"
            : $((fail++))
        fi
    done <<< "$osds"

    echo
    print_color $CYAN "=== Summary ==="
    printf "Set to %s:  %d OK, %d failed\n" "$set_value" "$ok" "$fail"
}

main() {
    local filter_type="none"
    local filter_name=""
    local set_value=""

    if [[ $# -eq 0 ]]; then
        usage
    fi

    while [[ $# -gt 0 ]]; do
        case $1 in
            --hosts)
                filter_type="hosts"
                filter_name="$2"
                shift 2
                ;;
            --root)
                filter_type="root"
                filter_name="$2"
                shift 2
                ;;
            --rack)
                filter_type="rack"
                filter_name="$2"
                shift 2
                ;;
            --tree)
                filter_type="tree"
                filter_name="$2"
                shift 2
                ;;
            --set)
                set_value="$2"
                shift 2
                ;;
            -h|--help)
                usage
                ;;
            *)
                print_color $RED "Error: Unknown option '$1'"
                usage
                ;;
        esac
    done

    if [[ "$filter_type" != "none" && -z "$filter_name" ]]; then
        print_color $RED "Error: Filter option requires a value"
        usage
    fi

    if [[ -n "$set_value" && ! "$set_value" =~ ^[0-9]+$ ]]; then
        print_color $RED "Error: --set VALUE must be a non-negative integer"
        exit 1
    fi

    check_ceph_cmd

    if [[ -n "$set_value" ]]; then
        apply_results "$filter_type" "$filter_name" "$set_value"
    else
        display_results "$filter_type" "$filter_name"
    fi
}

main "$@"
