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

ROOT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)
# shellcheck source=tests/lib/assert.sh
source "${ROOT_DIR}/tests/lib/assert.sh"

status_of() {
  set +e
  "$@" >/dev/null 2>&1
  local rc=$?
  set -e
  printf '%s\n' "${rc}"
}

maybe_enter_namespace() {
  if (( EUID == 0 )); then
    return 0
  fi

  if [[ ${IPF_IN_NAMESPACE:-0} == 1 ]]; then
    return 0
  fi

  if command -v unshare >/dev/null 2>&1 && unshare -Urn true >/dev/null 2>&1; then
    exec unshare -Urn env IPF_IN_NAMESPACE=1 bash "$0"
  fi

  printf 'SKIP: root 或 unshare 不可用，跳过集成测试。\n'
  exit 0
}

maybe_enter_namespace

for bin in ip iptables iptables-save iptables-restore nsenter python3; do
  if ! command -v "${bin}" >/dev/null 2>&1; then
    printf 'SKIP: %s 不可用，跳过集成测试。\n' "${bin}"
    exit 0
  fi
done

ip link set lo up >/dev/null 2>&1 || true

TMP_DIR=$(mktemp -d)
BACKUP_V4="${TMP_DIR}/iptables.v4.bak"
BACKUP_V6="${TMP_DIR}/iptables.v6.bak"
BACKUP_IPV4_FORWARD=$(cat /proc/sys/net/ipv4/ip_forward 2>/dev/null || printf '')
BACKUP_IPV6_FORWARD=$(cat /proc/sys/net/ipv6/conf/all/forwarding 2>/dev/null || printf '')
PERSIST_RULES_V4="${TMP_DIR}/rules.v4"
PERSIST_RULES_V6="${TMP_DIR}/rules.v6"
PERSIST_FIXTURE_LOG="${TMP_DIR}/persist.log"
TRAFFIC_LISTEN_PORT_V4=65432
TRAFFIC_TARGET_PORT_V4=18080
TRAFFIC_ROUTER_EDGE_IP4=10.23.0.1
TRAFFIC_CLIENT_IP4=10.23.0.2
TRAFFIC_ROUTER_BACKEND_IP4=10.24.0.1
TRAFFIC_TARGET_IP4=10.24.0.2
TRAFFIC_LISTEN_PORT_V6=65433
TRAFFIC_TARGET_PORT_V6=18081
TRAFFIC_ROUTER_EDGE_IP6=2001:db8:23::1
TRAFFIC_CLIENT_IP6=2001:db8:23::2
TRAFFIC_ROUTER_BACKEND_IP6=2001:db8:24::1
TRAFFIC_TARGET_IP6=2001:db8:24::2
TRAFFIC_IPV6_READY=0
TRAFFIC_CLIENT_PID=''
TRAFFIC_BACKEND_PID=''
TRAFFIC_SERVER_PID=''

if command -v ip6tables >/dev/null 2>&1 && command -v ip6tables-save >/dev/null 2>&1 && command -v ip6tables-restore >/dev/null 2>&1; then
  TRAFFIC_IPV6_READY=1
fi

restore_runtime_from_backup() {
  if [[ -f ${BACKUP_V4} ]]; then
    iptables-restore <"${BACKUP_V4}" >/dev/null 2>&1 || true
  fi
  if [[ -f ${BACKUP_V6} ]] && command -v ip6tables-restore >/dev/null 2>&1; then
    ip6tables-restore <"${BACKUP_V6}" >/dev/null 2>&1 || true
  fi
  if [[ -n ${BACKUP_IPV4_FORWARD} ]]; then
    printf '%s\n' "${BACKUP_IPV4_FORWARD}" >/proc/sys/net/ipv4/ip_forward 2>/dev/null || true
  fi
  if [[ -n ${BACKUP_IPV6_FORWARD} ]]; then
    printf '%s\n' "${BACKUP_IPV6_FORWARD}" >/proc/sys/net/ipv6/conf/all/forwarding 2>/dev/null || true
  fi
}

