#!/usr/bin/env bash
# Enable protection against thin pools running out of space
# Configures LVM to automatically protect thin pools when they approach capacity

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"

source "$PROJECT_ROOT/lib/common.sh"

# Load configuration (if available)
load_config 2>/dev/null || true

# Default thresholds (can be overridden via config file or environment variables)
THIN_POOL_AUTOEXTEND_THRESHOLD="${THIN_POOL_AUTOEXTEND_THRESHOLD:-80}"  # Warn/extend at 80% usage
THIN_POOL_AUTOEXTEND_PERCENT="${THIN_POOL_AUTOEXTEND_PERCENT:-20}"      # Extend by 20% when threshold reached

# LVM configuration file
LVM_CONF="/etc/lvm/lvm.conf"
LVM_CONF_BACKUP="${LVM_CONF}.backup.$(date +%Y%m%d_%H%M%S)"

# Check prerequisites
check_prerequisites() {
    log_info "Checking prerequisites..."
    
    # Must run as root
    check_root
    
    # Check for LVM commands
    if ! command_exists lvs; then
        error_exit "lvs command not found. LVM tools are required."
    fi
    
    if ! command_exists vgs; then
        error_exit "vgs command not found. LVM tools are required."
    fi
    
    log_success "Prerequisites check passed"
}

# List all thin pools on the system
list_thin_pools() {
    log_info "Scanning for LVM thin pools..."
    
    local thin_pools
    thin_pools=$(lvs --noheadings --select 'lv_attr=~[^t.*]' -o vg_name,lv_name,lv_size,data_percent,metadata_percent 2>/dev/null || true)
    
    if [[ -z "$thin_pools" ]]; then
        log_warn "No LVM thin pools found on this system"
        return 1
    fi
    
    log_info ""
    log_info "=== Found Thin Pools ==="
    echo "VG Name        LV Name              Size      Data%    Metadata%"
    echo "-------------------------------------------------------------------"
    while IFS= read -r line; do
        if [[ -n "$line" ]]; then
            echo "$line"
        fi
    done <<< "$thin_pools"
    
    return 0
}

# Backup LVM configuration
backup_lvm_conf() {
    if [[ -f "$LVM_CONF" ]]; then
        log_info "Creating backup of $LVM_CONF..."
        cp "$LVM_CONF" "$LVM_CONF_BACKUP"
        log_success "Backup created: $LVM_CONF_BACKUP"
    else
        log_warn "LVM configuration file not found at $LVM_CONF"
        log_info "Creating new configuration file..."
        mkdir -p "$(dirname "$LVM_CONF")"
    fi
}

# Configure thin pool protection in lvm.conf
configure_thin_pool_protection() {
    log_info "Configuring thin pool protection in $LVM_CONF..."
    
    backup_lvm_conf
    
    # Create a temporary Python script for reliable configuration editing
    local temp_python_script
    temp_python_script=$(mktemp)
    
    cat > "$temp_python_script" << 'PYEOF'
import sys
import re

lvm_conf = sys.argv[1]
threshold = sys.argv[2]
percent = sys.argv[3]

# Read existing file
try:
    with open(lvm_conf, 'r') as f:
        lines = f.readlines()
except FileNotFoundError:
    lines = []

# Configuration to add/update
new_threshold_line = f"    thin_pool_autoextend_threshold = {threshold}\n"
new_percent_line = f"    thin_pool_autoextend_percent = {percent}\n"
config_comment = "    # Thin pool protection (prevents running out of space)\n"

# Track if we're in activation section and if settings exist
in_activation = False
activation_start = -1
activation_end = -1
has_threshold = False
has_percent = False
brace_count = 0

# First pass: find activation section and check for existing settings
for i, line in enumerate(lines):
    if re.match(r'^\s*activation\s*\{', line):
        in_activation = True
        activation_start = i
        brace_count = 1
    elif in_activation:
        if '{' in line:
            brace_count += line.count('{')
        if '}' in line:
            brace_count -= line.count('}')
            if brace_count == 0:
                activation_end = i
                break
        if 'thin_pool_autoextend_threshold' in line:
            has_threshold = True
        if 'thin_pool_autoextend_percent' in line:
            has_percent = True

# Build new file content
new_lines = []

if activation_start >= 0 and activation_end >= 0:
    # Add lines before activation section
    new_lines.extend(lines[:activation_start])
    
    # Process activation section
    i = activation_start
    while i <= activation_end:
        line = lines[i]
        
        # Replace existing threshold/percent lines
        if 'thin_pool_autoextend_threshold' in line:
            new_lines.append(new_threshold_line)
            # Skip any comment lines that follow
            i += 1
            while i <= activation_end and lines[i].strip().startswith('#'):
                i += 1
            continue
        elif 'thin_pool_autoextend_percent' in line:
            new_lines.append(new_percent_line)
            i += 1
            continue
        elif i == activation_end - 1 and not (has_threshold and has_percent):
            # Before closing brace, add our settings if they don't exist
            new_lines.append(line)
            if not has_threshold or not has_percent:
                new_lines.append(config_comment)
            if not has_threshold:
                new_lines.append(new_threshold_line)
            if not has_percent:
                new_lines.append(new_percent_line)
            i += 1
            continue
        
        new_lines.append(line)
        i += 1
    
    # Add remaining lines after activation section
    new_lines.extend(lines[activation_end + 1:])
else:
    # No activation section, add one at the end
    new_lines.extend(lines)
    if new_lines and not new_lines[-1].endswith('\n'):
        new_lines.append('\n')
    new_lines.append('\n# Thin pool protection (prevents running out of space)\n')
    new_lines.append('activation {\n')
    new_lines.append(config_comment)
    new_lines.append(new_threshold_line)
    new_lines.append(new_percent_line)
    new_lines.append('}\n')

# Write back
with open(lvm_conf, 'w') as f:
    f.writelines(new_lines)

print("OK")
PYEOF
    
    # Try using Python first (more reliable)
    if command_exists python3 && python3 "$temp_python_script" "$LVM_CONF" "$THIN_POOL_AUTOEXTEND_THRESHOLD" "$THIN_POOL_AUTOEXTEND_PERCENT" >/dev/null 2>&1; then
        log_success "Thin pool protection configured"
        rm -f "$temp_python_script"
    else
        # Fallback: simple append approach (LVM merges sections)
        log_info "Using simple configuration method (LVM will merge sections)..."
        
        # Remove old settings if they exist (simple grep -v)
        if [[ -f "$LVM_CONF" ]]; then
            grep -v "thin_pool_autoextend_threshold" "$LVM_CONF" | \
            grep -v "thin_pool_autoextend_percent" > "${LVM_CONF}.tmp" 2>/dev/null || cp "$LVM_CONF" "${LVM_CONF}.tmp"
            mv "${LVM_CONF}.tmp" "$LVM_CONF"
        fi
        
        # Append new configuration (LVM automatically merges activation sections)
        {
            echo ""
            echo "# Thin pool protection (added $(date))"
            echo "# Prevents thin pools from running out of space"
            echo "activation {"
            echo "    thin_pool_autoextend_threshold = ${THIN_POOL_AUTOEXTEND_THRESHOLD}"
            echo "    thin_pool_autoextend_percent = ${THIN_POOL_AUTOEXTEND_PERCENT}"
            echo "}"
        } >> "$LVM_CONF"
        
        log_success "Thin pool protection configured (using fallback method)"
        log_warn "Note: LVM will automatically merge activation sections in lvm.conf"
        rm -f "$temp_python_script"
    fi
    
    log_info "  Threshold: ${THIN_POOL_AUTOEXTEND_THRESHOLD}% (protection triggers when pool reaches this usage)"
    log_info "  Extend by: ${THIN_POOL_AUTOEXTEND_PERCENT}% (extend pool by this amount when threshold is reached)"
}

