#!/bin/bash

# Personal Internet Cell - Network Configuration Script
# This script sets up proper routing and NAT for WireGuard VPN internet access

set -e

echo "🔧 Setting up Personal Internet Cell Network Configuration..."

# Configuration variables
WG_INTERFACE="wg0"
WG_NETWORK="10.0.0.0/24"
WG_CONTAINER="cell-wireguard"
HOST_INTERFACE="eth0"  # This will be auto-detected

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

log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

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

# Function to detect the main network interface
detect_interface() {
    # Try to detect the main interface (not loopback, not docker)
    HOST_INTERFACE=$(ip route | grep default | awk '{print $5}' | head -1)
    if [[ -z "$HOST_INTERFACE" ]]; then
        log_error "Could not detect main network interface"
        exit 1
    fi
    log_info "Detected main interface: $HOST_INTERFACE"
}

# Function to enable IP forwarding
enable_ip_forwarding() {
    log_info "Enabling IP forwarding..."
    echo 'net.ipv4.ip_forward = 1' >> /etc/sysctl.conf
    sysctl -p
    log_success "IP forwarding enabled"
}

# Function to configure WireGuard container networking
configure_wireguard_container() {
    log_info "Configuring WireGuard container networking..."
    
    # Check if container is running
    if ! docker ps | grep -q "$WG_CONTAINER"; then
        log_error "WireGuard container is not running"
        exit 1
    fi
    
    # Get container's main interface
    CONTAINER_INTERFACE=$(docker exec $WG_CONTAINER ip route | grep default | awk '{print $5}' | head -1)
    if [[ -z "$CONTAINER_INTERFACE" ]]; then
        CONTAINER_INTERFACE="eth0"
    fi
    
    log_info "Container interface: $CONTAINER_INTERFACE"
    
    # Configure iptables rules inside the container
    docker exec $WG_CONTAINER sh -c "
        # Enable IP forwarding
        echo 1 > /proc/sys/net/ipv4/ip_forward
        
        # Add default route if missing
        ip route add default via 172.20.0.1 dev $CONTAINER_INTERFACE 2>/dev/null || true
        
        # Clear existing rules (be careful!)
        iptables -t nat -F
        iptables -F FORWARD
        
        # Allow forwarding for WireGuard interface
        iptables -A FORWARD -i $WG_INTERFACE -j ACCEPT
        iptables -A FORWARD -o $WG_INTERFACE -j ACCEPT
        
        # NAT rule for internet access
        iptables -t nat -A POSTROUTING -s $WG_NETWORK -o $CONTAINER_INTERFACE -j MASQUERADE
        
        # Allow established and related connections
        iptables -A FORWARD -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT
        
        # Log the configuration
        echo 'Network configuration applied:'
        echo 'IP Forwarding:'
        cat /proc/sys/net/ipv4/ip_forward
        echo 'Routing Table:'
        ip route show
        echo 'NAT Rules:'
        iptables -t nat -L POSTROUTING -n
        echo 'Forwarding Rules:'
        iptables -L FORWARD -n
    "
    
    log_success "WireGuard container networking configured"
}

# Function to configure host networking
configure_host_networking() {
    log_info "Configuring host networking..."
    
    # Enable IP forwarding on host
    echo 'net.ipv4.ip_forward = 1' > /etc/sysctl.d/99-wireguard.conf
    sysctl -p /etc/sysctl.d/99-wireguard.conf
    
    # Configure iptables rules on host
    iptables -t nat -A POSTROUTING -s $WG_NETWORK -o $HOST_INTERFACE -j MASQUERADE
    iptables -A FORWARD -i $WG_INTERFACE -j ACCEPT
    iptables -A FORWARD -o $WG_INTERFACE -j ACCEPT
    
    # Save iptables rules
    if command -v iptables-save >/dev/null 2>&1; then
        mkdir -p /etc/iptables
        iptables-save > /etc/iptables/rules.v4
        log_info "iptables rules saved"
    fi
    
    log_success "Host networking configured"
}

