#!/bin/bash
# =====================================================
# SoftSwitch Platform - NFTables Configuration Generator
# =====================================================
# File: generate-nftables.sh
# Author: Rodrigo Cuadra - SoftSwitch LLC
# =====================================================
# Usage: ./generate-nftables.sh [--apply] [--backup]
# =====================================================

set -e

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

# Configuration
NFTABLES_DIR="/etc/nftables.d"
NFTABLES_CONF="/etc/nftables.conf"
BACKUP_DIR="/var/backups/nftables"
GEOIP_DIR="/var/lib/geoip"
VOIPBL_DIR="/var/lib/voipbl"
LOG_FILE="/var/log/SoftSwitch/nftables-generator.log"
ODBC_INI="/etc/odbc.ini"

# Read database config from odbc.ini
read_odbc_config() {
    local dsn="ss_admin"
    
    if [ -f "$ODBC_INI" ]; then
        # Use exact field names to avoid matching wrong lines (e.g., Description vs Database)
        # Also strip carriage returns (\r) in case of Windows line endings
        DB_HOST=$(grep -A 20 "^\[$dsn\]" "$ODBC_INI" | grep "^Server " | cut -d'=' -f2 | tr -d ' \r' | head -1)
        DB_PORT=$(grep -A 20 "^\[$dsn\]" "$ODBC_INI" | grep "^Port " | cut -d'=' -f2 | tr -d ' \r' | head -1)
        DB_USER=$(grep -A 20 "^\[$dsn\]" "$ODBC_INI" | grep "^Username " | cut -d'=' -f2 | tr -d ' \r' | head -1)
        DB_PASS=$(grep -A 20 "^\[$dsn\]" "$ODBC_INI" | grep "^Password " | cut -d'=' -f2 | tr -d ' \r' | head -1)
        DB_NAME=$(grep -A 20 "^\[$dsn\]" "$ODBC_INI" | grep "^Database " | cut -d'=' -f2 | tr -d ' \r' | head -1)
    fi
    
    # Fallback to environment variables or defaults
    DB_HOST="${DB_HOST:-${PGHOST:-localhost}}"
    DB_PORT="${DB_PORT:-${PGPORT:-5432}}"
    DB_USER="${DB_USER:-${PGUSER:-ss_user}}"
    DB_PASS="${DB_PASS:-${PGPASSWORD:-ss2025}}"
    DB_NAME="${DB_NAME:-${PGDATABASE:-ss_admin}}"
    
    # Export PGPASSWORD so psql doesn't ask for password
    export PGPASSWORD="$DB_PASS"
}

# =====================================================
# FUNCTIONS
# =====================================================

log() {
    echo -e "${BLUE}[INFO]${NC} $1"
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] [INFO] $1" >> "$LOG_FILE" 2>/dev/null || true
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] [SUCCESS] $1" >> "$LOG_FILE" 2>/dev/null || true
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] [WARNING] $1" >> "$LOG_FILE" 2>/dev/null || true
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] [ERROR] $1" >> "$LOG_FILE" 2>/dev/null || true
}

# Check if running as root
check_root() {
    if [ "$EUID" -ne 0 ]; then
        log_error "This script must be run as root"
        exit 1
    fi
}

# Create required directories
setup_directories() {
    mkdir -p "$NFTABLES_DIR"
    mkdir -p "$BACKUP_DIR"
    mkdir -p "$GEOIP_DIR"
    mkdir -p "$VOIPBL_DIR"
    mkdir -p "$(dirname $LOG_FILE)"
}

# Backup current configuration
backup_config() {
    local timestamp=$(date '+%Y%m%d-%H%M%S')
    local backup_file="$BACKUP_DIR/nftables-backup-$timestamp.tar.gz"
    
    log "Creating backup: $backup_file"
    
    if [ -f "$NFTABLES_CONF" ]; then
        tar -czf "$backup_file" "$NFTABLES_CONF" "$NFTABLES_DIR" 2>/dev/null || true
        log_success "Backup created: $backup_file"
    else
        log_warning "No existing configuration to backup"
    fi
}

# Query database for GeoFirewall blocked countries
get_geofirewall_blocked_countries() {
    export PGPASSWORD="$DB_PASS"
    
    psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -t -A -F'|' -c "
        SELECT 
            country_code,
            country_name
        FROM geo_firewall_rules
        WHERE enabled = true AND action = 'block'
        ORDER BY country_code;
    " 2>/dev/null
}