# Verify configuration
verify_configuration() {
    log_info "Verifying LVM configuration..."
    
    if ! lvm dumpconfig activation/thin_pool_autoextend_threshold >/dev/null 2>&1; then
        log_warn "Could not verify thin_pool_autoextend_threshold via lvm dumpconfig"
        log_info "Checking lvm.conf directly..."
        
        if grep -q "thin_pool_autoextend_threshold[[:space:]]*=[[:space:]]*${THIN_POOL_AUTOEXTEND_THRESHOLD}" "$LVM_CONF"; then
            log_success "Configuration verified in lvm.conf"
        else
            log_warn "Configuration may not be set correctly. Please verify manually."
        fi
    else
        local current_threshold
        current_threshold=$(lvm dumpconfig activation/thin_pool_autoextend_threshold 2>/dev/null | grep -oP '=\s*\K[0-9]+' || echo "")
        
        if [[ -n "$current_threshold" ]]; then
            log_success "Configuration verified: threshold = ${current_threshold}%"
        else
            log_warn "Could not retrieve threshold value, but configuration file has been updated"
        fi
    fi
}

# Show current thin pool status
show_thin_pool_status() {
    log_info ""
    log_info "=== Current Thin Pool Status ==="
    
    list_thin_pools || return 0
    
    log_info ""
    log_info "Protection Status:"
    log_info "  - When a thin pool reaches ${THIN_POOL_AUTOEXTEND_THRESHOLD}% usage, LVM will:"
    log_info "    1. Issue warnings about low space"
    log_info "    2. Attempt to auto-extend the pool by ${THIN_POOL_AUTOEXTEND_PERCENT}% if possible"
    log_info "    3. Prevent new allocations if the pool cannot be extended"
}

# Main function
main() {
    log_info "========================================="
    log_info "Enable Thin Pool Protection"
    log_info "========================================="
    log_info ""
    log_info "This script enables protection against thin pools running out of space"
    log_info "by configuring LVM to automatically manage thin pool capacity."
    log_info ""
    
    check_prerequisites
    
    log_info ""
    list_thin_pools || {
        log_warn "No thin pools found, but configuration will still be applied"
        log_info "Configuration will take effect for any thin pools created in the future"
    }
    
    log_info ""
    configure_thin_pool_protection
    
    log_info ""
    verify_configuration
    
    log_info ""
    show_thin_pool_status
    
    log_info ""
    log_success "Thin pool protection has been enabled!"
    log_info ""
    log_info "Next steps:"
    log_info "  1. Monitor thin pool usage: lvs -o vg_name,lv_name,data_percent,metadata_percent"
    log_info "  2. If needed, manually extend a thin pool: lvextend -L +<size> <vg>/<thin_pool>"
    log_info "  3. Check logs for thin pool warnings: dmesg | grep -i thin"
    log_info ""
    log_info "Configuration backup saved to: $LVM_CONF_BACKUP"
    log_info "To restore previous configuration: cp $LVM_CONF_BACKUP $LVM_CONF"
}

# Run main function
main