# Function to create persistent configuration
create_persistent_config() {
    log_info "Creating persistent configuration..."
    
    # Create systemd service for network configuration
    cat > /etc/systemd/system/pic-network.service << EOF
[Unit]
Description=Personal Internet Cell Network Configuration
After=docker.service
Requires=docker.service

[Service]
Type=oneshot
RemainAfterExit=yes
ExecStart=/opt/pic/scripts/setup-network.sh
ExecReload=/opt/pic/scripts/setup-network.sh

[Install]
WantedBy=multi-user.target
EOF

    # Enable the service
    systemctl daemon-reload
    systemctl enable pic-network.service
    
    log_success "Persistent configuration created"
}

# Function to test the configuration
test_configuration() {
    log_info "Testing network configuration..."
    
    # Check if WireGuard interface is up
    if docker exec $WG_CONTAINER ip link show $WG_INTERFACE >/dev/null 2>&1; then
        log_success "WireGuard interface is up"
    else
        log_error "WireGuard interface is not up"
        return 1
    fi
    
    # Check NAT rules
    if docker exec $WG_CONTAINER iptables -t nat -L POSTROUTING | grep -q MASQUERADE; then
        log_success "NAT rules are configured"
    else
        log_error "NAT rules are missing"
        return 1
    fi
    
    # Check IP forwarding
    if docker exec $WG_CONTAINER cat /proc/sys/net/ipv4/ip_forward | grep -q 1; then
        log_success "IP forwarding is enabled"
    else
        log_error "IP forwarding is not enabled"
        return 1
    fi
    
    log_success "Network configuration test passed"
}

# Function to show status
show_status() {
    log_info "Network Configuration Status:"
    echo "=================================="
    
    echo "WireGuard Interface:"
    docker exec $WG_CONTAINER ip addr show $WG_INTERFACE 2>/dev/null || echo "  Interface not found"
    
    echo -e "\nRouting Table:"
    docker exec $WG_CONTAINER ip route show
    
    echo -e "\nNAT Rules:"
    docker exec $WG_CONTAINER iptables -t nat -L POSTROUTING -n
    
    echo -e "\nForwarding Rules:"
    docker exec $WG_CONTAINER iptables -L FORWARD -n
    
    echo -e "\nIP Forwarding Status:"
    echo "  Container: $(docker exec $WG_CONTAINER cat /proc/sys/net/ipv4/ip_forward)"
    echo "  Host: $(cat /proc/sys/net/ipv4/ip_forward)"
}

# Main execution
main() {
    log_info "Starting Personal Internet Cell Network Setup..."
    
    check_root
    detect_interface
    
    case "${1:-setup}" in
        "setup")
            enable_ip_forwarding
            configure_wireguard_container
            configure_host_networking
            create_persistent_config
            test_configuration
            log_success "Network configuration completed successfully!"
            ;;
        "test")
            test_configuration
            ;;
        "status")
            show_status
            ;;
        "reset")
            log_warning "Resetting network configuration..."
            docker exec $WG_CONTAINER iptables -t nat -F
            docker exec $WG_CONTAINER iptables -F FORWARD
            iptables -t nat -D POSTROUTING -s $WG_NETWORK -o $HOST_INTERFACE -j MASQUERADE 2>/dev/null || true
            iptables -D FORWARD -i $WG_INTERFACE -j ACCEPT 2>/dev/null || true
            iptables -D FORWARD -o $WG_INTERFACE -j ACCEPT 2>/dev/null || true
            log_success "Network configuration reset"
            ;;
        *)
            echo "Usage: $0 {setup|test|status|reset}"
            echo "  setup  - Configure network (default)"
            echo "  test   - Test configuration"
            echo "  status - Show current status"
            echo "  reset  - Reset configuration"
            exit 1
            ;;
    esac
}

# Run main function
main "$@"