# Check if GeoFirewall is enabled
# Firewall rules are server-level, not tenant-level
# GeoFirewall is enabled if there are any blocked countries in the database
is_geofirewall_enabled() {
    export PGPASSWORD="$DB_PASS"
    
    # Check if there are any blocked countries
    local blocked_count=$(psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -t -A -c "
        SELECT COUNT(*) FROM geo_firewall_rules WHERE enabled = true AND action = 'block';
    " 2>/dev/null)
    
    if [ -n "$blocked_count" ] && [ "$blocked_count" -gt 0 ]; then
        echo "true"
    else
        echo "false"
    fi
}

# Download GeoIP database
download_geoip_database() {
    log "Downloading GeoIP database..."
    
    local geoip_url="https://download.db-ip.com/free/dbip-country-lite-$(date +%Y-%m).csv.gz"
    local geoip_file="$GEOIP_DIR/dbip-country-lite.csv"
    
    if curl -sL "$geoip_url" | gunzip > "$geoip_file" 2>/dev/null; then
        log_success "GeoIP database downloaded: $geoip_file"
        return 0
    else
        log_warning "Failed to download GeoIP database"
        return 1
    fi
}

# Query database for firewall rules
get_firewall_rules() {
    export PGPASSWORD="$DB_PASS"
    
    local query_result=$(psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -t -A -F'|' -c "
        SELECT 
            r.id,
            r.name,
            r.description,
            r.action,
            r.direction,
            r.priority,
            COALESCE(r.source_address, '') as source_address,
            COALESCE(r.destination_address, '') as destination_address,
            COALESCE(r.source_port, '') as source_port,
            COALESCE(r.destination_port, '') as destination_port,
            COALESCE(r.interface_name, '') as interface_name,
            r.enabled,
            COALESCE(s.protocol, 'tcp') as protocol,
            COALESCE(s.port, r.destination_port, '') as port,
            COALESCE(s.name, '') as service_name
        FROM firewall_rules r
        LEFT JOIN firewall_services s ON r.service_id = s.id
        WHERE r.enabled = true
        ORDER BY r.priority ASC, r.id ASC;
    " 2>&1)
    
    local psql_exit_code=$?
    
    if [ $psql_exit_code -ne 0 ]; then
        log_error "Failed to query firewall rules from database: $query_result"
        echo ""
        return 1
    fi
    
    echo "$query_result"
}

# Query database for access control entries (whitelist/blacklist)
get_access_control_rules() {
    export PGPASSWORD="$DB_PASS"
    
    # First check if table exists
    local table_exists=$(psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -t -A -c "
        SELECT EXISTS (
            SELECT FROM information_schema.tables 
            WHERE table_schema = 'public' 
            AND table_name = 'access_control'
        );
    " 2>&1)
    
    if [ $? -ne 0 ] || [ "$table_exists" != "t" ]; then
        # Table doesn't exist yet, return empty
        log "Access control table not found - skipping access control rules"
        echo ""
        return 0
    fi
    
    local query_result=$(psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -t -A -F'|' -c "
        SELECT 
            a.id,
            a.name,
            a.description,
            CASE 
                WHEN a.list_type = 'whitelist' THEN 'accept'
                WHEN a.list_type = 'blacklist' THEN 'drop'
                ELSE 'drop'
            END as action,
            a.direction,
            a.priority,
            CASE 
                WHEN a.direction = 'input' THEN a.ip_address
                ELSE ''
            END as source_address,
            CASE 
                WHEN a.direction = 'output' THEN a.ip_address
                ELSE ''
            END as destination_address,
            COALESCE(a.source_port, '') as source_port,
            COALESCE(a.destination_port, '') as destination_port,
            COALESCE(a.interface_name, '') as interface_name,
            a.enabled,
            COALESCE(a.protocol, 'all') as protocol,
            '' as port,
            a.list_type as service_name
        FROM access_control a
        WHERE a.enabled = true
        AND a.direction = 'input'
        ORDER BY a.priority ASC, a.id ASC;
    " 2>&1)
    
    local psql_exit_code=$?
    
    if [ $psql_exit_code -ne 0 ]; then
        log_warning "Failed to query access control rules from database: $query_result"
        echo ""
        return 1
    fi
    
    echo "$query_result"
}

# Generate firewall rules file from database
generate_firewall_rules() {
    log "Generating firewall rules from database..."
    
    local firewall_rules_file="$NFTABLES_DIR/98-firewall-rules.nft"
    
    # Get firewall rules from database
    local firewall_rules=$(get_firewall_rules)
    local get_rules_exit_code=$?
    
    # Get access control rules from database
    local access_control_rules=$(get_access_control_rules)
    local get_access_control_exit_code=$?
    
    if [ $get_rules_exit_code -ne 0 ]; then
        log_error "Failed to retrieve firewall rules from database"
        # Create empty file to prevent errors
        cat > "$firewall_rules_file" << EOF
# =====================================================
# SoftSwitch Platform - Firewall Rules (Error)
# =====================================================
# Auto-generated: $(date '+%Y-%m-%d %H:%M:%S')
# Error: Failed to retrieve firewall rules from database
# =====================================================
# Please check database connection and try again
EOF
        return 1
    fi
    
    # Check if we have any rules (filter out empty lines and error messages)
    local clean_rules=$(echo "$firewall_rules" | grep -v "^$" | grep -v "ERROR\|error\|Error" || true)
    
    # Add access control rules if available
    if [ $get_access_control_exit_code -eq 0 ] && [ -n "$access_control_rules" ]; then
        local clean_access_rules=$(echo "$access_control_rules" | grep -v "^$" | grep -v "ERROR\|error\|Error" || true)
        if [ -n "$clean_access_rules" ]; then
            clean_rules=$(echo -e "$clean_rules\n$clean_access_rules")
            log "Including access control rules (whitelist/blacklist)"
        fi
    fi
    
    if [ -z "$clean_rules" ]; then
        log "No firewall rules found - creating empty file"
        cat > "$firewall_rules_file" << EOF
# =====================================================
# SoftSwitch Platform - Firewall Rules (Empty)
# =====================================================
# Auto-generated: $(date '+%Y-%m-%d %H:%M:%S')
# No firewall rules configured in database
# =====================================================
# Add rules via UI to populate this file
EOF
        log_success "Empty firewall rules file created: $firewall_rules_file"
        return 0
    fi
    
    # Count rules (count lines with pipe separator)
    local rule_count=$(echo "$clean_rules" | grep -c '|' || echo "0")
    log "Processing $rule_count firewall rules..."
    
    # Generate the rules file
    cat > "$firewall_rules_file" << EOF
# =====================================================
# SoftSwitch Platform - Firewall Rules
# =====================================================
# Auto-generated: $(date '+%Y-%m-%d %H:%M:%S')
# Rules count: $rule_count
# =====================================================
# WARNING: This file is auto-generated - DO NOT EDIT MANUALLY
# To modify rules, use the SoftSwitch Admin UI
# =====================================================
# These rules are added to the existing 'input' chain
# =====================================================

EOF

    # Process each rule
    local processed_count=0
    while IFS='|' read -r id name description action direction priority source_addr dest_addr source_port dest_port interface enabled protocol port service_name; do
        # Skip empty lines
        [ -z "$id" ] && continue
        
        # Skip if rule is disabled (PostgreSQL returns 't' for true, 'f' for false)
        if [ "$enabled" != "t" ] && [ "$enabled" != "true" ] && [ "$enabled" != "1" ]; then
            log "Skipping disabled rule: $name"
            continue
        fi
        
        # Only process input direction rules (output/forward can be added later)
        if [ "$direction" != "input" ]; then
            log "Skipping non-input rule: $name (direction: $direction)"
            continue
        fi
        
        processed_count=$((processed_count + 1))
        
        # Build comment
        local comment="Rule: $name (Priority: $priority)"
        if [ -n "$description" ] && [ "$description" != "NULL" ]; then
            comment="$comment - $description"
        fi
        
        # Build rule components
        local rule_parts=""
        
        # Add interface if specified
        if [ -n "$interface" ] && [ "$interface" != "NULL" ] && [ "$interface" != "" ]; then
            rule_parts="$rule_parts iifname \"$interface\""
        fi
        
        # Add source address if specified
        if [ -n "$source_addr" ] && [ "$source_addr" != "NULL" ] && [ "$source_addr" != "" ]; then
            rule_parts="$rule_parts ip saddr $source_addr"
        fi
        
        # Add destination address if specified
        if [ -n "$dest_addr" ] && [ "$dest_addr" != "NULL" ] && [ "$dest_addr" != "" ]; then
            rule_parts="$rule_parts ip daddr $dest_addr"
        fi
        
        # Add protocol and port
        if [ -n "$protocol" ] && [ "$protocol" != "NULL" ] && [ "$protocol" != "all" ]; then
            if [ "$protocol" = "tcp/udp" ]; then
                # Handle tcp/udp - need to create two rules
                if [ -n "$port" ] && [ "$port" != "NULL" ] && [ "$port" != "" ]; then
                    # Remove leading space from rule_parts if it exists
                    rule_parts=$(echo "$rule_parts" | sed 's/^[[:space:]]*//')
                    # Check if port is a range
                    if [[ "$port" =~ ^[0-9]+-[0-9]+$ ]]; then
                        if [ -n "$rule_parts" ]; then
                            echo "add rule inet filter input $rule_parts tcp dport { $port } $action comment \"$comment (TCP)\"" >> "$firewall_rules_file"
                            echo "add rule inet filter input $rule_parts udp dport { $port } $action comment \"$comment (UDP)\"" >> "$firewall_rules_file"
                        else
                            echo "add rule inet filter input tcp dport { $port } $action comment \"$comment (TCP)\"" >> "$firewall_rules_file"
                            echo "add rule inet filter input udp dport { $port } $action comment \"$comment (UDP)\"" >> "$firewall_rules_file"
                        fi
                    else
                        if [ -n "$rule_parts" ]; then
                            echo "add rule inet filter input $rule_parts tcp dport $port $action comment \"$comment (TCP)\"" >> "$firewall_rules_file"
                            echo "add rule inet filter input $rule_parts udp dport $port $action comment \"$comment (UDP)\"" >> "$firewall_rules_file"
                        else
                            echo "add rule inet filter input tcp dport $port $action comment \"$comment (TCP)\"" >> "$firewall_rules_file"
                            echo "add rule inet filter input udp dport $port $action comment \"$comment (UDP)\"" >> "$firewall_rules_file"
                        fi
                    fi
                else
                    log_warning "Rule $name: tcp/udp protocol without port specified, skipping"
                    continue
                fi
            elif [ "$protocol" = "icmp" ] || [ "$protocol" = "icmpv6" ]; then
                # Remove leading space from rule_parts if it exists
                rule_parts=$(echo "$rule_parts" | sed 's/^[[:space:]]*//')
                if [ "$protocol" = "icmp" ]; then
                    if [ -n "$rule_parts" ]; then
                        echo "add rule inet filter input $rule_parts ip protocol icmp $action comment \"$comment\"" >> "$firewall_rules_file"
                    else
                        echo "add rule inet filter input ip protocol icmp $action comment \"$comment\"" >> "$firewall_rules_file"
                    fi
                else
                    if [ -n "$rule_parts" ]; then
                        echo "add rule inet filter input $rule_parts ip6 nexthdr icmpv6 $action comment \"$comment\"" >> "$firewall_rules_file"
                    else
                        echo "add rule inet filter input ip6 nexthdr icmpv6 $action comment \"$comment\"" >> "$firewall_rules_file"
                    fi
                fi
            else
                # TCP or UDP
                if [ -n "$port" ] && [ "$port" != "NULL" ] && [ "$port" != "" ]; then
                    # Check if port is a range
                    if [[ "$port" =~ ^[0-9]+-[0-9]+$ ]]; then
                        # Remove leading space from rule_parts if it exists
                        rule_parts=$(echo "$rule_parts" | sed 's/^[[:space:]]*//')
                        if [ -n "$rule_parts" ]; then
                            echo "add rule inet filter input $rule_parts $protocol dport { $port } $action comment \"$comment\"" >> "$firewall_rules_file"
                        else
                            echo "add rule inet filter input $protocol dport { $port } $action comment \"$comment\"" >> "$firewall_rules_file"
                        fi
                    else
                        # Remove leading space from rule_parts if it exists
                        rule_parts=$(echo "$rule_parts" | sed 's/^[[:space:]]*//')
                        if [ -n "$rule_parts" ]; then
                            echo "add rule inet filter input $rule_parts $protocol dport $port $action comment \"$comment\"" >> "$firewall_rules_file"
                        else
                            echo "add rule inet filter input $protocol dport $port $action comment \"$comment\"" >> "$firewall_rules_file"
                        fi
                    fi
                else
                    # No port specified - this is invalid for TCP/UDP, skip or use default
                    log_warning "Rule $name: TCP/UDP protocol without port specified, skipping"
                    continue
                fi
            fi
        else
            # Protocol is "all" or not specified - this should not happen for port-based rules
            # Remove leading space from rule_parts if it exists
            rule_parts=$(echo "$rule_parts" | sed 's/^[[:space:]]*//')
            if [ -n "$rule_parts" ]; then
                echo "add rule inet filter input $rule_parts $action comment \"$comment\"" >> "$firewall_rules_file"
            else
                log_warning "Rule $name: No protocol or port specified, skipping"
                continue
            fi
        fi
        
    done <<< "$clean_rules"
    
    log "Processed $processed_count firewall rules"
    log_success "Firewall rules generated: $firewall_rules_file"
}

# Generate GeoIP data file with blocked country IPs
generate_geoip_data() {
    log "Generating GeoIP data..."
    
    local blocked_countries=$(get_geofirewall_blocked_countries)
    local geoip_csv="$GEOIP_DIR/dbip-country-lite.csv"
    local geoip_data_file="$NFTABLES_DIR/geoip-data.nft"
    
    if [ -z "$blocked_countries" ]; then
        log "No countries blocked - creating empty GeoIP data file"
        cat > "$geoip_data_file" << EOF
# =====================================================
# SoftSwitch Platform - GeoIP Data (Empty)
# =====================================================
# No countries are blocked in GeoFirewall
# Generated: $(date '+%Y-%m-%d %H:%M:%S')
# =====================================================
EOF
        return
    fi
    
    # Check if GeoIP database exists
    if [ ! -f "$geoip_csv" ]; then
        log_warning "GeoIP database not found. Run: r2a-nftables --geoip-update"
        cat > "$geoip_data_file" << EOF
# =====================================================
# SoftSwitch Platform - GeoIP Data
# =====================================================
# GeoIP database not downloaded yet
# Run: r2a-nftables --geoip-update
# =====================================================
EOF
        return
    fi
    
    # Count blocked countries
    local country_count=$(echo "$blocked_countries" | grep -c '|' || echo "0")
    
    log "Processing $country_count blocked countries..."
    
    # Generate the data file with nft commands to populate sets
    cat > "$geoip_data_file" << EOF
# =====================================================
# SoftSwitch Platform - GeoIP Data
# =====================================================
# Auto-generated: $(date '+%Y-%m-%d %H:%M:%S')
# Blocked countries: $country_count
# =====================================================
# To apply: nft -f /etc/nftables.conf
# =====================================================

EOF

    # Process each blocked country
    while IFS='|' read -r country_code country_name; do
        [ -z "$country_code" ] && continue
        
        echo "# Country: $country_name ($country_code)" >> "$geoip_data_file"
        
        # Extract IPv4 ranges for this country from CSV
        # CSV format: start_ip,end_ip,country_code
        grep ",$country_code\$" "$geoip_csv" 2>/dev/null | head -500 | while IFS=',' read -r start_ip end_ip cc; do
            # Skip IPv6 (contains :)
            if [[ ! "$start_ip" =~ ":" ]]; then
                echo "add element inet filter geoip_blocked_v4 { $start_ip-$end_ip }" >> "$geoip_data_file"
            fi
        done
        
    done <<< "$blocked_countries"
    
    log_success "GeoIP data generated: $geoip_data_file"
}

# Check if VoIPBL is enabled
is_voipbl_enabled() {
    export PGPASSWORD="$DB_PASS"
    
    local enabled=$(psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -t -A -c "
        SELECT enabled FROM voipbl_settings LIMIT 1;
    " 2>/dev/null)
    
    if [ "$enabled" = "t" ] || [ "$enabled" = "true" ]; then
        echo "true"
    else
        echo "false"
    fi
}

# Download VoIPBL database
download_voipbl_database() {
    log "Downloading VoIPBL database..."
    
    mkdir -p "$VOIPBL_DIR"
    local voipbl_url="http://www.voipbl.org/update/"
    local voipbl_file="$VOIPBL_DIR/voipbl.txt"
    
    if curl -sL "$voipbl_url" > "$voipbl_file.tmp" 2>/dev/null; then
        grep -v "^#" "$voipbl_file.tmp" | grep -v "^$" > "$voipbl_file"
        rm -f "$voipbl_file.tmp"
        log_success "VoIPBL database downloaded: $voipbl_file"
        return 0
    else
        log_warning "Failed to download VoIPBL database"
        return 1
    fi
}

# Generate VoIPBL data file
generate_voipbl_data() {
    log "Generating VoIPBL data..."
    
    local voipbl_file="$VOIPBL_DIR/voipbl.txt"
    local voipbl_data_file="$NFTABLES_DIR/voipbl-data.nft"
    
    local voipbl_enabled=$(is_voipbl_enabled)
    if [ "$voipbl_enabled" != "true" ]; then
        log "VoIPBL disabled - creating empty VoIPBL data file"
        cat > "$voipbl_data_file" << EOF
# =====================================================
# SoftSwitch Platform - VoIPBL Data (Disabled)
# =====================================================
# VoIPBL protection is currently disabled
# Generated: $(date '+%Y-%m-%d %H:%M:%S')
# =====================================================
EOF
        return
    fi
    
    if [ ! -f "$voipbl_file" ]; then
        log_warning "VoIPBL list not found. Run: r2a-nftables --voipbl-update"
        cat > "$voipbl_data_file" << EOF
# =====================================================
# SoftSwitch Platform - VoIPBL Data
# =====================================================
# VoIPBL list not downloaded yet
# Run: r2a-nftables --voipbl-update
# Generated: $(date '+%Y-%m-%d %H:%M:%S')
# =====================================================
EOF
        return
    fi
    
    local ip_count=$(wc -l < "$voipbl_file")
    log "Processing $ip_count VoIPBL IPs/Subnets..."
    
    cat > "$voipbl_data_file" << EOF
# =====================================================
# SoftSwitch Platform - VoIPBL Data
# =====================================================
# Auto-generated by r2a-nftables --voipbl-update / --generate
# Total entries: $ip_count
# Generated: $(date '+%Y-%m-%d %H:%M:%S')
# =====================================================

table inet filter {
    set voipbl_blocked_v4 {
        type ipv4_addr; flags interval;
        elements = {
EOF

    awk '{print "            " $0 ","}' "$voipbl_file" >> "$voipbl_data_file"

    cat >> "$voipbl_data_file" << EOF
        }
    }
}

insert rule inet filter input ip saddr @voipbl_blocked_v4 counter drop comment "VoIPBL Blocked IPs"
EOF
    
    log_success "VoIPBL data generated: $voipbl_data_file"
}

# Generate configuration
generate_config() {
    log "Generating nftables configuration from database..."
    
    # Ensure base /etc/nftables.conf file exists
    if [ ! -f "$NFTABLES_CONF" ] || ! grep -q "table inet filter" "$NFTABLES_CONF"; then
        cat > "$NFTABLES_CONF" << EOF
#!/usr/sbin/nft -f

flush ruleset

table inet filter {
	chain input {
		type filter hook input priority filter; policy accept;
	}
	chain forward {
		type filter hook forward priority filter; policy accept;
	}
	chain output {
		type filter hook output priority filter; policy accept;
	}
}
EOF
    fi

    # Generate firewall rules from database
    generate_firewall_rules
    
    # Generate GeoIP data if enabled
    local geoip_enabled=$(is_geofirewall_enabled)
    if [ "$geoip_enabled" = "true" ]; then
        generate_geoip_data
    else
        log "GeoFirewall disabled, skipping GeoIP generation"
    fi

    # Generate VoIPBL data
    generate_voipbl_data
    
    log_success "Configuration generated successfully"
}

# Validate configuration
validate_config() {
    log "Validating nftables configuration..."
    
    if nft -c -f "$NFTABLES_CONF" 2>&1; then
        log_success "Configuration is valid"
        return 0
    else
        log_error "Configuration validation failed"
        return 1
    fi
}

# Apply configuration
apply_config() {
    log "Applying nftables configuration..."
    
    if nft -f "$NFTABLES_CONF"; then
        log_success "Configuration applied successfully"
        
        # Apply firewall rules file separately (add rule commands need chain to exist)
        if [ -f "$NFTABLES_DIR/98-firewall-rules.nft" ]; then
            log "Applying dynamic firewall rules..."
            if nft -f "$NFTABLES_DIR/98-firewall-rules.nft" 2>/dev/null; then
                log_success "Dynamic firewall rules applied"
            else
                log_warning "Failed to apply dynamic firewall rules"
            fi
        fi

        # Apply VoIPBL data if exists
        if [ -f "$NFTABLES_DIR/voipbl-data.nft" ]; then
            log "Loading VoIPBL data..."
            if nft -f "$NFTABLES_DIR/voipbl-data.nft" 2>/dev/null; then
                log_success "VoIPBL data loaded successfully"
            else
                log_warning "Failed to load VoIPBL data"
            fi
        fi
        
        # Apply GeoIP data if exists
        if [ -f "$NFTABLES_DIR/geoip-data.nft" ]; then
            log "Loading GeoIP data..."
            if nft -f "$NFTABLES_DIR/geoip-data.nft" 2>/dev/null; then
                log_success "GeoIP data loaded"
            else
                log_warning "GeoIP data file is empty or has no elements"
            fi
        fi
        
        return 0
    else
        log_error "Failed to apply configuration"
        return 1
    fi
}

# Reload nftables service
reload_service() {
    log "Reloading nftables service..."
    
    if systemctl reload nftables 2>/dev/null || service nftables reload 2>/dev/null; then
        log_success "nftables service reloaded"
        return 0
    else
        log_warning "Could not reload nftables service"
        return 1
    fi
}

# Show help
show_help() {
    cat << EOF
SoftSwitch NFTables Configuration Generator

Usage: $0 [OPTIONS]

Options:
    --generate      Generate configuration from database (default)
    --apply         Apply configuration after generation
    --backup        Create backup before changes
    --validate      Validate configuration without applying
    --reload        Reload nftables service
    --geoip-update  Update GeoIP database
    --status        Show current firewall status
    --help          Show this help message

Examples:
    $0 --generate --backup --apply
    $0 --geoip-update --generate --apply
    $0 --status

EOF
}

# Show status
show_status() {
    log "Current nftables status:"
    echo ""
    echo "=== Tables ==="
    nft list tables 2>/dev/null || echo "No tables configured"
    echo ""
    echo "=== GeoIP Sets ==="
    nft list set inet filter geoip_blocked_v4 2>/dev/null | head -20 || echo "GeoIP set not found"
    echo ""
    echo "=== Input Chain Rules (first 30) ==="
    nft list chain inet filter input 2>/dev/null | head -30 || echo "Input chain not found"
}

# =====================================================
# MAIN
# =====================================================

main() {
    local do_generate=false
    local do_apply=false
    local do_backup=false
    local do_validate=false
    local do_reload=false
    local do_geoip_update=false
    local do_voipbl_update=false
    local do_status=false
    
    # Parse arguments
    if [ $# -eq 0 ]; then
        do_generate=true
    fi
    
    while [ $# -gt 0 ]; do
        case "$1" in
            --generate)
                do_generate=true
                ;;
            --apply)
                do_apply=true
                ;;
            --backup)
                do_backup=true
                ;;
            --validate)
                do_validate=true
                ;;
            --reload)
                do_reload=true
                ;;
            --geoip-update)
                do_geoip_update=true
                ;;
            --voipbl-update)
                do_voipbl_update=true
                ;;
            --status)
                do_status=true
                ;;
            --help|-h)
                show_help
                exit 0
                ;;
            *)
                log_error "Unknown option: $1"
                show_help
                exit 1
                ;;
        esac
        shift
    done
    
    # Check root
    check_root
    
    # Read database configuration from odbc.ini
    read_odbc_config
    
    # Setup directories
    setup_directories
    
    echo "=========================================="
    echo " SoftSwitch NFTables Generator"
    echo "=========================================="
    
    # Execute actions
    if [ "$do_status" = true ]; then
        show_status
        exit 0
    fi
    
    if [ "$do_backup" = true ]; then
        backup_config
    fi
    
    if [ "$do_geoip_update" = true ]; then
        download_geoip_database
    fi
    
    if [ "$do_voipbl_update" = true ]; then
        download_voipbl_database
        generate_voipbl_data
    fi
    
    if [ "$do_generate" = true ]; then
        generate_config
    fi
    
    if [ "$do_validate" = true ]; then
        validate_config || exit 1
    fi
    
    if [ "$do_apply" = true ]; then
        if validate_config; then
            apply_config
        fi
    fi
    
    if [ "$do_reload" = true ]; then
        reload_service
    fi
    
    echo "=========================================="
    log_success "Done!"
}

# Run main function
main "$@"