traffic_stop_server() {
  [[ -n ${TRAFFIC_SERVER_PID} ]] || return 0
  kill "${TRAFFIC_SERVER_PID}" >/dev/null 2>&1 || true
  wait "${TRAFFIC_SERVER_PID}" 2>/dev/null || true
  TRAFFIC_SERVER_PID=''
}

traffic_cleanup_fixture() {
  traffic_stop_server
  for link in ipfc0 ipfb0; do
    ip link del "${link}" >/dev/null 2>&1 || true
  done
  for pid in "${TRAFFIC_CLIENT_PID}" "${TRAFFIC_BACKEND_PID}"; do
    [[ -n ${pid} ]] || continue
    kill "${pid}" >/dev/null 2>&1 || true
    wait "${pid}" 2>/dev/null || true
  done
  TRAFFIC_CLIENT_PID=''
  TRAFFIC_BACKEND_PID=''
}

cleanup() {
  traffic_cleanup_fixture
  restore_runtime_from_backup
  rm -rf "${TMP_DIR}"
}
trap cleanup EXIT

traffic_ns_exec() {
  local pid=$1
  shift
  nsenter -t "${pid}" -n "$@"
}

traffic_setup_fixture() {
  traffic_cleanup_fixture

  unshare -n bash -c 'sleep 1000' &
  TRAFFIC_CLIENT_PID=$!
  unshare -n bash -c 'sleep 1000' &
  TRAFFIC_BACKEND_PID=$!

  ip link add ipfc0 type veth peer name ipfc1
  ip link add ipfb0 type veth peer name ipfb1
  ip link set ipfc1 netns "${TRAFFIC_CLIENT_PID}"
  ip link set ipfb1 netns "${TRAFFIC_BACKEND_PID}"

  ip addr add "${TRAFFIC_ROUTER_EDGE_IP4}/24" dev ipfc0
  ip addr add "${TRAFFIC_ROUTER_BACKEND_IP4}/24" dev ipfb0
  ip link set ipfc0 up
  ip link set ipfb0 up

  traffic_ns_exec "${TRAFFIC_CLIENT_PID}" ip link set lo up
  traffic_ns_exec "${TRAFFIC_CLIENT_PID}" ip addr add "${TRAFFIC_CLIENT_IP4}/24" dev ipfc1
  traffic_ns_exec "${TRAFFIC_CLIENT_PID}" ip link set ipfc1 up
  traffic_ns_exec "${TRAFFIC_CLIENT_PID}" ip route add default via "${TRAFFIC_ROUTER_EDGE_IP4}"

  traffic_ns_exec "${TRAFFIC_BACKEND_PID}" ip link set lo up
  traffic_ns_exec "${TRAFFIC_BACKEND_PID}" ip addr add "${TRAFFIC_TARGET_IP4}/24" dev ipfb1
  traffic_ns_exec "${TRAFFIC_BACKEND_PID}" ip link set ipfb1 up
  traffic_ns_exec "${TRAFFIC_BACKEND_PID}" ip route add default via "${TRAFFIC_ROUTER_BACKEND_IP4}"

  if [[ ${TRAFFIC_IPV6_READY} == 1 && -n ${BACKUP_IPV6_FORWARD} ]]; then
    ip addr add "${TRAFFIC_ROUTER_EDGE_IP6}/64" dev ipfc0 nodad
    ip addr add "${TRAFFIC_ROUTER_BACKEND_IP6}/64" dev ipfb0 nodad
    traffic_ns_exec "${TRAFFIC_CLIENT_PID}" ip addr add "${TRAFFIC_CLIENT_IP6}/64" dev ipfc1 nodad
    traffic_ns_exec "${TRAFFIC_CLIENT_PID}" ip -6 route add default via "${TRAFFIC_ROUTER_EDGE_IP6}"
    traffic_ns_exec "${TRAFFIC_BACKEND_PID}" ip addr add "${TRAFFIC_TARGET_IP6}/64" dev ipfb1 nodad
    traffic_ns_exec "${TRAFFIC_BACKEND_PID}" ip -6 route add default via "${TRAFFIC_ROUTER_BACKEND_IP6}"
    printf '1\n' >/proc/sys/net/ipv6/conf/all/forwarding
  fi

  printf '1\n' >/proc/sys/net/ipv4/ip_forward
}

