#!/usr/bin/env bash

WGCTL_BINARY="$(command -v wgctl)"

# ============================================
# Lifecycle
# ============================================

function cmd::test::on_load() {
  flag::register --destructive
  flag::register --section
  flag::register --fn
  flag::register --function
}

function cmd::test::help() {
  cat <<EOF
Usage: wgctl test [options]

Run the wgctl test suite.

Options:
  --destructive     Include tests that modify state (add/remove/block)
  --section <name>  Run only a specific section (list, rules, groups, audit, logs, fw)

Examples:
  wgctl test
  wgctl test --section rules
  wgctl test --destructive
EOF
}

# ============================================
# Test helpers
# ============================================

function cmd::test::run_cmd() {
  local desc="$1"
  local expected="${2:-}"
  shift 2

  local tmp exit_code
  tmp=$(mktemp)
  
  set +e # disable exit on error (return 1)

  timeout 30 "$WGCTL_BINARY" "$@" > "$tmp" 2>&1 &
  local pid=$!
  wait $pid
  exit_code=$?

  set -e # re-enable exit on error

  if [[ $exit_code -eq 124 ]]; then
    test::warn "${desc} (timed out after 30s)"
    rm -f "$tmp"
    return 1
  fi

  if [[ $exit_code -ne 0 ]]; then
    test::fail "${desc}"
    if [[ "${WGCTL_TEST_VERBOSE:-false}" == "true" ]]; then
      printf "    Output: %s\n" "$(cat "$tmp")"
    fi
    rm -f "$tmp"
    return 1
  fi

  if [[ -n "$expected" ]] && ! grep -qF "$expected" "$tmp"; then
    local actual
    actual=$(head -3 "$tmp" | tr '\n' ' ' | sed 's/  */ /g' | cut -c1-100)
    test::fail "${desc} (expected '${expected}', got: '${actual}')"
    rm -f "$tmp"
    return 1
  fi

  test::pass "$desc"
  rm -f "$tmp"
}

function cmd::test::run_cmd_fails() {
  local desc="$1"
  shift

  set +e # disable exit on error (return 1)

  local tmp exit_code
  tmp=$(mktemp)
  timeout 10 setsid "$WGCTL_BINARY" "$@" > "$tmp" 2>&1
  exit_code=$?

  set -e # re-enable exit on error

  rm -f "$tmp"

  if [[ $exit_code -eq 124 ]]; then
    test::warn "${desc} (timed out)"
    return 1
  fi

  if [[ $exit_code -eq 0 ]]; then
    test::fail "${desc} (expected failure but succeeded)"
    return 1
  fi

  test::pass "$desc"
}

function cmd::test::run_function() {
  local fn="$1"

  local namespace
  namespace=$(echo "$fn" | cut -d':' -f3)
  load_command "$namespace" 2>/dev/null || true

  test::reset
  log::section "Function Test: ${fn}"

  case "$fn" in
    cmd::block::run)       cmd::test::fn_block        ;;
    cmd::unblock::run)     cmd::test::fn_unblock      ;;
    cmd::remove::run)      cmd::test::fn_remove       ;;
    cmd::rule::assign)     cmd::test::fn_rule_assign  ;;
    cmd::rename::run)      cmd::test::fn_rename       ;;
    cmd::remove::run)      cmd::test::fn_remove       ;;
    cmd::unblock::run)     cmd::test::fn_unblock      ;;
    *)
      log::error "No function test defined for: ${fn}"
      return 1
      ;;
  esac

  test::summary
}

# ============================================
# Test sections
# ============================================

function cmd::test::section_list() {
  test::section "List"
  cmd::test::run_cmd "list"                          "WireGuard Clients"   list
  cmd::test::run_cmd "list --online"                 ""                    list --online
  cmd::test::run_cmd "list --offline"                ""                    list --offline
  cmd::test::run_cmd "list --blocked"                ""                    list --blocked
  cmd::test::run_cmd "list --type phone"             ""                    list --type phone
  cmd::test::run_cmd "list --type guest"             ""                    list --type guest
  cmd::test::run_cmd "list --detailed"               "Client:"             list --detailed
  cmd::test::run_cmd "list --name phone-nuno"        "phone-nuno"          list --name phone-nuno
}

