#!/usr/bin/env bash
# identity.command.sh — manage peer identities
#
# Subcommands:
#   wgctl identity list
#   wgctl identity show    --name <name>
#   wgctl identity add     --name <name> --peer <peer>
#   wgctl identity remove  --name <name>
#   wgctl identity migrate [--dry-run]

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

function cmd::identity::on_load() {
  load_module identity
  load_module policy
 
  flag::register --name
  flag::register --peer
  flag::register --dry-run
  flag::register --force
  # rule subcommand flags
  flag::register --rule
  # options subcommand flags
  flag::register --policy
  flag::register --set-strict-rule
  flag::register --unset-strict-rule
  flag::register --set-auto-apply
  flag::register --unset-auto-apply
  flag::register --field
  flag::register --value
}

# ============================================
# Help
# ============================================

function cmd::identity::help() {
  cat <<EOF
Usage: wgctl identity <subcommand> [options]
 
Manage peer identities.
 
Subcommands:
  list                          List all identities
  show   --name <name>          Show identity details and device status
  add    --name <name>          Manually attach a peer to an identity
         --peer <peer>
  remove --name <name>          Remove identity and all associated peers
  migrate [--dry-run]           Create identities from existing peer names
 
  rule assign   --name <name>   Assign a rule to an identity
                --rule <rule>
  rule unassign --name <name>   Remove rule from an identity
  rule show     --name <name>   Show current identity rule
 
  options --name <name>         Set identity options
          [--policy <policy>]
          [--set-strict-rule | --unset-strict-rule]
          [--set-auto-apply  | --unset-auto-apply]
 
Examples:
  wgctl identity list
  wgctl identity show --name nuno
  wgctl identity rule assign --name nuno --rule admin
  wgctl identity rule unassign --name nuno
  wgctl identity options --name guests-identity --policy guest
  wgctl identity options --name nuno --set-strict-rule
EOF
}

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

function cmd::identity::run() {
  local subcmd="${1:-list}"
  shift || true

  case "$subcmd" in
    list)    cmd::identity::_list    "$@" ;;
    show)    cmd::identity::_show    "$@" ;;
    add)     cmd::identity::_add     "$@" ;;
    remove)  cmd::identity::_remove  "$@" ;;
    migrate) cmd::identity::_migrate "$@" ;;
    rule)    cmd::identity::_rule    "$@" ;;
    options) cmd::identity::_options "$@" ;;
    --help)  cmd::identity::help          ;;
    *)
      log::error "Unknown subcommand '${subcmd}'. Available: list, show, add, remove, migrate, rule, options"
      return 1
      ;;
  esac
}

# ============================================
# Subcommands
# ============================================

function cmd::identity::_list() {
  local data
  data=$(identity::list_data)
 
  if [[ -z "$data" ]]; then
    log::info "No identities found. Run 'wgctl identity migrate' to create from existing peers."
    return 0
  fi
 
  echo ""
  while IFS='|' read -r name peer_count types rules policy; do
    local rules_display
    rules_display=$(echo "$rules" | sed 's/,/, /g')
    ui::identity::list_row_compact "$name" "$peer_count" "$rules_display" "$policy"
  done <<< "$data"
  echo ""
}

function cmd::identity::_show() {
  local name=""
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --name) name="$2"; shift 2 ;;
      --help) cmd::identity::help; return ;;
      *) log::error "Unknown flag: $1"; return 1 ;;
    esac
  done
 
  [[ -z "$name" ]] && { log::error "Missing required flag: --name"; return 1; }
  identity::require_exists "$name" || return 1
 
  # Gather identity-level metadata
  local policy strict auto rules_list peer_count
  policy=$(identity::policy "$name")
  strict=$(identity::rule_flags "$name" "strict_rule")
  auto=$(identity::rule_flags   "$name" "auto_apply")
  rules_list=$(identity::rules  "$name" | tr '\n' ',' | sed 's/,$//' | sed 's/,/, /g')
 
  local data
  data=$(identity::show_data "$name")
  peer_count=$(echo "$data" | grep '^peer_count|' | cut -d'|' -f2)
 
  # Precompute handshakes once for all peers
  declare -A _id_handshakes=()
  while IFS=$'\t' read -r pk ts; do
    [[ -n "$pk" ]] && _id_handshakes["$pk"]="$ts"
  done < <(wg show "$(config::interface)" latest-handshakes 2>/dev/null)
 
  # Header
  echo ""
  ui::row "Identity"    "$name"
  ui::row "Policy"      "$policy"
  ui::row "Rules"       "${rules_list:-—}"
  ui::row "Strict rule" "$(ui::bool "$strict")"
  ui::row "Auto apply"  "$(ui::bool "$auto")"
  ui::row "Peers"       "$peer_count"
  echo ""
 
  # Device list
  while IFS='|' read -r key val type_val index_val; do
    case "$key" in
      name|peer_count) ;;
      device)
        local status=""
        status=$(cmd::identity::_device_status "$val" _id_handshakes)
        ui::identity::device_row "$val" "$type_val" "$index_val" "$status"
        ;;
    esac
  done <<< "$data"
 
  # Rules tree
  local identity_rules
  identity_rules=$(identity::rules "$name")
  if [[ -n "$identity_rules" ]]; then
    printf "\n  \033[2m── Rules \033[0m%s\n\n" \
    "$(printf '\033[2m─%.0s' {1..38})"
    ui::rule::identity_block "$name" "$strict" --no-header
  fi
 
  echo ""
}

