#!/usr/bin/env bash
set -euo pipefail

LAN_IF="end0"
WAN_IF="wlan0"
STATE_FILE="/var/lib/nymea/ngw-dedicated-lan-state.conf"
DNSMASQ_CONF="/etc/dnsmasq.d/nymea-experience-plugin-ngw.conf"
NFT_FILE="/etc/nftables-nymea-experience-plugin-ngw.nft"
PREFIX=24
LEASE_TIME="12h"
CON_NAME="nymea-ngw-dedicated-lan"
DHCP_CON_NAME="nymea-ngw-lan-dhcp"
NFT_FORWARD_COMMENT="nymea-ngw-lan forward"

usage() {
    local program
    program="$(basename "$0")"
    cat <<USAGE
Usage:
  ${program} apply --enabled true|false --gateway-services POLICY --internet-sharing true|false [OPTIONS]
  ${program} reapply [OPTIONS]
  ${program} status [OPTIONS]
  ${program} uninstall [OPTIONS]

Commands:
  apply
      Apply and persist a requested dedicated LAN state. If enabled, end0 is
      configured as a static dedicated LAN with DHCP and nftables rules. If
      disabled, the owned static LAN profile and generated firewall/DHCP files
      are removed and end0 is configured as a DHCP client.

  reapply
      Read the last helper state from the state file and apply it again. This is
      intended for systemd boot restore before nymead has finished starting.

  status
      Print the helper state file, NetworkManager device status, and owned
      nftables tables.

  uninstall
      Remove generated firewall/DHCP state, delete the helper state file, and
      restore end0 to the owned DHCP client profile.

Required apply options:
  --enabled true|false
      Enable or disable the dedicated LAN mode.

  --gateway-services disabled|configured-ports|all-ports
      Select which gateway services LAN clients may reach.
      disabled          Allow only DHCP and ICMP to the gateway.
      configured-ports  Allow only the TCP ports provided by --gateway-tcp-ports.
      all-ports         Allow all LAN-to-gateway traffic.

  --internet-sharing true|false
      Enable or disable full IPv4 NAT from the dedicated LAN through the WAN
      interface. When false, DHCP clients receive no router or DNS option.

Options:
  --gateway-tcp-ports LIST
      Comma-separated TCP port list used with --gateway-services
      configured-ports, for example "80,443,1883,2222,4444".

  --lan-if IFACE
      LAN interface to manage. Default: ${LAN_IF}

  --wan-if IFACE
      WAN interface used for NAT. Default: ${WAN_IF}

  --state-file FILE
      Helper state file. Default: ${STATE_FILE}

Files:
  ${STATE_FILE}
      Last helper state used by reapply.
  ${DNSMASQ_CONF}
      Generated dnsmasq configuration.
  ${NFT_FILE}
      Generated nftables rules.

Notes:
  This helper must run as root. It is normally called by nymead or by
  nymea-experience-plugin-ngw-lan.service, not by users directly.
USAGE
}

die() {
    echo "$*" >&2
    exit 1
}

log() {
    echo "[ngw-dedicated-lan] $*"
}

need_root() {
    [[ "${EUID}" -eq 0 ]] || die "This helper must run as root."
}

require_cmds() {
    local missing=0
    local required_cmd
    for required_cmd in nmcli ip nft dnsmasq systemctl sysctl awk sed grep mkdir chmod rm dirname seq cat; do
        if ! command -v "${required_cmd}" >/dev/null 2>&1; then
            echo "Missing command: ${required_cmd}" >&2
            missing=1
        fi
    done
    [[ "${missing}" -eq 0 ]] || exit 2
}

bool_arg() {
    case "$1" in
        true|false) echo "$1" ;;
        *) die "Invalid boolean: $1" ;;
    esac
}

services_arg() {
    case "$1" in
        disabled|configured-ports|all-ports) echo "$1" ;;
        *) die "Invalid gateway service policy: $1" ;;
    esac
}