traffic_start_server() {
  local family=$1
  local host=$2
  local port=$3
  traffic_stop_server
  traffic_ns_exec "${TRAFFIC_BACKEND_PID}" python3 -u - "${family}" "${host}" "${port}" <<'PY' &
import socket
import sys

family = socket.AF_INET6 if sys.argv[1] == '6' else socket.AF_INET
host = sys.argv[2]
port = int(sys.argv[3])
bind_addr = (host, port, 0, 0) if family == socket.AF_INET6 else (host, port)
with socket.socket(family, socket.SOCK_STREAM) as sock:
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind(bind_addr)
    sock.listen(1)
    conn, _addr = sock.accept()
    with conn:
        conn.recv(1024)
        conn.sendall(b'HTTP/1.0 200 OK\r\nContent-Length: 2\r\n\r\nOK')
PY
  TRAFFIC_SERVER_PID=$!
  sleep 0.2
}

traffic_client_http() {
  local family=$1
  local host=$2
  local port=$3
  traffic_ns_exec "${TRAFFIC_CLIENT_PID}" python3 - "${family}" "${host}" "${port}" <<'PY'
import socket
import sys

family = socket.AF_INET6 if sys.argv[1] == '6' else socket.AF_INET
host = sys.argv[2]
port = int(sys.argv[3])
remote = (host, port, 0, 0) if family == socket.AF_INET6 else (host, port)
chunks = []
with socket.socket(family, socket.SOCK_STREAM) as sock:
    sock.settimeout(2)
    sock.connect(remote)
    sock.sendall(b'GET / HTTP/1.0\r\nHost: ipf-test\r\n\r\n')
    while True:
        data = sock.recv(4096)
        if not data:
            break
        chunks.append(data)
sys.stdout.buffer.write(b''.join(chunks))
PY
}

traffic_request_ok() {
  local family=$1
  local listen_host=$2
  local listen_port=$3
  local target_host=$4
  local target_port=$5
  local output
  traffic_start_server "${family}" "${target_host}" "${target_port}"
  if ! output=$(traffic_client_http "${family}" "${listen_host}" "${listen_port}"); then
    traffic_stop_server
    return 1
  fi
  wait "${TRAFFIC_SERVER_PID}"
  TRAFFIC_SERVER_PID=''
  printf '%s\n' "${output}"
}

iptables-save >"${BACKUP_V4}"
if [[ ${TRAFFIC_IPV6_READY} == 1 ]]; then
  ip6tables-save >"${BACKUP_V6}"
fi

traffic_setup_fixture
iptables -P FORWARD DROP
if [[ ${TRAFFIC_IPV6_READY} == 1 && -n ${BACKUP_IPV6_FORWARD} ]]; then
  ip6tables -P FORWARD DROP
fi

export IPF_STORAGE_DIR="${TMP_DIR}/storage"
export IPF_STORAGE_DB="${IPF_STORAGE_DIR}/rules.db"
export IPF_LOCK_FILE="${IPF_STORAGE_DIR}/.lock"
export IPF_SKIP_ENV_CHECK=1
export IPF_FORCE_PLAIN_UI=1
export NETFILTER_PERSISTENT_BIN="${ROOT_DIR}/tests/mocks/persist-fixture.sh"
export PERSIST_RULES_V4
export PERSIST_RULES_V6
export PERSIST_FIXTURE_LOG

# shellcheck source=lib/common.sh
source "${ROOT_DIR}/lib/common.sh"
# shellcheck source=lib/storage.sh
source "${ROOT_DIR}/lib/storage.sh"
# shellcheck source=lib/persist.sh
source "${ROOT_DIR}/lib/persist.sh"
# shellcheck source=lib/iptables_ops.sh
source "${ROOT_DIR}/lib/iptables_ops.sh"
# shellcheck source=lib/rules_mgr.sh
source "${ROOT_DIR}/lib/rules_mgr.sh"

