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

# --- WireGuard Specific Logic ---

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

_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]}"
            config_ref["$key"]="$value"
        fi
    done < "$config_file"
}

_parse_dns() {
    local -a dns_servers=()
    local dns_ip

    while IFS=',' read -r dns_ip; do
        if _is_ip_address "${dns_ip}"; then
            bashio::log.warning "Ignoring invalid DNS server address: ${dns_ip}"
            continue
        fi
        dns_servers+=("${dns_ip}")
    done <<< $(bashio::config 'DNS_server' | tr -d ' ')

    if [ ${#dns_servers[@]} -eq 0 ]; then
        bashio::log.warning "No valid DNS servers configured. Using addon defaults."
        dns_servers=("8.8.8.8" "1.1.1.1")
    fi
    config["DnsServers"]="${dns_servers[*]}"
}

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

_is_ip_address() {
    if [ "$1" != "${1#*[0-9].[0-9]}" ]; then
        return 1  # IPv4
    elif [ "$1" != "${1#*:[0-9a-fA-F]}" ]; then
        return 2  # IPv6
    else
        return 0  # Not an IP address
    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 DNS servers: ${config["DnsServers"]}"
        if ! bashio::fs.file_exists "${resolv_backup}"; then
            cp "${resolv_conf}" "${resolv_backup}" 2>/dev/null || true
        fi
        bashio::log.warn "Overriding ${resolv_conf} with DNS servers: ${config["DnsServers"]}"
        local valid_dns="false"
        {
            local dns_ip
            for dns_ip in ${config["DnsServers"]}; do
                _is_ip_address "${dns_ip}"
                local is_ip=$?
                if [ "${is_ip}" -eq 1 ] && [ ${config["IPv4Enabled"]} = "true" ]; then
                    echo "nameserver ${dns_ip}"
                    valid_dns="true"
                elif [ "${is_ip}" -eq 2 ] && [ "${config["IPv6Enabled"]}" = "true" ]; then
                    echo "nameserver ${dns_ip}"
                    valid_dns="true"
                else
                    bashio::log.warning "Ignoring invalid DNS server address: ${dns_ip}"
                    continue
                fi
            done
        } > "${resolv_conf}"
        if [ "${valid_dns}" = "false" ]; then
            bashio::exit.nok "No valid DNS servers could be written to ${resolv_conf}."
        fi
    else
        bashio::exit.nok "Invalid resolvconf mode specified. Use 'update' or 'reset'."
    fi
}

_resolve_hostname() {
    local hostname=$1
    local ips=""

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

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

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

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

    return $ips
}

_routing_add() {
    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, ipv6
    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 route add default dev ${config["Interface"]} table ${config["Table"]}" || return 1
        _cmd "ip -6 rule add priority 1 from ${ipv6} table ${config["Table"]}" || return 1
    done

    local dns_ip
    for dns_ip in ${config["DnsServers"]}; do
        _is_ip_address "${dns_ip}"
        local is_ip=$?
        if [ "${is_ip}" -eq 0 ]; then
            bashio::log.warning "Ignoring invalid DNS server address: ${dns_ip}"
            continue
        elif [ "${is_ip}" -eq 1 ] && [ ${config["IPv4Enabled"]} = "true" ]; then
            #_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
        elif [ "${is_ip}" -eq 2 ] && [ "${config["IPv6Enabled"]}" = "true" ]; then
            #_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
        else
            bashio::log.warning "Failed to add route for DNS server: ${dns_ip}"
        fi
    done
    _resolvconf "update"
}

_routing_del() {
    _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_up() {
    bashio::log.warn "Bringing up Wireguard interface on ${config["Interface"]}..."
    bashio::log.warn "Using Wireguard configuration file: ${config["ConfigFile"]}"
    bashio::log.warn "This script force Wireguard to ignore any routes and DNS settings."
    bashio::log.warn "Default route will be inserted into custom routing table: ${config["Table"]}"
    bashio::log.warn "This routing table will be used for traffic from the VPN interface and to configured DNS servers."
    bashio::log.warn "Qbittorrent bittorrent client shall be set to use the VPN interface ${config["Interface"]} only."

    _cmd "ip link add dev ${config["Interface"]} type wireguard" || return 1
    local allowed_ips=""
    for local_ip in ${config["Address"]}; do
        _is_ip_address "${local_ip}"
        local is_ip=$?
        if [ "${is_ip}" -eq 1 ]; then
            allowed_ips="${allowed_ips},0.0.0.0/0"
            _cmd "ip addr add ${local_ip} dev ${config["Interface"]}" || return 1
        elif [ "${is_ip}" -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
    _cmd "wg set ${config["Interface"]} peer ${config["PublicKey"]} endpoint ${config["EndpointIP"]}:${config["EndpointPort"]} allowed-ips ${allowed_ips}" || return 1
    if [ -n "${config["PersistentKeepalive"]}" ]; then
        _cmd "wg set ${config["Interface"]} peer ${config["PublicKey"]} persistent-keepalive ${config["PersistentKeepalive"]}" || return 1
    fi
    _cmd "ip link set ${config["Interface"]} up" || return 1
    _routing_add
}

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

    _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

    echo ${config["PrivateKey"]} > ${WIREGUARD_STATE_DIR}/privatekey
    config["PrivateKey"]="${WIREGUARD_STATE_DIR}/privatekey"

    if [ "${mode}" = "up" ]; then
        bashio::log.info "Starting WireGuard interface ${config["Interface"]}..."
        if _is_ip_address ${config["EndpointHost"]}; then
            local 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.info "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 interface ${config["Interface"]}..."
        _wireguard_down
        bashio::log.info "WireGuard 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_up() {
    bashio::log.warn "Bringing up OpenVPN interface on ${config["Interface"]}..."
    bashio::log.warn "Using OpenVPN configuration file: ${config["ConfigFile"]}"
    bashio::log.warn "This script force OpenvPN to ignore any routes and DNS settings pushed by the server."
    bashio::log.warn "Default route will be inserted into custom routing table: ${config["Table"]}"
    bashio::log.warn "This routing table will be used for traffic from the VPN interface and to configured DNS servers."
    bashio::log.warn "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 "${config["MySelf"]} openvpn postup" > ${config["PostUpScript"]}
    chmod 755 ${config["PostUpScript"]}
    echo "${config["MySelf"]} openvpn postdown" > ${config["PostDownScript"]}
    chmod 755 ${config["PostDownScript"]}

    # Start OpenVPN in the background
    # (maybe use setsid instead of nohup to detach completely?)
    nohup /usr/sbin/openvpn \
        --config "${config["ConfigFile"]}" \
        --script-security 2 \
        --up ${config["PostUpScript"]} \
        --down ${config["PostDownScript"]} \
        --route-nopull \
        --pull-filter ignore "route" \
        --pull-filter ignore "redirect-gateway" \
        --pull-filter ignore "dhcp-option DNS" \
        --pull-filter ignore "route-ipv6" \
        --pull-filter ignore "redirect-gateway ipv6" \
        --pull-filter ignore "dhcp-option DNS6" \
    &
}

_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

    _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 with configuration file: ${config_file}..."
        _openvpn_up
        bashio::exit.ok 'OpenVPN started.'
    elif [ "${mode}" = "down" ]; then
        bashio::log.info "Stopping OpenVPN..."
        _openvpn_down
        bashio::exit.ok 'OpenVPN stopped.'
    elif [ "${mode}" = "postup" ]; then
        _routing_add
    elif [ "${mode}" = "postdown" ]; then
        _routing_del
    else
        bashio::log.error "Invalid OpenVPN mode specified. Use 'up', 'down', 'postup', or 'postdown'."
        bashio::exit.nok 'OpenVPN start failed.'
    fi
    bashio::log.info "Starting OpenVPN with configuration file: ${config_file}"
}


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
_parse_dns
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