validate_ports() {
    local ports="$1" port
    [[ -z "${ports}" ]] && return 0
    IFS=, read -r -a port_array <<<"${ports}"
    for port in "${port_array[@]}"; do
        [[ "${port}" =~ ^[0-9]+$ ]] || die "Invalid TCP port: ${port}"
        [[ "${port}" -ge 1 && "${port}" -le 65535 ]] || die "TCP port out of range: ${port}"
    done
}

detect_default_wan_if() {
    ip -4 route show default | awk -v lan_if="${LAN_IF}" '
        {
            for (i = 1; i <= NF; i++) {
                if ($i == "dev" && (i + 1) <= NF && $(i + 1) != lan_if) {
                    print $(i + 1)
                    exit
                }
            }
        }'
}

wan_has_default_route() {
    local wan_if="$1"
    ip -4 route show default | awk -v wan_if="${wan_if}" '
        {
            for (i = 1; i <= NF; i++) {
                if ($i == "dev" && (i + 1) <= NF && $(i + 1) == wan_if)
                    found = 1
            }
        }
        END { exit found ? 0 : 1 }'
}

resolve_wan_if() {
    local internet_sharing="$1" detected_wan_if
    [[ "${internet_sharing}" == "true" ]] || return 0

    [[ "${WAN_IF}" != "${LAN_IF}" ]] || die "WAN interface must be different from LAN interface: ${WAN_IF}"

    detected_wan_if="$(detect_default_wan_if)"
    if [[ -z "${WAN_IF}" || "${WAN_IF}" == "auto" ]]; then
        [[ -n "${detected_wan_if}" ]] || die "Could not detect a default-route WAN interface for internet sharing."
        log "Using detected default-route WAN interface ${detected_wan_if}."
        WAN_IF="${detected_wan_if}"
        return 0
    fi

    if ! wan_has_default_route "${WAN_IF}"; then
        [[ -n "${detected_wan_if}" ]] || die "Configured WAN interface ${WAN_IF} has no IPv4 default route and no alternative default-route interface was found."
        log "Configured WAN interface ${WAN_IF} has no IPv4 default route; using detected default-route interface ${detected_wan_if}."
        WAN_IF="${detected_wan_if}"
    fi
}

ip_to_int() {
    local ip="$1" a b c d
    IFS=. read -r a b c d <<<"${ip}"
    echo $(( (a << 24) + (b << 16) + (c << 8) + d ))
}

cidr_range() {
    local cidr="$1" ip prefix ip_int mask start size end
    ip="${cidr%/*}"
    prefix="${cidr#*/}"
    ip_int="$(ip_to_int "${ip}")"
    if [[ "${prefix}" -eq 0 ]]; then
        mask=0
    else
        mask=$(( (0xFFFFFFFF << (32 - prefix)) & 0xFFFFFFFF ))
    fi
    start=$(( ip_int & mask ))
    size=$(( 1 << (32 - prefix) ))
    end=$(( start + size - 1 ))
    echo "${start} ${end}"
}

cidr_overlaps() {
    local a_start a_end b_start b_end
    read -r a_start a_end < <(cidr_range "$1")
    read -r b_start b_end < <(cidr_range "$2")
    [[ "${a_start}" -le "${b_end}" && "${b_start}" -le "${a_end}" ]]
}

