#!/bin/bash

# ==============================================================================
# 1. SETTINGS & GLOBALS
# ==============================================================================
# Exit immediately if a command exits with a non-zero status
set -e 

# Determine the directory where this script resides
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

# Input arguments
client_name="$1"

# Global variables
server_pubkey=""
client_pubkey=""
folder_log=""
ips=""
users=""

# ==============================================================================
# 2. HELPER FUNCTIONS
# ==============================================================================
# Load environment variable of the script
environment_var() {
    echo "--- Starting Environment & Pre-flight Checks ---"
    # --- .env LOADING ---
    # This safely loads the variables from .env without touching other files
    if [ -f "$SCRIPT_DIR/.env" ]; then
        set -a
        source "$SCRIPT_DIR/.env"
        set +a
        echo "[✔] Environment variables loaded."
    else
        echo "[!] No .env file found. Using defaults."
    fi

    # --- Variable Checks ---
    server_pubkey=$(cat ${WG0_DIR}/publickey)
    ips=$(grep -i '^\s*AllowedIPs\s*=\s*' "${WG0_DIR}/wg0.conf" | sed 's/^\s*AllowedIPs\s*=\s*//I; s/\s*$//')
    users=$(grep -i '^\s*# BEGIN_PEER \s*' "${WG0_DIR}/wg0.conf" | sed 's/^\s*# BEGIN_PEER \s*//i; s/\s*$//')
}

