#!/bin/bash

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

# Function to print colored output
print_color() {
    local color=$1
    shift
    echo -e "${color}$@${NC}"
}

# Function to display usage information
show_usage() {
    print_color $BLUE "Usage: $0 <path_to_cfg_file> [interface_name] [peer_description] [zone] [-t|--test-route]"
    print_color $BLUE "Options:"
    print_color $BLUE "  -t, --test-route    Create a test route to 49.12.234.183"
    exit 1
}

# Function to get parameter from config file
get_param() {
    sed -n "/$1/{s/^[^=]*=//; s/^[[:space:]]*//; s/[[:space:]]*$//; p}" "$cfg_file"
}

# Function to add or update network interface
configure_interface() {
    if ! uci show network | grep -q "network.${interface_name}=interface"; then
        uci batch << EOI
set network.${interface_name}=interface
set network.${interface_name}.proto='amneziawg'
set network.${interface_name}.private_key='$(get_param 'PrivateKey')'
set network.${interface_name}.addresses='$(get_param 'Address')'
set network.${interface_name}.force_link='1'
set network.${interface_name}.delegate='0'
set network.${interface_name}.dns='$(get_param 'DNS' | tr ',' ' ')'
set network.${interface_name}.awg_jc='$(get_param 'Jc')'
set network.${interface_name}.awg_jmin='$(get_param 'Jmin')'
set network.${interface_name}.awg_jmax='$(get_param 'Jmax')'
set network.${interface_name}.awg_h1='$(get_param 'H1')'
set network.${interface_name}.awg_h2='$(get_param 'H2')'
set network.${interface_name}.awg_h3='$(get_param 'H3')'
set network.${interface_name}.awg_h4='$(get_param 'H4')'
EOI
        print_color $GREEN "Interface ${interface_name} created."
    else
        print_color $YELLOW "Interface ${interface_name} already exists."
    fi
}

# Function to add or update peer
configure_peer() {
    local peer_exists=false
    local peer_section=""

    # Check if peer already exists
    for peer in $(uci show network | grep "network.@amneziawg_${interface_name}\[" | cut -d. -f2 | cut -d= -f1); do
        if [ "$(uci get network.${peer}.description)" = "${peer_description}" ]; then
            peer_exists=true
            peer_section="${peer}"
            break
        fi
    done

    local allowed_ips=$(get_param 'AllowedIPs')
    IFS=',' read -ra ips <<< "$allowed_ips"

    if [ "$peer_exists" = false ]; then
        peer_section=$(uci add network amneziawg_${interface_name})
        uci batch << EOI
set network.${peer_section}.description='${peer_description}'
set network.${peer_section}.public_key='$(get_param 'PublicKey')'
set network.${peer_section}.route_allowed_ips='0'
set network.${peer_section}.persistent_keepalive='25'
set network.${peer_section}.endpoint_host='$(get_param 'Endpoint' | cut -d':' -f1)'
set network.${peer_section}.endpoint_port='$(get_param 'Endpoint' | cut -d':' -f2)'
set network.${peer_section}.persistent_keepalive='60'
EOI
        print_color $GREEN "Peer ${peer_description} added to interface ${interface_name}."
    else
        uci batch << EOI
set network.${peer_section}.public_key='$(get_param 'PublicKey')'
set network.${peer_section}.endpoint_host='$(get_param 'Endpoint' | cut -d':' -f1)'
set network.${peer_section}.endpoint_port='$(get_param 'Endpoint' | cut -d':' -f2)'
set network.${peer_section}.persistent_keepalive='60'
delete network.${peer_section}.allowed_ips
EOI
        print_color $YELLOW "Peer ${peer_description} updated in interface ${interface_name}."
    fi

    # Add allowed IPs
    for ip in "${ips[@]}"; do
        uci add_list network.${peer_section}.allowed_ips="${ip//[[:space:]]/}"
    done
}