function cmd::identity::_device_status() {
  local peer_name="${1:-}"
  local -n _handshakes="${2:-__empty_map}"
 
  local peer_ip
  peer_ip=$(peers::get_ip "$peer_name" 2>/dev/null) || return 0
  [[ -z "$peer_ip" ]] && return 0
 
  local is_blocked is_restricted pubkey handshake_ts
  peers::is_blocked "$peer_name" && is_blocked="true" || is_blocked="false"
  peers::is_restricted "$peer_name" && is_restricted="true" || is_restricted="false"
 
  pubkey="$(keys::public "$peer_name")"
  handshake_ts="${_handshakes[$pubkey]:-0}"
 
  local last_ts
  last_ts=$(peers::get_meta "$peer_name" "last_ts" 2>/dev/null) || last_ts=""
 
  local status
  status=$(peers::format_status_verbose \
    "$peer_name" "$pubkey" "$is_blocked" "$is_restricted" "$handshake_ts" "$last_ts")
  echo " — ${status}"
}

function cmd::identity::_add() {
  local name="" peer=""
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --name) name="$2"; shift 2 ;;
      --peer) peer="$2"; shift 2 ;;
      --help) cmd::identity::help; return ;;
      *) log::error "Unknown flag: $1"; return 1 ;;
    esac
  done

  [[ -z "$name" ]] && { log::error "Missing required flag: --name"; return 1; }
  [[ -z "$peer" ]] && { log::error "Missing required flag: --peer"; return 1; }

  cmd::identity::_require_peer_exists "$peer" || return 1

  local peer_type index
  peer_type=$(cmd::identity::_resolve_peer_type "$peer" "$name")
  index=$(identity::next_index "$name" "$peer_type")

  local id_file
  id_file=$(ctx::identity::path "${name}.identity")
  json::identity_add_peer "$id_file" "$name" "$peer" "$peer_type" "$index" </dev/null
  log::ok "Added '${peer}' to identity '${name}' (${peer_type} #${index})"
}

function cmd::identity::_require_peer_exists() {
  local peer="${1:-}"
  if [[ ! -f "$(ctx::clients)/${peer}.conf" ]]; then
    log::error "Peer '${peer}' not found"
    return 1
  fi
}

function cmd::identity::_resolve_peer_type() {
  local peer="${1:-}" identity_name="${2:-}"
  local inferred
  inferred=$(identity::infer "$peer")
  if [[ -n "$inferred" ]]; then
    echo "$inferred" | cut -d'|' -f2
  else
    peers::get_meta "$peer" "type" 2>/dev/null || echo "none"
  fi
}

function cmd::identity::_remove() {
  local name="" force=false
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --name)  name="$2";    shift 2 ;;
      --force) force=true;   shift   ;;
      --help)  cmd::identity::help; return ;;
      *) log::error "Unknown flag: $1"; return 1 ;;
    esac
  done

  [[ -z "$name" ]] && { log::error "Missing required flag: --name"; return 1; }
  identity::require_exists "$name" || return 1

  local peers
  peers=$(identity::peers "$name")

  if [[ -n "$peers" ]]; then
    local peer_list="${peers//$'\n'/, }"
    log::warn "This will permanently remove identity '${name}' and ALL associated peers:"
    log::warn "  ${peer_list}"

    if ! $force; then
      ui::confirm "Continue?" || { log::info "Aborted"; return 0; }
    fi

    cmd::identity::_remove_all_peers "$peers"
    peers::reload || return 1
  fi

  local id_file
  id_file=$(ctx::identity::path "${name}.identity")
  json::identity_remove "$id_file" </dev/null
  log::ok "Identity '${name}' removed"
}