used_cidrs() {
    ip -4 route show | awk '{print $1}' | while read -r dst; do
        [[ -z "${dst}" || "${dst}" == "default" ]] && continue
        case "${dst}" in
            */*) echo "${dst}" ;;
            *.*.*.*) echo "${dst}/32" ;;
        esac
    done
}

generate_candidates() {
    local b c
    for b in 77 78 79 137 138 139 201 202 203 240 241 242; do
        for c in $(seq 10 250); do
            echo "10.${b}.${c}.0/24"
        done
    done
    for b in $(seq 20 31); do
        for c in $(seq 10 250); do
            echo "172.${b}.${c}.0/24"
        done
    done
    for c in $(seq 200 254); do
        echo "192.168.${c}.0/24"
    done
}

choose_subnet() {
    local candidate u conflict
    mapfile -t used < <(used_cidrs)
    while read -r candidate; do
        conflict=0
        for u in "${used[@]}"; do
            if cidr_overlaps "${candidate}" "${u}"; then
                conflict=1
                break
            fi
        done
        if [[ "${conflict}" -eq 0 ]]; then
            echo "${candidate}"
            return 0
        fi
    done < <(generate_candidates)
    die "Could not find a free private /24 subnet."
}

load_state() {
    if [[ -r "${STATE_FILE}" ]]; then
        log "Loading state from ${STATE_FILE}."
        # shellcheck disable=SC1090
        source "${STATE_FILE}"
    else
        log "No state file found at ${STATE_FILE}."
    fi
}

write_state() {
    local enabled="$1" gateway_services="$2" internet_sharing="$3" tcp_ports="$4" cidr="$5"
    local base gw dhcp_start dhcp_end
    base="${cidr%/*}"
    gw="${base%.*}.1"
    dhcp_start="${base%.*}.2"
    dhcp_end="${base%.*}.254"

    log "Writing state to ${STATE_FILE}: enabled=${enabled}, gateway_services=${gateway_services}, internet_sharing=${internet_sharing}, lan_if=${LAN_IF}, wan_if=${WAN_IF}, cidr=${cidr}, tcp_ports=${tcp_ports:-none}."
    mkdir -p "$(dirname "${STATE_FILE}")"
    cat >"${STATE_FILE}" <<STATE
ENABLED=${enabled}
GATEWAY_SERVICES=${gateway_services}
INTERNET_SHARING=${internet_sharing}
GATEWAY_TCP_PORTS=${tcp_ports}
LAN_IF=${LAN_IF}
WAN_IF=${WAN_IF}
LAN_CIDR=${cidr}
LAN_GW=${gw}
DHCP_START=${dhcp_start}
DHCP_END=${dhcp_end}
PREFIX=${PREFIX}
LEASE_TIME=${LEASE_TIME}
STATE
    chmod 0644 "${STATE_FILE}"
}

configure_nm_static() {
    local cidr="$1" gw
    gw="${cidr%/*}"
    gw="${gw%.*}.1"

    log "Configuring ${LAN_IF} as static dedicated LAN address ${gw}/${PREFIX}."
    nmcli device set "${LAN_IF}" managed yes
    nmcli connection delete "${DHCP_CON_NAME}" >/dev/null 2>&1 || true

    if nmcli -t -f NAME connection show | grep -Fxq "${CON_NAME}"; then
        log "Updating NetworkManager connection ${CON_NAME}."
        nmcli connection modify "${CON_NAME}" connection.interface-name "${LAN_IF}" ipv4.method manual ipv4.addresses "${gw}/${PREFIX}" ipv4.never-default yes ipv4.ignore-auto-dns yes ipv6.method disabled connection.autoconnect yes
    else
        log "Creating NetworkManager connection ${CON_NAME}."
        nmcli connection add type ethernet ifname "${LAN_IF}" con-name "${CON_NAME}" ipv4.method manual ipv4.addresses "${gw}/${PREFIX}" ipv4.never-default yes ipv4.ignore-auto-dns yes ipv6.method disabled connection.autoconnect yes
    fi

    log "Bringing up NetworkManager connection ${CON_NAME}."
    nmcli connection up "${CON_NAME}"
}

configure_nm_dhcp() {
    log "Configuring ${LAN_IF} as DHCP client."
    nmcli device set "${LAN_IF}" managed yes || true
    nmcli connection delete "${CON_NAME}" >/dev/null 2>&1 || true

    if nmcli -t -f NAME connection show | grep -Fxq "${DHCP_CON_NAME}"; then
        log "Updating NetworkManager connection ${DHCP_CON_NAME}."
        nmcli connection modify "${DHCP_CON_NAME}" connection.interface-name "${LAN_IF}" ipv4.method auto ipv6.method disabled connection.autoconnect yes
    else
        log "Creating NetworkManager connection ${DHCP_CON_NAME}."
        nmcli connection add type ethernet ifname "${LAN_IF}" con-name "${DHCP_CON_NAME}" ipv4.method auto ipv6.method disabled connection.autoconnect yes
    fi

    log "Bringing up NetworkManager connection ${DHCP_CON_NAME}."
    nmcli connection up "${DHCP_CON_NAME}" || true
}

write_dnsmasq_conf() {
    local internet_sharing="$1"
    load_state
    log "Writing dnsmasq configuration to ${DNSMASQ_CONF}."
    {
        echo "# Managed by nymea-experience-plugin-ngw"
        echo "interface=${LAN_IF}"
        echo "bind-dynamic"
        echo "port=0"
        echo "dhcp-authoritative"
        echo "dhcp-range=set:ngw-lan,${DHCP_START},${DHCP_END},255.255.255.0,${LEASE_TIME}"
        if [[ "${internet_sharing}" == "true" ]]; then
            echo "dhcp-option=tag:ngw-lan,option:router,${LAN_GW}"
            echo "dhcp-option=tag:ngw-lan,option:dns-server,1.1.1.1,9.9.9.9"
        else
            echo "dhcp-option=tag:ngw-lan,option:router"
            echo "dhcp-option=tag:ngw-lan,option:dns-server"
        fi
    } >"${DNSMASQ_CONF}"

    log "Validating dnsmasq configuration."
    dnsmasq --test --conf-file="${DNSMASQ_CONF}" >/dev/null
    log "Enabling and restarting dnsmasq."
    systemctl enable dnsmasq >/dev/null
    systemctl restart dnsmasq
}

port_rule() {
    local ports="$1"
    [[ -z "${ports}" ]] && return 0
    echo "    iifname \"${LAN_IF}\" tcp dport { ${ports} } accept"
}

write_nft_file() {
    local gateway_services="$1" internet_sharing="$2" tcp_ports="$3"
    load_state

    log "Writing nftables rules to ${NFT_FILE}: gateway_services=${gateway_services}, internet_sharing=${internet_sharing}, tcp_ports=${tcp_ports:-none}."
    cat >"${NFT_FILE}" <<NFT
#!/usr/sbin/nft -f
table inet nymea_ngw_lan {
  chain input {
    type filter hook input priority filter; policy accept;
    iifname "${LAN_IF}" ct state established,related accept
    iifname "${LAN_IF}" udp sport 68 udp dport 67 accept
    iifname "${LAN_IF}" ip protocol icmp accept
    iifname "${LAN_IF}" ip6 nexthdr icmpv6 accept
NFT

    if [[ "${gateway_services}" == "all-ports" ]]; then
        echo "    iifname \"${LAN_IF}\" accept" >>"${NFT_FILE}"
    elif [[ "${gateway_services}" == "configured-ports" ]]; then
        port_rule "${tcp_ports}" >>"${NFT_FILE}"
    fi

    cat >>"${NFT_FILE}" <<NFT
    iifname "${LAN_IF}" counter drop
  }

  chain forward {
    type filter hook forward priority filter; policy accept;
    ct state established,related accept
NFT

    if [[ "${internet_sharing}" == "true" ]]; then
        echo "    iifname \"${LAN_IF}\" oifname \"${WAN_IF}\" ip saddr ${LAN_CIDR} accept" >>"${NFT_FILE}"
    fi

    cat >>"${NFT_FILE}" <<NFT
    iifname "${LAN_IF}" counter drop
  }
}
NFT

    if [[ "${internet_sharing}" == "true" ]]; then
        cat >>"${NFT_FILE}" <<NFT

table ip nymea_ngw_lan_nat {
  chain postrouting {
    type nat hook postrouting priority srcnat; policy accept;
    ip saddr ${LAN_CIDR} oifname "${WAN_IF}" masquerade
  }
}
NFT
    fi
}

delete_compat_forward_rules() {
    local handle
    if ! nft list chain ip filter FORWARD >/dev/null 2>&1; then
        return 0
    fi

    while read -r handle; do
        [[ -n "${handle}" ]] || continue
        nft delete rule ip filter FORWARD handle "${handle}" 2>/dev/null || true
    done < <(nft -a list chain ip filter FORWARD | awk -v comment="${NFT_FORWARD_COMMENT}" '
        index($0, comment) {
            for (i = 1; i <= NF; i++) {
                if ($i == "handle" && (i + 1) <= NF)
                    print $(i + 1)
            }
        }')
}

install_compat_forward_rules() {
    local internet_sharing="$1"
    [[ "${internet_sharing}" == "true" ]] || return 0

    if ! nft list chain ip filter FORWARD >/dev/null 2>&1; then
        log "No ip filter FORWARD chain found; relying on nymea dedicated LAN forward chain."
        return 0
    fi

    log "Installing compatibility forward rules in ip filter FORWARD."
    nft insert rule ip filter FORWARD iifname "${WAN_IF}" oifname "${LAN_IF}" ip daddr "${LAN_CIDR}" ct state established,related accept comment "${NFT_FORWARD_COMMENT}"
    nft insert rule ip filter FORWARD iifname "${LAN_IF}" oifname "${WAN_IF}" ip saddr "${LAN_CIDR}" accept comment "${NFT_FORWARD_COMMENT}"
}

apply_firewall() {
    local internet_sharing="$1"
    log "Replacing nftables rules."
    delete_compat_forward_rules
    nft delete table inet nymea_ngw_lan 2>/dev/null || true
    nft delete table ip nymea_ngw_lan_nat 2>/dev/null || true
    nft -f "${NFT_FILE}"
    install_compat_forward_rules "${internet_sharing}"
}

set_ip_forwarding() {
    local internet_sharing="$1"
    if [[ "${internet_sharing}" == "true" ]]; then
        log "Enabling IPv4 forwarding."
        sysctl -w net.ipv4.ip_forward=1 >/dev/null
    else
        log "IPv4 forwarding unchanged because internet sharing is disabled."
    fi
}

disable_lan() {
    log "Removing dedicated LAN nftables and dnsmasq state."
    delete_compat_forward_rules
    nft delete table inet nymea_ngw_lan 2>/dev/null || true
    nft delete table ip nymea_ngw_lan_nat 2>/dev/null || true
    rm -f "${NFT_FILE}" "${DNSMASQ_CONF}"
    log "Restarting dnsmasq after removing dedicated LAN configuration."
    systemctl restart dnsmasq 2>/dev/null || true
    configure_nm_dhcp
}

apply_cmd() {
    local enabled="" gateway_services="" internet_sharing="" tcp_ports="" cidr=""
    local requested_lan_if="" requested_wan_if=""
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --enabled) enabled="$(bool_arg "$2")"; shift 2 ;;
            --gateway-services) gateway_services="$(services_arg "$2")"; shift 2 ;;
            --internet-sharing) internet_sharing="$(bool_arg "$2")"; shift 2 ;;
            --gateway-tcp-ports)
                shift
                if [[ $# -gt 0 && "$1" != --* ]]; then
                    tcp_ports="$1"
                    shift
                fi
                ;;
            --lan-if) LAN_IF="$2"; shift 2 ;;
            --wan-if) WAN_IF="$2"; shift 2 ;;
            --state-file) STATE_FILE="$2"; shift 2 ;;
            *) die "Unknown argument: $1" ;;
        esac
    done

    [[ -n "${enabled}" ]] || die "Missing --enabled."
    [[ -n "${gateway_services}" ]] || die "Missing --gateway-services."
    [[ -n "${internet_sharing}" ]] || die "Missing --internet-sharing."
    validate_ports "${tcp_ports}"

    log "Applying dedicated LAN configuration: enabled=${enabled}, gateway_services=${gateway_services}, internet_sharing=${internet_sharing}, lan_if=${LAN_IF}, wan_if=${WAN_IF}, tcp_ports=${tcp_ports:-none}."
    requested_lan_if="${LAN_IF}"
    requested_wan_if="${WAN_IF}"
    load_state
    LAN_IF="${requested_lan_if}"
    WAN_IF="${requested_wan_if}"
    resolve_wan_if "${internet_sharing}"
    if [[ -n "${LAN_CIDR:-}" ]]; then
        cidr="${LAN_CIDR}"
        log "Reusing LAN subnet ${cidr} from state."
    else
        log "Choosing a free private /24 subnet."
        cidr="$(choose_subnet)"
        log "Selected LAN subnet ${cidr}."
    fi

    write_state "${enabled}" "${gateway_services}" "${internet_sharing}" "${tcp_ports}" "${cidr}"

    if [[ "${enabled}" == "false" ]]; then
        disable_lan
        echo "Dedicated LAN disabled on ${LAN_IF}."
        return 0
    fi

    configure_nm_static "${cidr}"
    write_dnsmasq_conf "${internet_sharing}"
    write_nft_file "${gateway_services}" "${internet_sharing}" "${tcp_ports}"
    apply_firewall "${internet_sharing}"
    set_ip_forwarding "${internet_sharing}"
    log "Enabling boot reapply service nymea-experience-plugin-ngw-lan.service."
    systemctl enable nymea-experience-plugin-ngw-lan.service >/dev/null 2>&1 || true

    echo "Dedicated LAN applied on ${LAN_IF}."
}

reapply_cmd() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --state-file) STATE_FILE="$2"; shift 2 ;;
            *) die "Unknown argument: $1" ;;
        esac
    done
    [[ -r "${STATE_FILE}" ]] || exit 0
    log "Reapplying dedicated LAN configuration from ${STATE_FILE}."
    load_state
    apply_cmd --enabled "${ENABLED}" --gateway-services "${GATEWAY_SERVICES}" --internet-sharing "${INTERNET_SHARING}" --gateway-tcp-ports "${GATEWAY_TCP_PORTS:-}" --lan-if "${LAN_IF}" --wan-if "${WAN_IF}" --state-file "${STATE_FILE}"
}

status_cmd() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --state-file) STATE_FILE="$2"; shift 2 ;;
            *) die "Unknown argument: $1" ;;
        esac
    done
    [[ -r "${STATE_FILE}" ]] && cat "${STATE_FILE}" || echo "No state file found: ${STATE_FILE}"
    nmcli device status || true
    nft list chain ip filter FORWARD || true
    nft list table inet nymea_ngw_lan || true
    nft list table ip nymea_ngw_lan_nat || true
}

uninstall_cmd() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --state-file) STATE_FILE="$2"; shift 2 ;;
            *) die "Unknown argument: $1" ;;
        esac
    done
    load_state || true
    log "Uninstalling dedicated LAN helper state."
    disable_lan
    rm -f "${STATE_FILE}"
}

main() {
    local cmd="${1:-}"
    shift || true

    case "${cmd}" in
        -h|--help|help|"") usage; return 0 ;;
    esac

    need_root
    require_cmds

    case "${cmd}" in
        apply) apply_cmd "$@" ;;
        reapply) reapply_cmd "$@" ;;
        status) status_cmd "$@" ;;
        uninstall) uninstall_cmd "$@" ;;
        *) die "Unknown command: ${cmd}" ;;
    esac
}

main "$@"