uuid_v4=$("${ROOT_DIR}/iptables-forward.sh" --batch add tcp "${TRAFFIC_LISTEN_PORT_V4}" "${TRAFFIC_TARGET_IP4}" "${TRAFFIC_TARGET_PORT_V4}" 4 'integration-v4')
assert_contains "$(iptables-save)" "MGMT:${uuid_v4}" 'IPv4 rule should appear in iptables-save output'
assert_contains "$("${ROOT_DIR}/iptables-forward.sh" --batch list)" "uuid=${uuid_v4}" 'batch list should include managed rule'
assert_file_contains "${PERSIST_RULES_V4}" "MGMT:${uuid_v4}" 'persist save should write IPv4 rules snapshot'
assert_file_contains "${PERSIST_FIXTURE_LOG}" 'persist-fixture.sh save' 'adding a rule should call persist save'
assert_contains "$(iptables -S FORWARD)" "--sport ${TRAFFIC_TARGET_PORT_V4}" 'IPv4 rule should include reply-direction FORWARD allow'
assert_contains "$(traffic_request_ok 4 "${TRAFFIC_ROUTER_EDGE_IP4}" "${TRAFFIC_LISTEN_PORT_V4}" "${TRAFFIC_TARGET_IP4}" "${TRAFFIC_TARGET_PORT_V4}")" 'OK' 'managed IPv4 rule should forward actual TCP traffic'

iptables -t nat -D PREROUTING \
  -p tcp --dport "${TRAFFIC_LISTEN_PORT_V4}" \
  -j DNAT --to-destination "${TRAFFIC_TARGET_IP4}:${TRAFFIC_TARGET_PORT_V4}" \
  -m comment --comment "MGMT:${uuid_v4}"
list_output=$(cmd_list 0)
[[ ${list_output} =~ [[:space:]]![[:space:]] ]] || fail 'partial runtime loss should show degraded health in cmd_list'
assert_status 1 "$(status_of traffic_client_http 4 "${TRAFFIC_ROUTER_EDGE_IP4}" "${TRAFFIC_LISTEN_PORT_V4}")" 'missing PREROUTING rule should break actual TCP forwarding'

persist_reload
list_output=$(cmd_list 0)
[[ ${list_output} =~ [[:space:]]✓[[:space:]] ]] || fail 'persist_reload should restore healthy status in cmd_list'
assert_contains "$(traffic_request_ok 4 "${TRAFFIC_ROUTER_EDGE_IP4}" "${TRAFFIC_LISTEN_PORT_V4}" "${TRAFFIC_TARGET_IP4}" "${TRAFFIC_TARGET_PORT_V4}")" 'OK' 'persist_reload should restore actual TCP forwarding'

ipt_remove_rule "${uuid_v4}" tcp "${TRAFFIC_LISTEN_PORT_V4}" "${TRAFFIC_TARGET_IP4}" "${TRAFFIC_TARGET_PORT_V4}" 4
assert_status 1 "$(status_of grep -F "MGMT:${uuid_v4}" <(iptables-save))" 'manual runtime removal should clear managed IPv4 rule'
assert_status 1 "$(status_of traffic_client_http 4 "${TRAFFIC_ROUTER_EDGE_IP4}" "${TRAFFIC_LISTEN_PORT_V4}")" 'runtime rule removal should stop actual TCP forwarding'
persist_reload
assert_file_contains "${PERSIST_FIXTURE_LOG}" 'persist-fixture.sh reload' 'persist_reload should call persistence wrapper'
assert_contains "$(iptables-save)" "MGMT:${uuid_v4}" 'persist_reload should restore IPv4 rule from snapshot'
assert_contains "$(traffic_request_ok 4 "${TRAFFIC_ROUTER_EDGE_IP4}" "${TRAFFIC_LISTEN_PORT_V4}" "${TRAFFIC_TARGET_IP4}" "${TRAFFIC_TARGET_PORT_V4}")" 'OK' 'restored IPv4 rule should resume actual TCP forwarding'

iptables -A FORWARD -p tcp --dport 65000 -j ACCEPT

