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

MODE="${1:-status}"
TAB_LIMIT="${AGENT_BROWSER_TAB_LIMIT:-10}"
RSS_LIMIT_MB="${AGENT_BROWSER_RSS_LIMIT_MB:-1500}"
INCLUDE_HEADED="${AGENT_BROWSER_CLEANUP_INCLUDE_HEADED:-0}"

profileForPort() {
  case "$1" in
    9223) printf '%s\n' "$HOME/.agent-browser-profile" ;;
    9226) printf '%s\n' "$HOME/.agent-browser-profile-headless" ;;
    *) printf '' ;;
  esac
}

usage() {
  cat <<'USAGE'
Usage: scripts/agent-browser-cleanup.sh [status|soft|hard|auto]

Modes:
  status  只查看 9223/9226 的标签页数量与内存占用
  soft    关闭 about:blank、newtab、devtools 等无业务标签页
  hard    优雅退出 agent-browser 专用 Chrome（9223/9226），释放内存
  auto    超阈值时自动释放 9226；9223 默认只提示，避免误关用户正在看的页面

Env:
  AGENT_BROWSER_TAB_LIMIT      默认 10
  AGENT_BROWSER_RSS_LIMIT_MB   默认 1500
  AGENT_BROWSER_CLEANUP_INCLUDE_HEADED=1 时，auto 也允许释放 9223
USAGE
}

if [[ "$MODE" == "-h" || "$MODE" == "--help" ]]; then
  usage
  exit 0
fi

if [[ "$MODE" != "status" && "$MODE" != "soft" && "$MODE" != "hard" && "$MODE" != "auto" ]]; then
  usage >&2
  exit 2
fi

portStatus() {
  local portNumber="$1"
  local profileDir=""
  local tabCount="0"
  local rssMb="0"
  local profileArg=""

  profileDir=$(profileForPort "$portNumber")
  profileArg="--user-data-dir=${profileDir}"
  rssMb=$(ps -axo rss=,command= | awk -v profileArg="$profileArg" '
    {
      for (fieldIndex = 2; fieldIndex <= NF; fieldIndex++) {
        if ($fieldIndex == profileArg) {
          sum += $1
          break
        }
      }
    }
    END {printf "%d", sum / 1024}
  ')

  if curl -sf "http://127.0.0.1:${portNumber}/json/list" >/tmp/agent-browser-tabs-${portNumber}.json 2>/dev/null; then
    tabCount=$(jq 'length' "/tmp/agent-browser-tabs-${portNumber}.json" 2>/dev/null || echo "0")
  fi

  rm -f "/tmp/agent-browser-tabs-${portNumber}.json"
  printf '%s %s %s\n' "$portNumber" "$tabCount" "$rssMb"
}

printStatus() {
  printf 'port\ttabs\trssMB\n'
  portStatus 9223 | awk '{printf "%s\t%s\t%s\n", $1, $2, $3}'
  portStatus 9226 | awk '{printf "%s\t%s\t%s\n", $1, $2, $3}'
}

closeChromeByPort() {
  local portNumber="$1"
  local profileDir=""
  local profileArg=""
  local processIds=""
  profileDir=$(profileForPort "$portNumber")
  profileArg="--user-data-dir=${profileDir}"
  processIds=$(ps -axo pid=,command= | awk -v profileArg="$profileArg" '
    {
      hasProfile = 0
      for (fieldIndex = 2; fieldIndex <= NF; fieldIndex++) {
        if ($fieldIndex == profileArg) hasProfile = 1
      }
      if (hasProfile) print $1
    }
  ')
  if [[ -z "$processIds" ]]; then
    echo "port ${portNumber}: no agent-browser Chrome process"
    return
  fi

  echo "port ${portNumber}: graceful quit"
  printf '%s\n' "$processIds" | xargs -n1 /bin/kill -TERM 2>/dev/null || true
}

softCleanup() {
  local portNumber="$1"
  if ! curl -sf "http://127.0.0.1:${portNumber}/json/list" >/tmp/agent-browser-tabs-${portNumber}.json 2>/dev/null; then
    echo "port ${portNumber}: not reachable"
    return
  fi

  jq -r '.[] | select(.url | test("^(about:blank|chrome://newtab|devtools://|chrome://)")).id' "/tmp/agent-browser-tabs-${portNumber}.json" |
    while IFS= read -r tabId; do
      [[ -z "$tabId" ]] && continue
      curl -sf "http://127.0.0.1:${portNumber}/json/close/${tabId}" >/dev/null || true
      echo "port ${portNumber}: closed idle tab ${tabId}"
    done

  rm -f "/tmp/agent-browser-tabs-${portNumber}.json"
}

autoCleanup() {
  while read -r portNumber tabCount rssMb; do
    if (( tabCount <= TAB_LIMIT && rssMb <= RSS_LIMIT_MB )); then
      continue
    fi

    echo "port ${portNumber}: over limit tabs=${tabCount}/${TAB_LIMIT} rssMB=${rssMb}/${RSS_LIMIT_MB}"
    if [[ "$portNumber" == "9226" ]]; then
      closeChromeByPort "$portNumber"
      continue
    fi

    if [[ "$portNumber" == "9223" && "$INCLUDE_HEADED" == "1" ]]; then
      closeChromeByPort "$portNumber"
      continue
    fi

    echo "port ${portNumber}: headed Chrome skipped; set AGENT_BROWSER_CLEANUP_INCLUDE_HEADED=1 to close it automatically"
  done < <(portStatus 9223 && portStatus 9226)
}

case "$MODE" in
  status)
    printStatus
    ;;
  soft)
    softCleanup 9223
    softCleanup 9226
    printStatus
    ;;
  hard)
    closeChromeByPort 9223
    closeChromeByPort 9226
    ;;
  auto)
    printStatus
    autoCleanup
    ;;
esac