# Function to configure firewall
configure_firewall() {
    print_color $BLUE "Configuring firewall..."
    
    # Create or update the zone for the new interface
    zone_config=$(uci show firewall | grep "@zone\[.*\].name='${zone_name}'")
    if [ -z "$zone_config" ]; then
        # Create new zone
        new_zone=$(uci add firewall zone)
        uci batch << EOI
set firewall.$new_zone.name='${zone_name}'
set firewall.$new_zone.input='REJECT'
set firewall.$new_zone.output='ACCEPT'
set firewall.$new_zone.forward='REJECT'
set firewall.$new_zone.masq='1'
set firewall.$new_zone.mtu_fix='1'
set firewall.$new_zone.network='${interface_name}'
EOI
        print_color $GREEN "Created new firewall zone ${zone_name} for interface ${interface_name}."
    else
        # Update existing zone
        existing_zone=$(echo "$zone_config" | cut -d. -f2 | cut -d= -f1)
        uci set firewall.$existing_zone.network="${interface_name}"
        print_color $YELLOW "Updated existing firewall zone ${zone_name} with interface ${interface_name}."
    fi

    # Remove all existing forwarding rules for this zone
    existing_forwards=$(uci show firewall | grep "@forwarding\[[0-9]\+\]" | grep "\.dest='${zone_name}'")
    if [ -n "$existing_forwards" ]; then
        print_color $YELLOW "Removing existing forwarding rules for ${zone_name}..."
        echo "$existing_forwards" | while read -r line; do
            section=$(echo "$line" | cut -d. -f2 | cut -d= -f1)
            uci delete firewall.$section
        done
    fi

    # Add a single new forwarding rule
    new_forward=$(uci add firewall forwarding)
    uci set firewall.$new_forward.src='lan'
    uci set firewall.$new_forward.dest="${zone_name}"
    print_color $GREEN "Added new forwarding rule from LAN to ${zone_name}."

    print_color $GREEN "Firewall settings for ${zone_name} zone have been configured."
}

# Function to create test route
create_test_route() {
    local address=$(get_param 'Address')
    local gateway=$(echo $address | awk -F'[./]' '{print $1"."$2"."$3".1"}')
    
    uci batch << EOI
set network.test_route=route
set network.test_route.interface='${interface_name}'
set network.test_route.target='49.12.234.183/32'
set network.test_route.gateway='${gateway}'
EOI
    
    print_color $GREEN "Test route created:"
    print_color $BLUE "  Target: 49.12.234.183/32"
    print_color $BLUE "  Gateway: ${gateway}"
    print_color $BLUE "  Interface: ${interface_name}"
    print_color $YELLOW "You can test the connection using: ${NC}curl 49.12.234.183"
}

# Main execution
# Parse command line arguments
CREATE_TEST_ROUTE=false
while [[ $# -gt 0 ]]; do
    case $1 in
        -t|--test-route)
            CREATE_TEST_ROUTE=true
            shift
            ;;
        -h|--help)
            show_usage
            ;;
        *)
            if [ -z "$cfg_file" ]; then
                cfg_file="$1"
            elif [ -z "$interface_name" ]; then
                interface_name="$1"
            elif [ -z "$peer_description" ]; then
                peer_description="$1"
            elif [ -z "$zone_name" ]; then
                zone_name="$1"
            else
                print_color $RED "Unknown argument: $1"
                show_usage
            fi
            shift
            ;;
    esac
done

# Check for minimum required arguments
[ -z "$cfg_file" ] && show_usage

# Set default values if not provided
interface_name="${interface_name:-AWG}"
peer_description="${peer_description:-openwrt_router}"
zone_name="${zone_name:-awg}"

print_color $BLUE "Starting Amnezia WireGuard setup..."

configure_interface
configure_peer
uci commit network

configure_firewall
uci commit firewall

if [ "$CREATE_TEST_ROUTE" = true ]; then
    create_test_route
    uci commit network
fi

print_color $YELLOW "Restarting network..."
/etc/init.d/network restart

print_color $GREEN "Amnezia WireGuard setup completed."