function cmd::test::section_inspect() {
  test::section "Inspect"
  cmd::test::run_cmd "inspect --name phone-nuno"           "IP:"         inspect --name phone-nuno
  cmd::test::run_cmd "inspect --name nuno --type phone"    "IP:"         inspect --name nuno --type phone
  cmd::test::run_cmd "inspect --name phone-nuno --config"  "PrivateKey"  inspect --name phone-nuno --config
  cmd::test::run_cmd_fails "inspect nonexistent peer"                    inspect --name nonexistent-peer
}

function cmd::test::section_config() {
  test::section "Config & QR"
  cmd::test::run_cmd "config --name phone-nuno"       "PrivateKey"   config --name phone-nuno
  cmd::test::run_cmd "config --name nuno --type phone" "PrivateKey"  config --name nuno --type phone
  cmd::test::run_cmd "qr --name phone-nuno"           ""             qr --name phone-nuno
}

function cmd::test::section_rules() {
  test::section "Rules"
  cmd::test::run_cmd "rule list"                      "guest"        rule list
  cmd::test::run_cmd "rule show --name guest"         "Description"  rule show --name guest
  cmd::test::run_cmd "rule show --name user"          "Description"  rule show --name user
  cmd::test::run_cmd "rule show --name admin"         "Description"  rule show --name admin
  cmd::test::run_cmd_fails "rule show nonexistent"                   rule show --name nonexistent
}

function cmd::test::section_groups() {
  test::section "Groups"
  cmd::test::run_cmd "group list"                     "Groups"       group list
  cmd::test::run_cmd "group show --name family"       "Peers:"       group show --name family
  cmd::test::run_cmd_fails "group show nonexistent"                  group show --name nonexistent
}

function cmd::test::section_audit() {
  test::section "Audit"
  cmd::test::run_cmd "audit"                          "passed"       audit
  cmd::test::run_cmd "audit --peer phone-nuno"        "passed"       audit --peer phone-nuno
  cmd::test::run_cmd "audit --type phone"             "passed"       audit --type phone
}

function cmd::test::section_logs() {
  test::section "Logs"
  cmd::test::run_cmd "logs"                           "Activity"     logs
  cmd::test::run_cmd "logs --name phone-nuno"         "Activity"     logs --name phone-nuno
  cmd::test::run_cmd "logs --type guest"              "Activity"     logs --type guest
  cmd::test::run_cmd "logs --fw"                      "Activity"     logs --fw
  cmd::test::run_cmd "logs --wg"                      "Activity"     logs --wg
}

function cmd::test::section_fw() {
  test::section "Firewall"
  cmd::test::run_cmd "fw list"                        "FORWARD"      fw list
  cmd::test::run_cmd "fw list --peer phone-nuno"      ""             fw list --peer phone-nuno
  cmd::test::run_cmd "fw list --no-nflog"             ""             fw list --no-nflog
  cmd::test::run_cmd "fw list --no-accept"            ""             fw list --no-accept
  cmd::test::run_cmd "fw list --no-drop"              ""             fw list --no-drop
  cmd::test::run_cmd "fw nat"                         "PREROUTING"   fw nat
  cmd::test::run_cmd "fw count"                       "TOTAL"        fw count
}
function cmd::test::section_net() {
  test::section "Net"

  "$WGCTL_BINARY" net rm --name test-svc --force > /dev/null 2>&1 || true

  cmd::test::run_cmd "net add service"      "added"    \
    net add --name test-svc --ip 10.0.0.99 --desc "Test service"
  cmd::test::run_cmd "net add port"         "Added"    \
    net add --name test-svc:web --port 9999:tcp
  cmd::test::run_cmd "net list"             "test-svc" \
    net list
  cmd::test::run_cmd "net list --detailed"  "web"      \
    net list --detailed
  cmd::test::run_cmd "net show"             "9999"     \
    net show --name test-svc
  cmd::test::run_cmd "net rm port"          "Removed"  \
    net rm --name test-svc:web --force
  cmd::test::run_cmd "net add port again"   "Added"    \
    net add --name test-svc:web --port 9999:tcp
  cmd::test::run_cmd "net rm all ports"     "Removed"  \
    net rm --name test-svc:ports --force
  cmd::test::run_cmd "net rm service"       "Removed"  \
    net rm --name test-svc --force
  cmd::test::run_cmd_fails "net show nonexistent" \
    net show --name nonexistent-svc
  cmd::test::run_cmd_fails "net add port no service" \
    net add --name nonexistent:web --port 80:tcp
}

