#!/bin/bash

set -euo pipefail

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

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

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

# Check if running as root or with sudo
check_privileges() {
    if [[ $EUID -ne 0 ]] && ! sudo -n true 2>/dev/null; then
        error_exit "This script requires root privileges or passwordless sudo access"
    fi
}

# Check if SELinux is enabled
check_selinux() {
    if ! command -v getenforce >/dev/null 2>&1; then
        error_exit "SELinux tools not found. Is SELinux installed?"
    fi
    
    local selinux_status
    selinux_status=$(getenforce 2>/dev/null || echo "Disabled")
    
    if [[ "$selinux_status" == "Disabled" ]]; then
        error_exit "SELinux is disabled. This script requires SELinux to be enabled."
    fi
    
    log "SELinux status: $selinux_status"
}

# Check if required tools are available
check_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 tools: ${missing_tools[*]}. Please install policycoreutils-python-utils or equivalent package."
    fi
}

# 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
}

# Create completion flag file
create_completion_flag() {
    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:
# - /var/lib/k0s/bin/containerd.* (container_runtime_exec_t)
# - /var/lib/k0s/bin/runc (container_runtime_exec_t)
# - /var/lib/k0s/containerd directory tree (container_var_lib_t)
# - /var/lib/k0s/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: $(readlink -f "$0" 2>/dev/null || echo "$0")
EOF
    
    if [[ $? -eq 0 ]]; then
        log "SUCCESS: Created completion flag at $COMPLETION_FLAG"
    else
        log "WARNING: Failed to create completion flag at $COMPLETION_FLAG"
    fi
}
check_data_dir() {
    if [[ ! -d "$DATA_DIR" ]]; then
        error_exit "Data directory $DATA_DIR does not exist"
    fi
    log "Data directory $DATA_DIR exists"
}

# Check if fcontext rule already exists
fcontext_exists() {
    local pattern="$1"
    local context="$2"
    
    sudo semanage fcontext -l | grep -q "^${pattern}.*${context}" 2>/dev/null
}

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

# Restore SELinux contexts
restore_contexts() {
    local path="$1"
    local description="$2"
    
    if [[ ! -e "$path" ]]; then
        log "SKIP: $description - path $path does not exist"
        return 0
    fi
    
    log "RESTORING: $description"
    if sudo restorecon -R -v "$path" 2>&1 | while read -r line; do
        [[ -n "$line" ]] && log "  $line"
    done; then
        log "SUCCESS: Restored contexts for $path"
        return 0
    else
        error_exit "Failed to restore contexts for $path"
    fi
}

# Verify that contexts were applied correctly
verify_contexts() {
    local path="$1"
    local expected_pattern="$2"
    local description="$3"
    
    if [[ ! -e "$path" ]]; then
        log "SKIP VERIFY: $description - path $path does not exist"
        return 0
    fi
    
    log "VERIFYING: $description"
    
    # Get the actual context
    local actual_contexts
    actual_contexts=$(find "$path" -exec ls -Z {} \; 2>/dev/null | awk '{print $1}' | sort -u)
    
    if echo "$actual_contexts" | grep -q "$expected_pattern"; then
        log "SUCCESS: Verified contexts for $path contain $expected_pattern"
        return 0
    else
        log "WARNING: Expected pattern $expected_pattern not found in contexts for $path"
        log "Actual contexts found: $actual_contexts"
        return 1
    fi
}

# Main execution
main() {
    log "Starting $SCRIPT_NAME"
    
    # Check if already completed
    check_completion_flag
    
    # Pre-flight checks
    check_privileges
    check_selinux
    check_tools
    check_data_dir
    
    # Add file context rules (idempotent)
    add_fcontext "${DATA_DIR}/bin/containerd.*" "container_runtime_exec_t" "containerd executables"
    add_fcontext "${DATA_DIR}/bin/runc" "container_runtime_exec_t" "runc executable"
    add_fcontext "${DATA_DIR}/containerd(/.*)?", "container_var_lib_t" "containerd directory"
    add_fcontext "${DATA_DIR}/containerd/io.containerd.snapshotter.*/snapshots(/.*)?" "container_ro_file_t" "containerd snapshots"
    
    # Restore contexts
    restore_contexts "${DATA_DIR}/bin" "k0s binaries"
    restore_contexts "${DATA_DIR}/containerd" "containerd directory"
    
    # Verify contexts were applied (optional verification)
    local verification_failed=0
    
    if [[ -d "${DATA_DIR}/bin" ]]; then
        verify_contexts "${DATA_DIR}/bin" "container_runtime_exec_t" "k0s binaries" || verification_failed=1
    fi
    
    if [[ -d "${DATA_DIR}/containerd" ]]; then
        verify_contexts "${DATA_DIR}/containerd" "container_" "containerd directory" || verification_failed=1
    fi
    
    if [[ $verification_failed -eq 1 ]]; then
        log "WARNING: Some context verifications failed. Check the logs above."
        log "You may want to run 'sudo restorecon -R -v $DATA_DIR' manually."
    fi
    
    # Create completion flag to prevent future runs
    create_completion_flag
    
    log "Completed $SCRIPT_NAME successfully"
}

# Run main function
main "$@"