#!/bin/bash

set -euo pipefail

# Configuration
readonly DATA_DIR="/var/lib/k0s"
readonly SCRIPT_NAME="$(basename "$0")"
readonly COMPLETION_FLAG="$HOME/.k0s-selinuxsetup-complete"

# Logging function
log() {
    printf "[%s] %s\n" "$(date '+%Y-%m-%d %H:%M:%S')" "$*" >&2
}

# Error handling function
error_exit() {
    log "ERROR: $1"
    exit 1
}

# Check if script has already been run successfully
check_completion_flag() {
    if [[ -f "$COMPLETION_FLAG" ]]; then
        log "SKIP: SELinux setup has already been completed successfully"
        log "Completion flag found at: $COMPLETION_FLAG"
        log "If you need to re-run this setup, remove the flag file and run again:"
        log "  rm '$COMPLETION_FLAG'"
        exit 0
    fi
}

# Check if running with sufficient privileges
check_privileges() {
    if ! sudo -n true 2>/dev/null; then
        error_exit "This script requires passwordless sudo access"
    fi
}

# Check if SELinux is available and enabled
check_selinux() {
    if ! command -v getenforce >/dev/null 2>&1; then
        error_exit "getenforce command not found - SELinux may not be installed"
    fi
    
    local selinux_status
    if ! selinux_status=$(getenforce 2>/dev/null); then
        error_exit "Failed to get SELinux status"
    fi
    
    case "$selinux_status" in
        "Enforcing"|"Permissive")
            log "SELinux status: $selinux_status"
            ;;
        *)
            error_exit "SELinux is not enabled (status: $selinux_status)"
            ;;
    esac
}

# Check if required SELinux management tools are available
check_selinux_tools() {
    local missing_tools=()
    
    for tool in semanage restorecon; do
        if ! command -v "$tool" >/dev/null 2>&1; then
            missing_tools+=("$tool")
        fi
    done
    
    if [[ ${#missing_tools[@]} -gt 0 ]]; then
        error_exit "Missing required SELinux tools: ${missing_tools[*]}"
    fi
}

# Check if target directory exists
check_target_directory() {
    if [[ ! -d "$DATA_DIR" ]]; then
        error_exit "Target directory $DATA_DIR does not exist"
    fi
    log "Target directory $DATA_DIR exists"
}

# Check if fcontext rule already exists (safer grep with fixed strings)
fcontext_exists() {
    local pattern="$1"
    local context_type="$2"
    
    # Use semanage to list and grep for exact matches
    if sudo semanage fcontext -l 2>/dev/null | grep -F "$pattern" | grep -q "$context_type"; then
        return 0
    else
        return 1
    fi
}

# Add SELinux file context rule (idempotent)
add_fcontext_rule() {
    local pattern="$1"
    local context_type="$2"
    local description="$3"
    
    if fcontext_exists "$pattern" "$context_type"; then
        log "SKIP: $description - fcontext rule already exists"
        return 0
    fi
    
    log "ADDING: $description ($pattern -> $context_type)"
    
    if sudo semanage fcontext -a -t "$context_type" "$pattern" 2>/dev/null; then
        log "SUCCESS: Added fcontext rule for $pattern"
        return 0
    else
        log "WARNING: Failed to add fcontext rule for $pattern"
        return 1
    fi
}

# Restore SELinux contexts (non-fatal)
restore_contexts() {
    local target_path="$1"
    local description="$2"
    
    if [[ ! -e "$target_path" ]]; then
        log "SKIP: $description - path $target_path does not exist"
        return 0
    fi
    
    log "RESTORING: $description"
    
    # Run restorecon and capture output, but don't fail the script
    if sudo restorecon -R -v "$target_path" 2>&1 | while IFS= read -r line; do
        [[ -n "$line" ]] && log "  restorecon: $line"
    done; then
        log "SUCCESS: Context restoration completed for $target_path"
        return 0
    else
        log "WARNING: restorecon reported issues for $target_path (continuing anyway)"
        return 1
    fi
}

# Show what contexts are actually present (informational only)
show_contexts() {
    local target_path="$1"
    local description="$2"
    
    if [[ ! -e "$target_path" ]]; then
        log "SKIP INFO: $description - path $target_path does not exist"
        return 0
    fi
    
    log "CONTEXTS: $description"
    
    # Show context summary without making assumptions about what's "correct"
    local context_info
    if context_info=$(find "$target_path" -exec ls -Z {} \; 2>/dev/null | awk '{print $1}' | sort | uniq -c 2>/dev/null); then
        if [[ -n "$context_info" ]]; then
            echo "$context_info" | while IFS= read -r line; do
                log "  $line"
            done
        else
            log "  (no context information available)"
        fi
    else
        log "  (unable to retrieve context information)"
    fi
}

# Create completion flag file
create_completion_flag() {
    # Use a here-document with proper quoting
    if cat > "$COMPLETION_FLAG" << EOF
# k0s SELinux Setup Completion Flag
#
# This file indicates that the k0s SELinux configuration script has been
# run successfully. It prevents the script from running multiple times.
#
# The script configures SELinux file contexts for:
# - ${DATA_DIR}/bin/containerd.* (container_runtime_exec_t)
# - ${DATA_DIR}/bin/runc (container_runtime_exec_t)
# - ${DATA_DIR}/containerd directory tree (container_var_lib_t)
# - ${DATA_DIR}/containerd snapshots (container_ro_file_t)
#
# If you remove this file, the SELinux script will run again on the next
# k0sapply execution.
#
# Created: $(date)
# Script: $0
EOF
    then
        log "SUCCESS: Created completion flag at $COMPLETION_FLAG"
    else
        log "WARNING: Failed to create completion flag at $COMPLETION_FLAG"
        return 1
    fi
}

# Main execution function
main() {
    log "Starting $SCRIPT_NAME"
    
    # Early exit if already completed
    check_completion_flag
    
    # System checks
    check_privileges
    check_selinux
    check_selinux_tools
    check_target_directory
    
    # Track any failures
    local rule_failures=0
    local restore_failures=0
    
    # Add SELinux file context rules
    add_fcontext_rule "${DATA_DIR}/bin/containerd.*" "container_runtime_exec_t" "containerd executables" || ((rule_failures++))
    add_fcontext_rule "${DATA_DIR}/bin/runc" "container_runtime_exec_t" "runc executable" || ((rule_failures++))
    add_fcontext_rule "${DATA_DIR}/containerd(/.*)?" "container_var_lib_t" "containerd directory" || ((rule_failures++))
    add_fcontext_rule "${DATA_DIR}/containerd/io.containerd.snapshotter.*/snapshots(/.*)?" "container_ro_file_t" "containerd snapshots" || ((rule_failures++))
    
    # Apply contexts to existing files
    restore_contexts "${DATA_DIR}/bin" "k0s binaries" || ((restore_failures++))
    restore_contexts "${DATA_DIR}/containerd" "containerd directory" || ((restore_failures++))
    
    # Show what contexts are actually present (informational)
    show_contexts "${DATA_DIR}/bin" "k0s binaries"
    show_contexts "${DATA_DIR}/containerd" "containerd directory"
    
    # Report any issues but don't fail
    if [[ $rule_failures -gt 0 ]]; then
        log "WARNING: $rule_failures fcontext rule(s) failed to add"
    fi
    
    if [[ $restore_failures -gt 0 ]]; then
        log "WARNING: $restore_failures context restoration(s) had issues"
    fi
    
    # Create completion flag regardless of warnings
    create_completion_flag
    
    log "Completed $SCRIPT_NAME"
}

# Execute main function
main "$@"