#!/bin/bash
################################################################################
# JetsonHacks Post-Fix Script
#
# Applies critical kernel module and networking fixes after apt update.
# Ensures Docker + NemoClaw sandbox reliability.
#
# Fixes applied:
# 1. Load br_netfilter kernel module (required for Docker networking)
# 2. Switch iptables to legacy mode (fixes OpenClaw/k3s compatibility)
# 3. Persist both changes across reboots
# 4. Test + rollback on failure
#
# Usage: sudo bash jetson-postfix.sh [--dry-run] [--skip-tests]
# Exit codes: 0 = success, 1 = br_netfilter failed, 2 = iptables failed, 3 = tests failed
################################################################################

set -euo pipefail

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

DRY_RUN="${1:---dry-run=false}"
SKIP_TESTS="${2:---skip-tests=false}"

if [[ "$DRY_RUN" == "--dry-run" ]] || [[ "$DRY_RUN" == "--dry-run=true" ]]; then
    DRY_RUN=true
else
    DRY_RUN=false
fi

if [[ "$SKIP_TESTS" == "--skip-tests" ]] || [[ "$SKIP_TESTS" == "--skip-tests=true" ]]; then
    SKIP_TESTS=true
else
    SKIP_TESTS=false
fi

log_info() { echo -e "${BLUE}[INFO]${NC} $1"; }
log_success() { echo -e "${GREEN}[✓]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[!]${NC} $1"; }
log_error() { echo -e "${RED}[✗]${NC} $1"; }

if [[ $EUID -ne 0 ]]; then
    log_error "This script must be run as root. Use: sudo bash jetson-postfix.sh"
    exit 1
fi

echo ""
echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"
echo -e "${BLUE}JetsonHacks Post-Fix: br_netfilter + iptables-legacy${NC}"
echo -e "${CYAN}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}"

if [[ "$DRY_RUN" == "true" ]]; then
    log_warn "DRY-RUN MODE: No actual changes will be made"
fi

# ============================================================================
# Fix 1: Load br_netfilter Kernel Module
# ============================================================================

log_info "Fix 1: Loading br_netfilter kernel module..."

if [[ "$DRY_RUN" == "false" ]]; then
    if modprobe br_netfilter; then
        log_success "br_netfilter loaded"
    else
        log_error "Failed to load br_netfilter"
        exit 1
    fi

    # Persist br_netfilter across reboots
    if ! grep -q "br_netfilter" /etc/modules-load.d/br_netfilter.conf 2>/dev/null; then
        mkdir -p /etc/modules-load.d/
        echo "br_netfilter" > /etc/modules-load.d/br_netfilter.conf
        log_success "br_netfilter persisted to /etc/modules-load.d/br_netfilter.conf"
    else
        log_info "br_netfilter already persisted"
    fi
else
    log_info "[DRY-RUN] Would load: modprobe br_netfilter"
fi

# ============================================================================
# Fix 2: Switch iptables to Legacy Mode
# ============================================================================

log_info "Fix 2: Switching iptables to legacy mode..."

if [[ "$DRY_RUN" == "false" ]]; then
    CURRENT_IPTABLES=$(update-alternatives --query iptables | grep "^Value:" | awk '{print $2}')
    log_info "Current iptables: $CURRENT_IPTABLES"

    if [[ "$CURRENT_IPTABLES" != *"iptables-legacy"* ]]; then
        log_info "Switching to iptables-legacy..."
        if update-alternatives --set iptables /usr/sbin/iptables-legacy 2>/dev/null; then
            log_success "iptables switched to legacy mode"
        else
            log_error "Failed to switch iptables to legacy mode"
            exit 2
        fi
    else
        log_info "iptables already in legacy mode"
    fi

    # Also set ip6tables to legacy for consistency
    CURRENT_IP6TABLES=$(update-alternatives --query ip6tables 2>/dev/null | grep "^Value:" | awk '{print $2}' || echo "")
    if [[ ! "$CURRENT_IP6TABLES" == *"iptables-legacy"* ]]; then
        log_info "Switching to ip6tables-legacy..."
        update-alternatives --set ip6tables /usr/sbin/ip6tables-legacy 2>/dev/null || log_warn "ip6tables-legacy not available"
    fi