function cmd::identity::_remove_all_peers() {
  local peers="${1:-}"
  while IFS= read -r peer_name; do
    [[ -z "$peer_name" ]] && continue
    cmd::identity::_remove_peer "$peer_name"
  done <<< "$peers"
}

function cmd::identity::_remove_peer() {
  local peer_name="${1:-}"
  local client_ip was_blocked=false

  client_ip=$(peers::get_ip "$peer_name" 2>/dev/null) || client_ip=""
  peers::is_blocked "$peer_name" && was_blocked=true

  peers::purge "$peer_name" "$client_ip" "$was_blocked" || return 1
  log::ok "Removed peer '${peer_name}'"
}

function cmd::identity::_migrate() {
  local dry_run="false"
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --dry-run) dry_run="true"; shift ;;
      --help)    cmd::identity::help; return ;;
      *) log::error "Unknown flag: $1"; return 1 ;;
    esac
  done

  [[ "$dry_run" == "true" ]] && log::info "Dry run — no files will be written"
  echo ""

  [[ "$dry_run" == "false" ]] && mkdir -p "$(ctx::identities)"

  local created=0 skipped=0 output
  output=$(json::identity_migrate \
    "$(ctx::identities)" \
    "$(ctx::clients)" \
    "$(ctx::meta)" \
    "$dry_run")

  while IFS='|' read -r action identity_name peer_name peer_type index; do
    case "$action" in
      create)
        ui::identity::migrate_create "$peer_name" "$identity_name" "$peer_type" "$index"
        (( created++ )) || true
        ;;
      skip)
        ui::identity::migrate_skip "$peer_name"
        (( skipped++ )) || true
        ;;
    esac
  done <<< "$output"

  ui::identity::migrate_summary "$created" "$skipped" "$dry_run"
}

function cmd::identity::_rule() {
  local subcmd="${1:-show}"
  shift || true
 
  case "$subcmd" in
    assign)   cmd::identity::_rule_assign   "$@" ;;
    unassign) cmd::identity::_rule_unassign "$@" ;;
    show)     cmd::identity::_rule_show     "$@" ;;
    *)
      log::error "Unknown rule subcommand '${subcmd}'. Available: assign, unassign, show"
      return 1
      ;;
  esac
}
 
function cmd::identity::_rule_assign() {
  local name="" rule=""
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --name) name="$2"; shift 2 ;;
      --rule) rule="$2"; shift 2 ;;
      *) log::error "Unknown flag: $1"; return 1 ;;
    esac
  done
 
  [[ -z "$name" ]] && { log::error "Missing required flag: --name"; return 1; }
  [[ -z "$rule" ]] && { log::error "Missing required flag: --rule"; return 1; }
  identity::require_exists "$name" || return 1
  rule::exists "$rule" || { log::error "Rule not found: ${rule}"; return 1; }
 
  local exit_code
  identity::add_rule "$name" "$rule" || exit_code=$?  
 
  if [[ $exit_code -eq 2 ]]; then
    log::warn "Rule '${rule}' is already assigned to identity '${name}'"
    return 0
  fi
 
  log::ok "Rule '${rule}' assigned to identity '${name}'"
 
  # Reapply rules if auto_apply
  local auto
  auto=$(identity::rule_flags "$name" "auto_apply")
  if [[ "$auto" != "false" ]]; then
    log::info "Reapplying rules for all peers in identity '${name}'..."
    identity::reapply_rules "$name"
    log::ok "Rules reapplied"
  fi
 
  # Warn about strict_rule
  if policy::strict_rule "$(identity::policy "$name")"; then
    log::warn "Strict rule is enabled for identity '${name}' — peer rules will not be additive"
  fi
}
 