function cmd::test::section_destructive() {
  test::section "Destructive (modifying state)"

  # ── Cleanup from any previous failed run ──
  "$WGCTL_BINARY" remove --name phone-testunit  --force > /dev/null 2>&1 || true
  "$WGCTL_BINARY" group  remove --name testgroup  --force > /dev/null 2>&1 || true
  "$WGCTL_BINARY" group  remove --name testgroup2 --force > /dev/null 2>&1 || true
  "$WGCTL_BINARY" net    rm --name test-block-svc --force > /dev/null 2>&1 || true
  "$WGCTL_BINARY" unblock --name phone-testunit > /dev/null 2>&1 || true
  
  # ── Add test peer ──────────────────────────
  cmd::test::run_cmd "add phone peer" "added successfully" \
    add --name testunit --type phone

  # ── Direct block/unblock ───────────────────
  cmd::test::run_cmd "block peer"          "blocked"   block --name phone-testunit
  cmd::test::run_cmd "list shows blocked"  "blocked"   list --blocked
  cmd::test::run_cmd "unblock peer"        "unblocked" unblock --name phone-testunit

  # ── Specific IP block/unblock ──────────────
  cmd::test::run_cmd "block peer --ip" "blocked for" \
    block --name phone-testunit --ip 10.0.0.99
  cmd::test::run_cmd "list shows restricted" "restricted" \
    list --name phone-testunit
  cmd::test::run_cmd "unblock peer --ip" "unblocked" \
    unblock --name phone-testunit --ip 10.0.0.99

  # ── Service block/unblock ──────────────────
  "$WGCTL_BINARY" net add --name test-block-svc \
    --ip 10.0.0.99 > /dev/null 2>&1
  "$WGCTL_BINARY" net add --name test-block-svc:web \
    --port 9999:tcp > /dev/null 2>&1
  cmd::test::run_cmd "block peer --service (ip)" "blocked" \
    block --name phone-testunit --service test-block-svc
  cmd::test::run_cmd "block already blocked service" "already" \
    block --name phone-testunit --service test-block-svc
  cmd::test::run_cmd "unblock peer --service (ip)" "unblocked" \
    unblock --name phone-testunit --service test-block-svc
  cmd::test::run_cmd "unblock not blocked service" "not blocked" \
    unblock --name phone-testunit --service test-block-svc
  cmd::test::run_cmd "block peer --service (port)" "blocked" \
    block --name phone-testunit --service test-block-svc:web
  cmd::test::run_cmd "unblock peer --service (port)" "unblocked" \
    unblock --name phone-testunit --service test-block-svc:web
  "$WGCTL_BINARY" net rm --name test-block-svc --force > /dev/null 2>&1 || true

  # ── Rule assign/unassign ───────────────────
  cmd::test::run_cmd "rule assign" "Assigned" \
    rule assign --name user --peer phone-testunit
  cmd::test::run_cmd "rule unassign" "Unassigned" \
    rule unassign --peer phone-testunit
  "$WGCTL_BINARY" rule assign --name user --peer phone-testunit \
    > /dev/null 2>&1 || true

  # ── Group basic operations ─────────────────
  cmd::test::run_cmd "group add" "created" \
    group add --name testgroup --desc "Test group"
  cmd::test::run_cmd "group peer add" "Added" \
    group peer add --name testgroup --peer phone-testunit
  cmd::test::run_cmd "group block" "blocked" \
    group block --name testgroup
  cmd::test::run_cmd "group unblock" "unblocked" \
    group unblock --name testgroup

  # ── M:N group block tracking ───────────────
  "$WGCTL_BINARY" group add --name testgroup2 \
    --desc "Test group 2" > /dev/null 2>&1
  "$WGCTL_BINARY" group peer add --name testgroup2 \
    --peer phone-testunit > /dev/null 2>&1

  cmd::test::run_cmd "group block first group"  "blocked" \
    group block --name testgroup
  cmd::test::run_cmd "group block second group" "blocked" \
    group block --name testgroup2

  "$WGCTL_BINARY" group unblock --name testgroup > /dev/null 2>&1
  cmd::test::run_cmd "peer stays blocked after partial unblock" "blocked" \
    list --blocked

  "$WGCTL_BINARY" group unblock --name testgroup2 > /dev/null 2>&1
  cmd::test::run_cmd "peer unblocked after all groups unblock" "phone-testunit" \
    list --allowed

  # ── Direct block overrides group block ─────
  "$WGCTL_BINARY" group block --name testgroup > /dev/null 2>&1
  cmd::test::run_cmd "direct unblock overrides group block" "unblocked" \
    unblock --name phone-testunit

  # ── Cleanup ────────────────────────────────
  cmd::test::run_cmd "group remove" "removed" \
    group remove --name testgroup --force
  "$WGCTL_BINARY" group remove --name testgroup2 --force > /dev/null 2>&1 || true

  cmd::test::run_cmd "remove phone peer" "removed" \
    remove --name phone-testunit --force
}