"${ROOT_DIR}/iptables-forward.sh" --batch delete "${uuid_v4}"
assert_status 1 "$(status_of grep -F "MGMT:${uuid_v4}" <(iptables-save))" 'deleted IPv4 rule should disappear from iptables-save'
assert_status 0 "$(status_of grep -F -- '--dport 65000 -j ACCEPT' <(iptables-save))" 'unmanaged rule should remain after deleting managed rule'
assert_not_contains "$(cat "${PERSIST_RULES_V4}")" "MGMT:${uuid_v4}" 'deleting IPv4 rule should refresh persisted snapshot'
assert_status 1 "$(status_of traffic_client_http 4 "${TRAFFIC_ROUTER_EDGE_IP4}" "${TRAFFIC_LISTEN_PORT_V4}")" 'deleting managed IPv4 rule should stop actual TCP forwarding'

if [[ ${TRAFFIC_IPV6_READY} == 1 && -n ${BACKUP_IPV6_FORWARD} ]]; then
  uuid_v6=$("${ROOT_DIR}/iptables-forward.sh" --batch add tcp "${TRAFFIC_LISTEN_PORT_V6}" "${TRAFFIC_TARGET_IP6}" "${TRAFFIC_TARGET_PORT_V6}" 6 'integration-v6')
  assert_contains "$(ip6tables-save)" "MGMT:${uuid_v6}" 'IPv6 rule should appear in ip6tables-save output'
  assert_file_contains "${PERSIST_RULES_V6}" "MGMT:${uuid_v6}" 'persist save should write IPv6 rules snapshot'
  assert_contains "$(ip6tables -S FORWARD)" "--sport ${TRAFFIC_TARGET_PORT_V6}" 'IPv6 rule should include reply-direction FORWARD allow'
  assert_contains "$(traffic_request_ok 6 "${TRAFFIC_ROUTER_EDGE_IP6}" "${TRAFFIC_LISTEN_PORT_V6}" "${TRAFFIC_TARGET_IP6}" "${TRAFFIC_TARGET_PORT_V6}")" 'OK' 'managed IPv6 rule should forward actual TCP traffic'
  ipt_remove_rule "${uuid_v6}" tcp "${TRAFFIC_LISTEN_PORT_V6}" "${TRAFFIC_TARGET_IP6}" "${TRAFFIC_TARGET_PORT_V6}" 6
  assert_status 1 "$(status_of grep -F "MGMT:${uuid_v6}" <(ip6tables-save))" 'manual runtime removal should clear managed IPv6 rule'
  assert_status 1 "$(status_of traffic_client_http 6 "${TRAFFIC_ROUTER_EDGE_IP6}" "${TRAFFIC_LISTEN_PORT_V6}")" 'runtime IPv6 rule removal should stop actual TCP forwarding'
  persist_reload
  assert_contains "$(ip6tables-save)" "MGMT:${uuid_v6}" 'persist_reload should restore IPv6 rule from snapshot'
  assert_contains "$(traffic_request_ok 6 "${TRAFFIC_ROUTER_EDGE_IP6}" "${TRAFFIC_LISTEN_PORT_V6}" "${TRAFFIC_TARGET_IP6}" "${TRAFFIC_TARGET_PORT_V6}")" 'OK' 'persist_reload should restore actual IPv6 TCP forwarding'
  "${ROOT_DIR}/iptables-forward.sh" --batch delete "${uuid_v6}"
  assert_status 1 "$(status_of grep -F "MGMT:${uuid_v6}" <(ip6tables-save))" 'deleted IPv6 rule should disappear from ip6tables-save'
  assert_not_contains "$(cat "${PERSIST_RULES_V6}")" "MGMT:${uuid_v6}" 'deleting IPv6 rule should refresh persisted snapshot'
  assert_status 1 "$(status_of traffic_client_http 6 "${TRAFFIC_ROUTER_EDGE_IP6}" "${TRAFFIC_LISTEN_PORT_V6}")" 'deleting managed IPv6 rule should stop actual TCP forwarding'
fi

pass 'test_integration.sh'