#!/usr/bin/with-contenv bashio
# shellcheck shell=bash

# --- Common Functions ---

declare -A config
config["MySelf"]="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/$(basename "${BASH_SOURCE[0]}")"

declare -a dns_servers_ipv4=()
declare -a dns_servers_ipv6=()

log_level="$(bashio::config "log_level")"
[[ "$log_level" == "debug" ]] && bashio::log.warning "--- Debug mode is active ---"
[[ "$log_level" == "debug" ]] && set -x

_parse_config() {
    local -n config_ref="$1"
    local config_file="$2"
    local line

    while IFS= read -r line || [[ -n "$line" ]]; do
        # Skip comments and empty lines
        [[ "$line" =~ ^[#!] ]] && continue
        # Extract key and value using regex (trim spaces)
        #if [[ "$line" =~ ^[[:space:]]*([^ =]+)[[:space:]]*=[[:space:]]*(.*)[[:space:]]* ]]; then
        if [[ "$line" =~ ^[[:space:]]*([^=[:space:]]+)[=[:space:]]+(.*)[[:space:]]* ]]; then
            local key="${BASH_REMATCH[1]}"
            local value="${BASH_REMATCH[2]}"
            if [[ "$key" == "Address" ]]; then
                if [[ -n "${config_ref["Address"]:-}" ]]; then
                    config_ref["Address"]+=",${value}"
                else
                    config_ref["Address"]="${value}"
                fi
            else
                config_ref["$key"]="$value"
            fi
        fi
    done < "$config_file"
}

_parse_dns() {
    local -a dns_conf=()
    local -a dns_backup_ipv4=("8.8.8.8" "1.1.1.1")
    local -a dns_backup_ipv6=("2001:4860:4860::8888" "2606:4700:4700::1111")
    local dns_servers=$(bashio::config 'DNS_server')

    mapfile -d ',' -t dns_conf < <(echo "${dns_servers}" | tr -d ' ' | tr -d '\n')
    if [ ${config["IPv4Enabled"]} = "true" ]; then
        for dns_ip in "${dns_conf[@]}"; do
            local result=0
            _check_host "${dns_ip}" || result=$?
            if [ "${result}" -eq 1 ]; then
                dns_servers_ipv4+=("${dns_ip}")
            fi
        done
        if [ ${#dns_servers_ipv4[@]} -eq 0 ]; then
            bashio::log.warning "No valid IPv4 DNS servers configured. Using addon defaults ${dns_backup_ipv4[@]}"
            dns_servers_ipv4=("${dns_backup_ipv4[@]}")
        fi
    fi
    if [ ${config["IPv6Enabled"]} = "true" ]; then
        for dns_ip in "${dns_conf[@]}"; do
            local result=0
            _check_host "${dns_ip}" || result=$?
            if [ "${result}" -eq 2 ]; then
                dns_servers_ipv6+=("${dns_ip}")
            fi
        done
        if [ ${#dns_servers_ipv6[@]} -eq 0 ]; then
            bashio::log.warning "No valid IPv6 DNS servers configured. Using addon defaults ${dns_backup_ipv6[@]}"
            dns_servers_ipv6=("${dns_backup_ipv6[@]}")
        fi
    fi
}

_cmd() {
    cmd="$1"
    bashio::log.debug "Executing command: ${cmd}"
    eval "${cmd}"
}

_check_host() {
    if ipcalc -c -4 "$1" >/dev/null 2>&1; then
        return 1  # IPv4
    elif ipcalc -c -6 "$1" >/dev/null 2>&1; then
        return 2  # IPv6
    elif getent ahosts "$1" >/dev/null 2>&1; then
        return 3  # resolvable hostnamee
    else
        return 0  # neither IP nor resolvable hostname
    fi
}

_resolvconf() {
    local mode=$1
    local resolv_conf="/etc/resolv.conf"
    local resolv_backup="/etc/resolv.conf.bak"

    if [ "${mode}" = "reset" ]; then
        bashio::log.info "Resetting ${resolv_conf} to default DNS servers."
        if bashio::fs.file_exists "${resolv_backup}"; then
            cp "${resolv_backup}" "${resolv_conf}"
        else
            bashio::log.warning "No original resolv.conf backup found. Leaving as is."
        fi
    elif [ "${mode}" = "update" ]; then
        bashio::log.info "Updating ${resolv_conf} with VPN DNS servers."
        if ! bashio::fs.file_exists "${resolv_backup}"; then
            bashio::log.debug "Creating backup of original resolv.conf at ${resolv_backup}"
            cp "${resolv_conf}" "${resolv_backup}" 2>/dev/null || true
        fi
        bashio::log.debug "Updating ${resolv_conf} with DNS servers: ${dns_servers_ipv4[*]} ${dns_servers_ipv6[*]}"
        {
            echo "# Generated by vpn script"
            local dns_ip
            for dns_ip in ${dns_servers_ipv4[@]} ${dns_servers_ipv6[@]}; do
                echo "nameserver ${dns_ip}"
            done
        } > "${resolv_conf}"
    else
        bashio::exit.nok "Invalid resolvconf mode specified. Use 'update' or 'reset'."
    fi
}

_resolve_hostname() {
    local hostname=$1
    local -a ips=()
    local -a ipv4_candidates=()
    local -a ipv6_candidates=()

    # Resolve hostname to IPv4
    mapfile -t ipv4_candidates < <(getent ahostsv4 "${hostname}" | awk '{print $1}' | uniq)

    # Resolve hostname to IPv6
    mapfile -t ipv6_candidates < <(getent ahostsv6 "${hostname}" | awk '{print $1}' | uniq)

    if [ ${#ipv4_candidates[@]} -gt 0 ]; then
        bashio::log.debug "Resolved ${hostname} to ${ipv4_candidates[@]}"
        ips+=("${ipv4_candidates[@]}")
    fi

    if [ ${#ipv6_candidates[@]} -gt 0 ]; then
        bashio::log.debug "Resolved ${hostname} to ${ipv6_candidates[@]}"
        ips+=("${ipv6_candidates[@]}")
    fi

    echo "${ips[@]}"
}

_routing_add() {
    bashio::log.info "Adding routing rules for VPN interface ${config["Interface"]}..."

    local local_ipv4=$(ip addr show ${config["Interface"]} | grep 'inet ' | awk '{print $2}' | cut -d'/' -f1)
    local local_ipv6=$(ip addr show ${config["Interface"]} | grep 'inet6 ' | awk '{print $2}' | cut -d'/' -f1)
    local ipv4
    local ipv6

    # add routing rules for local IPs
    for ipv4 in ${local_ipv4}; do
        config["IPv4Enabled"]="true"
        _cmd "ip -4 route add default dev ${config["Interface"]} table ${config["Table"]}" || return 1
        _cmd "ip -4 rule add priority 1 from ${ipv4} table ${config["Table"]}" || return 1
    done
    for ipv6 in ${local_ipv6}; do
        config["IPv6Enabled"]="true"
        _cmd "ip -6 rule add priority 1 from ${ipv6} table ${config["Table"]}" || return 1
    done
    if [ "${config["IPv6Enabled"]}" = "true" ]; then
        _cmd "ip -6 route add default dev ${config["Interface"]} table ${config["Table"]}" || true
    fi

    # get valid DNS servers
    _parse_dns

    # add routing rules for DNS servers
    for dns_ip in "${dns_servers_ipv4[@]}"; do
        #_cmd "ip -4 route add ${dns_ip} dev ${config["Interface"]}" || return 1
        _cmd "ip -4 rule add priority 1 to ${dns_ip} table ${config["Table"]}" || return 1
    done
    for dns_ip in "${dns_servers_ipv6[@]}"; do
        #_cmd "ip -6 route add ${dns_ip} dev ${config["Interface"]}" || return 1
        _cmd "ip -6 rule add priority 1 to ${dns_ip} table ${config["Table"]}" || return 1
    done

    # Update resolv.conf with VPN DNS servers
    _resolvconf "update"
}

_routing_del() {
    bashio::log.info "Removing routing rules for VPN interface ${config["Interface"]}..."

    _resolvconf "reset"
    while _cmd "ip -4 rule del priority 1 from all table ${config["Table"]} 2>/dev/null"; do :; done
    while _cmd "ip -4 rule del priority 1 to all table ${config["Table"]} 2>/dev/null"; do :; done
    while _cmd "ip -4 route del default dev ${config["Interface"]} table ${config["Table"]} 2>/dev/null"; do :; done
    while _cmd "ip -6 rule del priority 1 from all table ${config["Table"]} 2>/dev/null"; do :; done
    while _cmd "ip -6 rule del priority 1 to all table ${config["Table"]} 2>/dev/null"; do :; done
    while _cmd "ip -6 route del default dev ${config["Interface"]} table ${config["Table"]} 2>/dev/null"; do :; done
}

# --- WireGuard Specific Logic ---

_wg_wait_handshake() {
    local timeout="${1:-20}"
    local iface="${config["Interface"]}"
    local peer_pk="${config["PublicKey"]}"
    local deadline ts

    deadline=$(( $(date +%s) + timeout ))

    while [ "$(date +%s)" -lt "${deadline}" ]; do
        ping -I "${iface}" -c1 -W1 1.1.1.1 >/dev/null 2>&1 || true

        ts="$(wg show "${iface}" latest-handshakes 2>/dev/null | awk -v pk="${peer_pk}" '$1==pk{print $2; exit}')"
        if [ -n "${ts}" ] && [ "${ts}" -gt 0 ] 2>/dev/null; then
            return 0
        fi
        sleep 1
    done

    bashio::log.error "WireGuard handshake not established after ${timeout}s (latest-handshake=${ts:-0})."
    wg show "${iface}" 2>&1 | while IFS= read -r l; do bashio::log.error "${l}"; done
    return 1
}

_wireguard_up() {
    bashio::log.warning "This script force Wireguard to ignore any routes and DNS settings."
    bashio::log.warning "Default route will be inserted into custom routing table: ${config["Table"]}"
    bashio::log.warning "This routing table will be used for traffic from the VPN interface and to the configured DNS servers."
    bashio::log.warning "Qbittorrent bittorrent client shall be set to use the VPN interface ${config["Interface"]} only."

    for key in "Interface" "ListenPort" "PrivateKey" "PublicKey" "EndpointIP" "EndpointPort" ; do
        if [ ! -v config[$key] ] || [ -z "${config[$key]}" ]; then
            bashio::log.error "Missing required WireGuard configuration parameter: ${key}"
            return 1
        fi
    done

    _cmd "ip link add ${config["Interface"]} type wireguard" || return 1
    local allowed_ips=""
    local -a local_ips=()
    mapfile -d ',' -t local_ips < <(echo "${config["Address"]}" | tr -d ' ')
    for local_ip in ${local_ips[@]}; do
        local result=0
        _check_host "${local_ip}" || result=$?
        if [ "${result}" -eq 1 ]; then
            allowed_ips="${allowed_ips},0.0.0.0/0"
            _cmd "ip addr add ${local_ip} dev ${config["Interface"]}" || return 1
        elif [ "${result}" -eq 2 ]; then
            allowed_ips="${allowed_ips},::/0"
            _cmd "ip addr add ${local_ip} dev ${config["Interface"]}" || return 1
        else
            bashio::log.warning "Ignoring invalid local IP address: ${local_ip}"
        fi
    done
    allowed_ips="${allowed_ips#,}"
    if [ -z "${allowed_ips}" ]; then
        bashio::log.error "No valid local IP addresses configured."
        return 1
    fi

    _cmd "wg set ${config["Interface"]} listen-port ${config["ListenPort"]} private-key ${config["PrivateKey"]}" || return 1
    local endpoint="${config["EndpointIP"]}:${config["EndpointPort"]}"
    if [[ "${config["EndpointIP"]}" == *:* ]]; then
      endpoint="[${config["EndpointIP"]}]:${config["EndpointPort"]}"
    fi
    local peer_cmd="wg set ${config["Interface"]} peer ${config["PublicKey"]} endpoint ${endpoint} allowed-ips ${allowed_ips}"
    if [ -n "${config["PresharedKey"]:-}" ]; then
      peer_cmd="${peer_cmd} preshared-key ${config["PresharedKey"]}"
    fi
    if [ -n "${config["PersistentKeepalive"]:-}" ]; then
      peer_cmd="${peer_cmd} persistent-keepalive ${config["PersistentKeepalive"]}"
    fi
    _cmd "${peer_cmd}" || return 1

    if [ -v config["MTU"] ] && [ -n "${config["MTU"]}" ]; then
        _cmd "ip link set ${config["Interface"]} mtu ${config["MTU"]}" || return 1
    fi

    _cmd "ip link set ${config["Interface"]} up" || return 1
    _routing_add
    _wg_wait_handshake 10 || return 1
}

_wireguard_down() {
    _routing_del
    _cmd "ip link set ${config["Interface"]} down" 2>/dev/null || true
    _cmd "ip link del ${config["Interface"]}" 2>/dev/null || true
}

wireguard() {
    local mode=$1
    local interface
    local config_file
    local WIREGUARD_STATE_DIR="/var/run/wireguard"

    if ! bashio::fs.file_exists "${WIREGUARD_STATE_DIR}/interface"; then
        bashio::exit.nok 'WireGuard runtime configuration not prepared. Please restart the add-on.'
    fi
    interface=$(cat "${WIREGUARD_STATE_DIR}/interface")
    if [ -z "${interface}" ]; then
        bashio::exit.nok 'WireGuard runtime configuration not prepared. Please restart the add-on.'
    fi
    if ! bashio::fs.file_exists "${WIREGUARD_STATE_DIR}/config"; then
        bashio::exit.nok 'WireGuard runtime configuration not prepared. Please restart the add-on.'
    fi
    config_file=$(cat "${WIREGUARD_STATE_DIR}/config")
    if [ -z "${config_file}" ]; then
        bashio::exit.nok 'WireGuard runtime configuration not prepared. Please restart the add-on.'
    fi

    bashio::log.info "Using Wireguard configuration file: ${config_file}"

    _parse_config config "${config_file}"
    config["Interface"]="${interface}"
    config["ConfigFile"]="${config_file}"
    config["Table"]="${config["Table"]:-1000}"
    config["ListenPort"]="${config["ListenPort"]:-51820}"
    config["EndpointHost"]="${config["Endpoint"]%:*}"
    config["EndpointPort"]="${config["Endpoint"]##*:}"
    config["IPv4Enabled"]="false"
    config["IPv6Enabled"]="false"
    for key in "${!config[@]}"; do
        bashio::log.debug "${key}: ${config[$key]}"
    done

    printf '%s\n' "${config["PrivateKey"]}" > "${WIREGUARD_STATE_DIR}/privatekey"
    chmod 600 "${WIREGUARD_STATE_DIR}/privatekey" || true
    config["PrivateKey"]="${WIREGUARD_STATE_DIR}/privatekey"
    
    if [ -n "${config["PresharedKey"]:-}" ]; then
      printf '%s\n' "${config["PresharedKey"]}" > "${WIREGUARD_STATE_DIR}/presharedkey"
      chmod 600 "${WIREGUARD_STATE_DIR}/presharedkey" || true
      config["PresharedKey"]="${WIREGUARD_STATE_DIR}/presharedkey"
    fi

    if [ "${mode}" = "up" ]; then
        bashio::log.info "Starting WireGuard on interface ${config["Interface"]}..."
        local result=0
        _check_host "${config["EndpointHost"]}" || result=$?
        if [ "${result}" -eq 0 ]; then
            bashio::log.error "WireGuard endpoint ${config["EndpointHost"]} is neither a valid IP address nor a resolvable hostname."
            bashio::exit.nok 'WireGuard start failed.'
        elif [ "${result}" -eq 3 ]; then
            local -a endpoint_ips=()
            mapfile -d ' ' -t endpoint_ips < <(_resolve_hostname ${config["EndpointHost"]})
            if [ ${#endpoint_ips[@]} -eq 0 ]; then
                bashio::log.error "Failed to resolve WireGuard endpoint hostname: ${config["EndpointHost"]}"
                bashio::exit.nok 'WireGuard start failed.'
            fi
            for endpoint_ip in "${endpoint_ips[@]}"; do
                bashio::log.info "Resolved WireGuard endpoint hostname ${config["EndpointHost"]} to IP: ${endpoint_ip}"
                config["EndpointIP"]="${endpoint_ip}"
                if _wireguard_up; then
                    bashio::log.info "WireGuard interface ${config["Interface"]} is up."
                    bashio::exit.ok 'WireGuard started.'
                fi
                bashio::log.error 'WireGuard failed to establish connection.'
                _wireguard_down
            done
        else
            bashio::log.debug "WireGuard endpoint ${config["EndpointHost"]} is a valid IP address. Using as is."
            config["EndpointIP"]="${config["EndpointHost"]}"
            if _wireguard_up; then
                bashio::log.info "WireGuard interface ${config["Interface"]} is up."
                bashio::exit.ok 'WireGuard started.'
            fi
            bashio::log.error 'WireGuard failed to establish connection.'
            _wireguard_down
        fi
    elif [ "${mode}" = "down" ]; then
        bashio::log.info "Stopping WireGuard on interface ${config["Interface"]}..."
        _wireguard_down
        bashio::log.info "WireGuard on interface ${config["Interface"]} is down."
        bashio::exit.ok 'WireGuard stopped.'
    else
        bashio::log.error "Invalid WireGuard mode specified. Use 'up' or 'down'."
        bashio::exit.nok 'WireGuard start failed.'
    fi

    bashio::exit.nok 'WireGuard start failed.'
}

# --- OpenVPN Specific Logic ---

_openvpn_up() {
    bashio::log.warning "This script force OpenvPN to ignore any routes and DNS settings pushed by the server."
    bashio::log.warning "Default route will be inserted into custom routing table: ${config["Table"]}"
    bashio::log.warning "This routing table will be used for traffic from the VPN interface and to the configured DNS servers."
    bashio::log.warning "Qbittorrent bittorrent client shall be set to use the VPN interface ${config["Interface"]} only."

    # Register this script as OpenVPN up/down handlers to manage routing
    echo '#!/bin/bash' > ${config["PostUpScript"]}
    echo "${config["MySelf"]} openvpn postup" >> ${config["PostUpScript"]}
    chmod 755 ${config["PostUpScript"]}
    echo '#!/bin/bash' > ${config["PostDownScript"]}
    echo "${config["MySelf"]} openvpn postdown" >> ${config["PostDownScript"]}
    chmod 755 ${config["PostDownScript"]}

    # Define logging
    log_path="/dev/null"
    [[ "$log_level" == "debug" ]] && log_path="/proc/1/fd/1"

    # Start OpenVPN in the background
    _cmd "/usr/sbin/openvpn \
        --config "${config["ConfigFile"]}" \
        --client \
        --daemon \
        --log "$log_path" \
        --script-security 2 \
        --auth-user-pass "${OPENVPN_STATE_DIR}/credentials.conf" \
        --auth-retry none \
        --up ${config["PostUpScript"]} \
        --down ${config["PostDownScript"]} \
        --up-delay  \
        --up-restart \
        --route-nopull \
        --route-noexec" || return 1

    #wait for slow OpenVPN interface to come up
    for i in {1..10}; do
        if ip link show "${config["Interface"]}" > /dev/null 2>&1 ; then
            return 0
        fi
        sleep 2
    done
    bashio::log.error "OpenVPN interface ${config["Interface"]} failed to come up."
    return 1
}

_openvpn_down() {
    # Terminate OpenVPN process
    pkill -f "openvpn --config ${config["ConfigFile"]}" || true
    _routing_del
}

openvpn() {
    local mode=$1
    local interface
    local config_file
    local OPENVPN_STATE_DIR="/var/run/openvpn"

    if ! bashio::fs.file_exists "${OPENVPN_STATE_DIR}/interface"; then
        bashio::exit.nok 'OpenVPN runtime configuration not prepared. Please restart the add-on.'
    fi
    interface=$(cat "${OPENVPN_STATE_DIR}/interface")
    if [ -z "${interface}" ]; then
        bashio::exit.nok 'OpenVPN runtime configuration not prepared. Please restart the add-on.'
    fi
    if ! bashio::fs.file_exists "${OPENVPN_STATE_DIR}/config"; then
        bashio::exit.nok 'OpenVPN runtime configuration not prepared. Please restart the add-on.'
    fi
    config_file=$(cat "${OPENVPN_STATE_DIR}/config")
    if [ -z "${config_file}" ]; then
        bashio::exit.nok 'OpenVPN runtime configuration not prepared. Please restart the add-on.'
    fi

    bashio::log.warning "Using OpenVPN configuration file: ${config_file}"

    _parse_config config "${config_file}"
    config["Interface"]="${interface}"
    config["ConfigFile"]="${config_file}"
    config["Table"]="${config["Table"]:-1000}"
    config["PostUpScript"]="${OPENVPN_STATE_DIR}/up.sh"
    config["PostDownScript"]="${OPENVPN_STATE_DIR}/down.sh"

    if [ "${mode}" = "up" ]; then
        # register up and down scripts
        bashio::log.info "Starting OpenVPN on interface ${config["Interface"]}..."
        if _openvpn_up; then
            bashio::log.info "OpenVPN interface ${config["Interface"]} is up."
            bashio::exit.ok 'OpenVPN started.'
        fi
        bashio::log.error 'OpenVPN failed to establish connection.'
        _openvpn_down
    elif [ "${mode}" = "down" ]; then
        bashio::log.info "Stopping OpenVPN on interface ${config["Interface"]}..."
        _openvpn_down
        bashio::log.info "OpenVPN on interface ${config["Interface"]} is down."
        bashio::exit.ok 'OpenVPN stopped.'
    elif [ "${mode}" = "postup" ]; then
        _routing_add
        bashio::exit.ok 'OpenVPN routes added.'
    elif [ "${mode}" = "postdown" ]; then
        _routing_del
        bashio::exit.ok 'OpenVPN routes deleted.'
    else
        bashio::log.error "Invalid OpenVPN mode specified. Use 'up', 'down', 'postup', or 'postdown'."
        bashio::exit.nok 'OpenVPN start failed.'
    fi

    bashio::exit.nok 'OpenVPN start failed.'
}

# --- Entry Point ---

if [ $# -ne 2 ]; then
    bashio::log.error "Invalid number of arguments. Usage: vpn.sh <wireguard|openvpn> <up|down>"
    bashio::exit.nok 'VPN start failed.'
fi
if [[ "$1" == "wireguard" ]]; then
    wireguard "$2"
elif [[ "$1" == "openvpn" ]]; then
    openvpn "$2"
else
    bashio::log.error "Invalid VPN type specified. Use 'wireguard' or 'openvpn'."
    bashio::exit.nok 'VPN start failed.'
fi