# ============================================
# Function Blocks
# ============================================

function cmd::test::fn_block() {
  test::section "cmd::block::run"

  # Setup
  "$WGCTL_BINARY" remove --name phone-testunit --force > /dev/null 2>&1 || true
  "$WGCTL_BINARY" add --name testunit --type phone > /dev/null 2>&1

  # Tests
  cmd::test::run_cmd "block peer"              "blocked"  block --name phone-testunit
  cmd::test::run_cmd "block already blocked"   "already"  block --name phone-testunit

  "$WGCTL_BINARY" unblock --name phone-testunit > /dev/null 2>&1 || true
  cmd::test::run_cmd "block with --type" "blocked" block --name testunit --type phone

  cmd::test::run_cmd_fails "block nonexistent"            block --name truly-nonexistent-xyz

  # Cleanup
  "$WGCTL_BINARY" unblock --name phone-testunit > /dev/null 2>&1 || true
  "$WGCTL_BINARY" remove --name phone-testunit --force > /dev/null 2>&1 || true
}

function cmd::test::fn_remove() {
  test::section "cmd::remove::run"

  # Setup
  "$WGCTL_BINARY" remove --name phone-testunit --force \
    > /dev/null 2>&1 || true
  "$WGCTL_BINARY" add --name testunit --type phone \
    > /dev/null 2>&1

  # Tests
  # Skip the interactive prompt test — it hangs waiting for input
  # cmd::test::run_cmd_fails "remove without --force" \
  #   remove --name phone-testunit
  cmd::test::run_cmd "remove with --force" "removed" \
    remove --name phone-testunit --force
  cmd::test::run_cmd_fails "remove nonexistent" \
    remove --name nonexistent-peer --force
  cmd::test::run_cmd_fails "remove missing --name" \
    remove --force

  # Cleanup already done by tests
}