# Find the last registered client IP
find_last_ip() {
    # Find the last ip an create the next ip
    if [ -n "$ips" ]; then
        last=$(printf "%s\n" "$ips" | tail -1)
        addr=${last%/*}   # strip /32 if present
        IFS=. read -r a b c d <<< "$addr"
        next_ip="10.0.0.$((d + 1))"
        printf "Last CIDR: %s\nNext IP: %s/32\n" "$last" "$next_ip"
    else
        next_ip="10.0.0.2"
        printf "No peers found, starting at %s/32\n" "$next_ip"
    fi
}

# Display QR Code and Generate PNG
display_qr_code() {
    local config_file="$1"
    # Derive user folder path from the already established global TOOLS_DIR
    local user_folder_path="${TOOLS_DIR}/logs/${client_name}"
    local base_name="${client_name}"

    if command -v qrencode >/dev/null 2>&1; then
        printf '\n--- Scan QR Code with WireGuard App ---\n\n'
        
        # 1. Output to terminal in large ANSI UTF-8 format
        qrencode -t ansiutf8 < "$config_file"
        printf '\n'

        # 2. Generate PNG version in the user's log folder
        local png_file="${user_folder_path}/${base_name}.png"
        
        # -t PNG: Output format
        # -o: Output file path
        # -s 10: Size (makes a high-quality scan image)
        # -l M: Standard error correction
        printf 'Generating QR Code PNG image...\n'
        qrencode -t PNG -s 10 -l M -o "$png_file" < "$config_file"
        printf 'QR Code PNG saved to: %s\n' "$png_file"

        # Make it accessible
        chmod 777 $png_file
    else
        printf '\n[Notice] "qrencode" is not installed. QR codes could not be generated.\n\n'
    fi
}

# Create new client
client_create_new() {
    # Make sure that the client name doesn't have spaces
    if [[ "$client_name" =~ [[:space:]] ]]; then
        printf "Error: Client name cannot contain spaces. Use underscores instead of spaces.\n"
        exit 1
    elif [[ "$client_name" == "" ]]; then
        printf "Error: Client name argument is mandatory. Please provide one first.\n"
        exit 1
    fi

    # Enforce exact match tags
    local peer_begin="# BEGIN_PEER ${client_name}"

    # Check if client already exists
    if grep -i "^${peer_begin}$" "${WG0_DIR}/wg0.conf"; then
        printf 'Error: Client "%s" already exists\n' "$client_name"
        exit 1
    fi

    # Log client config with timestamp
    folder_log="${TOOLS_DIR}/logs/${client_name}"
    mkdir -p "${folder_log}"

    # Generate client keys
    printf '\n--- Generating keys for client "%s" ---\n' "$client_name"
    umask 077 && wg genkey | tee "${folder_log}/privatekey" | wg pubkey > "${folder_log}/publickey"
    client_privkey=$(cat "${folder_log}/privatekey")
    client_pubkey=$(cat "${folder_log}/publickey")

    printf 'Client Private Key: %s\n' "$client_privkey"
    printf 'Client Public Key: %s\n' "$client_pubkey"

    # Show template from .client_template file
    template_file="${TOOLS_DIR}/client_template"
    final_conf="${folder_log}/${client_name}.conf"

    if [ -f "$template_file" ]; then
        printf '\n--- Wireguard client config (generated from %s) ---\n' "$template_file"

        # Process template: replace placeholders
        sed -e "s|<%CLIENT_PRIVATE_KEY%>|${client_privkey}|g" \
            -e "s|<%CLIENT_IP%>|${next_ip}|g" \
            -e "s|<%SERVER_PUBLIC_KEY%>|${server_pubkey}|g" \
            -e "s|<%DOMAIN_NAME%>|${DOMAIN_NAME}|g" \
            "$template_file" > "${final_conf}"

        printf '\n--- Client config saved to: %s ---\n' "${final_conf}"
        cat "${final_conf}"

        # Make it accessible
        chmod 777 ${final_conf}

        # Output QR code to terminal and save PNG
        display_qr_code "${final_conf}"
    else
        printf '\n--- Error: Template file %s not found ---\n' "$template_file"
        exit 1
    fi

}

# Add the new client to wg0.conf
client_add_new() {
    # Add peer to wg0.conf
    printf '\n--- Adding peer to wg0.conf ---\n'

    # Create peer section
    peer_entry="# BEGIN_PEER ${client_name}
[Peer]
PublicKey = ${client_pubkey}
AllowedIPs = ${next_ip}/32
# END_PEER ${client_name}"

    # Append to wg0.conf
    echo "" >> "${WG0_DIR}/wg0.conf"
    echo "$peer_entry" >> "${WG0_DIR}/wg0.conf"

    printf 'Peer "%s" added to wg0.conf\n' "$client_name"
    printf 'AllowedIPs: %s/32\n' "$next_ip"

    # Apply config without restart
    printf 'Applying config to running WireGuard interface...\n'
    if wg syncconf wg0 <(wg-quick strip wg0); then
        printf 'Config applied successfully.\n'
        peer_status="active"
    else
        printf 'Error: wg syncconf failed\n'
        printf 'Check config syntax with: wg-quick strip wg0\n'
        peer_status="failed"
        exit 1
    fi

    # Log to CSV
    csv_file="${TOOLS_DIR}/logs/client_log.csv"
    timestamp=$(date -u +"%Y-%m-%dT%H:%M:%SZ")

    if [ ! -f "$csv_file" ]; then
        echo "timestamp,client_name,client_public_key,client_private_key,client_ip,status" > "$csv_file"
    fi

    printf '%s,"%s","%s","%s","%s","%s"\n' \
        "$timestamp" \
        "$client_name" \
        "$client_pubkey" \
        "$client_privkey" \
        "$next_ip/32" \
        "$peer_status" >> "$csv_file"

    printf '\n--- Client logged to: %s (status: %s) ---\n' "$csv_file" "$peer_status"
    printf '\n--- Done. WireGuard is running with the new peer ---\n'
}

# ==============================================================================
# 3. MAIN LOGIC
# ==============================================================================
main() {
    # Load environment variables first
    environment_var

    # Print current values
    printf "Public Key of the server: %s\n" "$server_pubkey"
    printf "List of current clients:\n"
    paste -d '|' <(echo "$users") <(echo "$ips") | while IFS='|' read -r user ip; do
        echo "User: $user -> IP: $ip"
    done
    printf "\n"

    # Find the last IP
    find_last_ip
    # Create new client configuration file
    client_create_new
    # Add new client to wg0.conf
    client_add_new
}

# Invoke main
main "$@"