function cmd::identity::_rule_unassign() {
  local name="" rule="" all=false
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --name) name="$2"; shift 2 ;;
      --rule) rule="$2"; shift 2 ;;
      --all)  all=true;  shift   ;;
      *) log::error "Unknown flag: $1"; return 1 ;;
    esac
  done
 
  [[ -z "$name" ]] && { log::error "Missing required flag: --name"; return 1; }
  identity::require_exists "$name" || return 1
 
  if $all; then
    local rules
    rules=$(identity::rules "$name")
    if [[ -z "$rules" ]]; then
      log::warn "Identity '${name}' has no rules assigned"
      return 0
    fi
    identity::clear_rules "$name"
    log::ok "All rules removed from identity '${name}'"
    cmd::identity::_reapply_after_unassign "$name"
    return 0
  fi
 
  [[ -z "$rule" ]] && {
    log::error "Missing required flag: --rule (or use --all to remove all)"
    return 1
  }
 
  identity::remove_rule "$name" "$rule"
  local exit_code=$?
  if [[ $exit_code -ne 0 ]]; then
    log::error "Rule '${rule}' is not assigned to identity '${name}'"
    return 1
  fi
 
  log::ok "Rule '${rule}' removed from identity '${name}'"
  cmd::identity::_reapply_after_unassign "$name"
}

function cmd::identity::_reapply_after_unassign() {
  local name="${1:-}"
  local auto
  auto=$(identity::rule_flags "$name" "auto_apply")
  if [[ "$auto" != "false" ]]; then
    log::info "Reapplying rules for all peers in identity '${name}'..."
    identity::reapply_rules "$name"
    log::ok "Rules reapplied"
  else
    log::info "Note: auto_apply is disabled — run 'wgctl audit --fix' to update fw rules"
  fi
}
 
function cmd::identity::_rule_show() {
  local name=""
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --name) name="$2"; shift 2 ;;
      *) log::error "Unknown flag: $1"; return 1 ;;
    esac
  done
 
  [[ -z "$name" ]] && { log::error "Missing required flag: --name"; return 1; }
  identity::require_exists "$name" || return 1
 
  local rules policy strict auto
  rules=$(identity::rules "$name")
  policy=$(identity::policy "$name")
  strict=$(identity::rule_flags "$name" "strict_rule")
  auto=$(identity::rule_flags "$name" "auto_apply")
 
  echo ""
  ui::row "Identity"    "$name"
  ui::row "Policy"      "$policy"
  ui::row "Strict rule" "$(ui::bool "$strict")"
  ui::row "Auto apply"  "$(ui::bool "$auto")"
  echo ""
 
  if [[ -z "$rules" ]]; then
    ui::row "Rules" "— none assigned"
  else
    printf "  %-20s\n" "Rules:"
    while IFS= read -r rule_name; do
      [[ -z "$rule_name" ]] && continue
      printf "    · %s\n" "$rule_name"
    done <<< "$rules"
  fi
  echo ""
}

function cmd::identity::_options() {
  local name="" new_policy=""
  local set_strict="" set_auto=""
 
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --name)             name="$2";        shift 2 ;;
      --policy)           new_policy="$2";  shift 2 ;;
      --set-strict-rule)  set_strict="true";  shift ;;
      --unset-strict-rule) set_strict="false"; shift ;;
      --set-auto-apply)   set_auto="true";    shift ;;
      --unset-auto-apply) set_auto="false";   shift ;;
      *) log::error "Unknown flag: $1"; return 1 ;;
    esac
  done
 
  [[ -z "$name" ]] && { log::error "Missing required flag: --name"; return 1; }
  identity::require_exists "$name" || return 1
 
  local changed=false
 
  if [[ -n "$new_policy" ]]; then
    policy::require_exists "$new_policy" || return 1
    identity::set_policy "$name" "$new_policy"
    log::ok "Policy set to '${new_policy}' for identity '${name}'"
    changed=true
  fi
 
  if [[ -n "$set_strict" ]]; then
    identity::set_rule_flag "$name" "strict_rule" "$set_strict"
    if [[ "$set_strict" == "true" ]]; then
      log::ok "Strict rule enabled for identity '${name}' — peer rules will not be additive"
    else
      log::ok "Strict rule disabled for identity '${name}' — peer rules will be additive"
    fi
    changed=true
  fi
 
  if [[ -n "$set_auto" ]]; then
    identity::set_rule_flag "$name" "auto_apply" "$set_auto"
    if [[ "$set_auto" == "true" ]]; then
      log::ok "Auto apply enabled for identity '${name}'"
    else
      log::ok "Auto apply disabled for identity '${name}'"
    fi
    changed=true
  fi
 
  if ! $changed; then
    cmd::identity::_rule_show --name "$name"
  fi
}