else
    log_info "[DRY-RUN] Would set: update-alternatives --set iptables /usr/sbin/iptables-legacy"
fi

# ============================================================================
# Fix 3: Restart Docker if Any Changes Made
# ============================================================================

log_info "Fix 3: Restarting Docker service..."

if [[ "$DRY_RUN" == "false" ]]; then
    if systemctl is-active docker &>/dev/null; then
        log_info "Stopping Docker..."
        systemctl stop docker || log_warn "Docker stop returned exit code $?"
        sleep 2

        log_info "Starting Docker..."
        if systemctl start docker; then
            log_success "Docker restarted"
            sleep 3  # Wait for Docker to fully initialize
        else
            log_error "Failed to restart Docker"
            exit 2
        fi
    else
        log_warn "Docker service not active, skipping restart"
    fi
else
    log_info "[DRY-RUN] Would restart Docker"
fi

# ============================================================================
# Test Suite: Verify Fixes
# ============================================================================

if [[ "$SKIP_TESTS" == "false" ]]; then
    log_info "Running post-fix test suite..."
    echo ""

    TESTS_PASSED=0
    TESTS_FAILED=0

    if [[ "$DRY_RUN" == "false" ]]; then
        # Test 1: br_netfilter is loaded
        log_info "Test 1: Verify br_netfilter is loaded..."
        if lsmod | grep -q br_netfilter; then
            log_success "✓ br_netfilter module loaded"
            ((TESTS_PASSED++))
        else
            log_error "✗ br_netfilter module not loaded"
            ((TESTS_FAILED++))
        fi

        # Test 2: iptables is in legacy mode
        log_info "Test 2: Verify iptables is legacy mode..."
        if ls -la /usr/sbin/iptables | grep -q iptables-legacy; then
            log_success "✓ iptables is in legacy mode"
            ((TESTS_PASSED++))
        else
            CURRENT=$(readlink /usr/sbin/iptables || echo "unknown")
            log_warn "⚠ iptables points to: $CURRENT (may still be nf_tables mode)"
        fi

        # Test 3: Docker service is running
        log_info "Test 3: Verify Docker service is running..."
        if systemctl is-active docker &>/dev/null; then
            log_success "✓ Docker service is running"
            ((TESTS_PASSED++))
        else
            log_error "✗ Docker service not running"
            ((TESTS_FAILED++))
        fi

        # Test 4: iptables DOCKER-USER chain exists
        log_info "Test 4: Verify iptables DOCKER-USER chain..."
        if iptables -L DOCKER-USER &>/dev/null 2>&1; then
            log_success "✓ iptables DOCKER-USER chain exists (Docker can bind)"
            ((TESTS_PASSED++))
        else
            log_warn "⚠ iptables DOCKER-USER chain not found (Docker may have issues)"
        fi

        # Test 5: Docker can start a container
        log_info "Test 5: Docker container test..."
        if timeout 10 docker run --rm hello-world &>/dev/null 2>&1; then
            log_success "✓ Docker can start containers"
            ((TESTS_PASSED++))
        else
            log_error "✗ Docker container test failed"
            ((TESTS_FAILED++))
        fi

        echo ""
        log_info "Test Results: $TESTS_PASSED passed, $TESTS_FAILED failed"

        if [[ $TESTS_FAILED -gt 0 ]]; then
            log_error "Some tests failed. Attempting rollback..."
            log_warn "Reverting iptables to nftables mode..."
            update-alternatives --set iptables /usr/sbin/iptables-nft 2>/dev/null || log_warn "Rollback failed"
            systemctl restart docker 2>/dev/null || log_warn "Docker restart failed"
            exit 3
        else
            log_success "All tests passed!"
        fi
    else
        log_info "[DRY-RUN] Would run test suite"
    fi
else
    log_warn "Tests skipped (--skip-tests)"
fi

echo ""
echo -e "${GREEN}✓ JetsonHacks post-fix completed successfully${NC}"
echo ""