function cmd::test::fn_rename() {
  test::section "cmd::rename::run"

  # Setup
  "$WGCTL_BINARY" remove --name phone-testunit --force \
    > /dev/null 2>&1 || true
  "$WGCTL_BINARY" remove --name phone-testunit2 --force \
    > /dev/null 2>&1 || true
  "$WGCTL_BINARY" add --name testunit --type phone \
    > /dev/null 2>&1

  # Tests
  cmd::test::run_cmd "rename peer" "renamed" \
    rename --name phone-testunit --new-name phone-testunit2
  cmd::test::run_cmd_fails "rename to existing" \
    rename --name phone-testunit2 --new-name phone-nuno
  cmd::test::run_cmd_fails "rename nonexistent" \
    rename --name phone-nonexistent --new-name phone-testunit
  cmd::test::run_cmd_fails "rename missing --new-name" \
    rename --name phone-testunit2

  # Cleanup
  "$WGCTL_BINARY" remove --name phone-testunit2 --force \
    > /dev/null 2>&1 || true
}

function cmd::test::fn_unblock() {
  test::section "cmd::unblock::run"

  # Setup — add and block a peer
  "$WGCTL_BINARY" remove --name phone-testunit --force > /dev/null 2>&1 || true
  "$WGCTL_BINARY" add --name testunit --type phone > /dev/null 2>&1
  "$WGCTL_BINARY" block --name phone-testunit > /dev/null 2>&1

  # Tests
  cmd::test::run_cmd "unblock peer"              "unblocked"  unblock --name phone-testunit
  cmd::test::run_cmd "unblock not blocked"       "not blocked" unblock --name phone-testunit
  cmd::test::run_cmd_fails "unblock nonexistent"              unblock --name nonexistent-peer
  cmd::test::run_cmd_fails "unblock missing --name"           unblock

  # Cleanup
  "$WGCTL_BINARY" remove --name phone-testunit --force > /dev/null 2>&1 || true
}

function cmd::test::fn_rule_assign() {
  test::section "cmd::rule::assign"
  
  "$WGCTL_BINARY" remove --name phone-testunit --force > /dev/null 2>&1 || true
  "$WGCTL_BINARY" add --name testunit --type phone > /dev/null 2>&1
  
  cmd::test::run_cmd "rule assign" "Assigned" \
    rule assign --name admin --peer phone-testunit

  "$WGCTL_BINARY" remove --name phone-testunit --force > /dev/null 2>&1 || true
}

# ============================================
# Run
# ============================================

function cmd::test::run() {
  local destructive=false section=""
  local fn=""

  while [[ $# -gt 0 ]]; do
    case "$1" in
      --destructive) destructive=true; shift ;;
      --section)
        util::require_flag "--section" "${2:-}" || return 1
        section="$2"; shift 2 
      ;;
      --fn|--function) fn="$2"; shift 2 ;;
      --verbose|-v) WGCTL_TEST_VERBOSE=true; shift ;;
      --help) cmd::test::help; return ;;
      *)
        log::error "Unknown flag: $1"
        return 1
        ;;
    esac
  done

  # After flag parsing:
  if [[ -n "$fn" ]]; then
    cmd::test::run_function "$fn"
    return
  fi

  test::reset
  log::section "wgctl Test Suite"

  if [[ -n "$section" ]]; then
    case "$section" in
      list)         cmd::test::section_list         ;;
      inspect)      cmd::test::section_inspect      ;;
      config)       cmd::test::section_config       ;;
      rules)        cmd::test::section_rules        ;;
      groups)       cmd::test::section_groups       ;;
      audit)        cmd::test::section_audit        ;;
      logs)         cmd::test::section_logs         ;;
      fw)           cmd::test::section_fw           ;;
      net)          cmd::test::section_net          ;;
      destructive)  cmd::test::section_destructive  ;;
      *)
        log::error "Unknown section: $section"
        return 1
        ;;
    esac
  else
    cmd::test::section_list
    cmd::test::section_inspect
    cmd::test::section_config
    cmd::test::section_rules
    cmd::test::section_groups
    cmd::test::section_audit
    cmd::test::section_logs
    cmd::test::section_fw
  fi

  if $destructive; then
    cmd::test::section_destructive
  fi

  test::summary
}