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

# 原始日志可能包含敏感字段，确保新建目录和文件默认只允许当前用户读取。
umask 077

usage() {
  cat <<'EOF'
用法:
  capture-device-raw-logs.sh DEVICE_ID [DURATION_SECONDS] [MESSAGE_COUNT] [BASE_URL] [OUTPUT_DIR]

参数:
  DEVICE_ID          必填，设备ID
  DURATION_SECONDS   可选，采集秒数，1~600，默认60
  MESSAGE_COUNT      可选，样本条数，1~50，默认10
  BASE_URL           可选，默认 http://127.0.0.1:8289/xdrrule
  OUTPUT_DIR         可选，默认当前目录下 raw-log-captures

可选环境变量:
  RAW_CAPTURE_AUTH_HEADER       例如 Authorization: Bearer xxx
  RAW_CAPTURE_POLL_SECONDS      查询间隔，默认2秒
  RAW_CAPTURE_WAIT_TIMEOUT_SECONDS  总等待超时；默认采集时间加420秒
  RAW_CAPTURE_HTTP_MAX_TIME_SECONDS 单次HTTP超时，默认120秒
EOF
}

fail() {
  printf 'ERROR: %s\n' "$*" >&2
  exit 1
}

require_command() {
  command -v "$1" >/dev/null 2>&1 || fail "缺少命令: $1"
}

positive_integer() {
  [[ "$1" =~ ^[0-9]+$ ]] && (( 10#$1 > 0 ))
}

if [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then
  usage
  exit 0
fi

DEVICE_ID="${1:-}"
DURATION_SECONDS="${2:-60}"
MESSAGE_COUNT="${3:-10}"
BASE_URL="${4:-http://127.0.0.1:8289/xdrrule}"
OUTPUT_DIR="${5:-$(pwd)/raw-log-captures}"
POLL_SECONDS="${RAW_CAPTURE_POLL_SECONDS:-2}"
HTTP_MAX_TIME_SECONDS="${RAW_CAPTURE_HTTP_MAX_TIME_SECONDS:-120}"

[[ -n "$DEVICE_ID" ]] || { usage >&2; fail "必须提供 DEVICE_ID"; }
positive_integer "$DEVICE_ID" || fail "DEVICE_ID 必须是大于0的整数"
positive_integer "$DURATION_SECONDS" || fail "DURATION_SECONDS 必须是大于0的整数"
positive_integer "$MESSAGE_COUNT" || fail "MESSAGE_COUNT 必须是大于0的整数"
positive_integer "$POLL_SECONDS" || fail "RAW_CAPTURE_POLL_SECONDS 必须是大于0的整数"
positive_integer "$HTTP_MAX_TIME_SECONDS" || fail "RAW_CAPTURE_HTTP_MAX_TIME_SECONDS 必须是大于0的整数"
(( 10#$DURATION_SECONDS <= 600 )) || fail "DURATION_SECONDS 不能超过600"
(( 10#$MESSAGE_COUNT <= 50 )) || fail "MESSAGE_COUNT 不能超过50"
(( 10#$POLL_SECONDS <= 30 )) || fail "RAW_CAPTURE_POLL_SECONDS 不能超过30"

DEFAULT_WAIT_TIMEOUT=$((10#$DURATION_SECONDS + 420))
WAIT_TIMEOUT_SECONDS="${RAW_CAPTURE_WAIT_TIMEOUT_SECONDS:-$DEFAULT_WAIT_TIMEOUT}"
positive_integer "$WAIT_TIMEOUT_SECONDS" || fail "RAW_CAPTURE_WAIT_TIMEOUT_SECONDS 必须是大于0的整数"

require_command curl
require_command python3
require_command mktemp

BASE_URL="${BASE_URL%/}"
if [[ "$BASE_URL" == */api/syslog/raw-captures ]]; then
  CAPTURE_API="$BASE_URL"
else
  CAPTURE_API="$BASE_URL/api/syslog/raw-captures"
fi

mkdir -p -- "$OUTPUT_DIR"
TEMP_DIR="$(mktemp -d "${TMPDIR:-/tmp}/hb-raw-capture.XXXXXX")"
SESSION_ID=""
CAPTURE_DIR=""

cleanup() {
  rm -rf -- "$TEMP_DIR"
}
trap cleanup EXIT

CURL_HEADERS=(-H 'Accept: application/json')
if [[ -n "${RAW_CAPTURE_AUTH_HEADER:-}" ]]; then
  CURL_HEADERS+=(-H "$RAW_CAPTURE_AUTH_HEADER")
fi

http_request() {
  local method="$1" url="$2" output="$3" body="${4:-}"
  local args=(-sS -o "$output" -w '%{http_code}' --connect-timeout 10
    --max-time "$HTTP_MAX_TIME_SECONDS" -X "$method" "${CURL_HEADERS[@]}")
  if [[ -n "$body" ]]; then
    args+=(-H 'Content-Type: application/json' --data "$body")
  fi
  curl "${args[@]}" "$url"
}

START_RESPONSE="$TEMP_DIR/start.json"
REQUEST_BODY="{\"deviceId\":$((10#$DEVICE_ID)),\"durationSeconds\":$((10#$DURATION_SECONDS)),\"messageCount\":$((10#$MESSAGE_COUNT))}"

printf '创建采集会话: deviceId=%s, duration=%ss, count=%s\n' \
  "$DEVICE_ID" "$DURATION_SECONDS" "$MESSAGE_COUNT"
if ! HTTP_CODE="$(http_request POST "$CAPTURE_API" "$START_RESPONSE" "$REQUEST_BODY")"; then
  fail "无法连接采集API: $CAPTURE_API"
fi
if [[ "$HTTP_CODE" != "202" && "$HTTP_CODE" != "200" ]]; then
  printf 'API响应(%s):\n' "$HTTP_CODE" >&2
  python3 -m json.tool "$START_RESPONSE" 2>/dev/null >&2 || sed -n '1,20p' "$START_RESPONSE" >&2
  fail "创建采集会话失败"
fi

SESSION_ID="$(python3 - "$START_RESPONSE" <<'PY'
import json, sys
with open(sys.argv[1], encoding="utf-8") as source:
    value = json.load(source)
print(value.get("sessionId", ""))
PY
)"
[[ -n "$SESSION_ID" ]] || fail "API响应缺少 sessionId"

START_STAMP="$(date '+%Y%m%d-%H%M%S')"
CAPTURE_DIR="$OUTPUT_DIR/device-${DEVICE_ID}_${START_STAMP}_${SESSION_ID}"
mkdir -p -- "$CAPTURE_DIR/raw"
cp -- "$START_RESPONSE" "$CAPTURE_DIR/start.json"
printf '%s\n' "$SESSION_ID" > "$CAPTURE_DIR/session-id.txt"
printf '会话已创建: %s\n结果目录: %s\n' "$SESSION_ID" "$CAPTURE_DIR"

LATEST_TEMP="$TEMP_DIR/latest.json"
DEADLINE=$(( $(date +%s) + 10#$WAIT_TIMEOUT_SECONDS ))
FINAL_STATUS=""

while (( $(date +%s) <= DEADLINE )); do
  if ! HTTP_CODE="$(http_request GET "$CAPTURE_API/$SESSION_ID" "$LATEST_TEMP")"; then
    printf 'WARN: 本次查询失败，%s秒后重试\n' "$POLL_SECONDS" >&2
    sleep "$POLL_SECONDS"
    continue
  fi
  if [[ "$HTTP_CODE" != "200" ]]; then
    printf 'WARN: 查询返回HTTP %s，%s秒后重试\n' "$HTTP_CODE" "$POLL_SECONDS" >&2
    sleep "$POLL_SECONDS"
    continue
  fi

  cp -- "$LATEST_TEMP" "$CAPTURE_DIR/latest.json.tmp"
  mv -- "$CAPTURE_DIR/latest.json.tmp" "$CAPTURE_DIR/latest.json"
  FINAL_STATUS="$(python3 - "$LATEST_TEMP" <<'PY'
import json, sys
with open(sys.argv[1], encoding="utf-8") as source:
    print(json.load(source).get("status", "UNKNOWN"))
PY
)"
  COUNTS="$(python3 - "$LATEST_TEMP" <<'PY'
import json, sys
with open(sys.argv[1], encoding="utf-8") as source:
    value = json.load(source)
print("samples=%s, anomalies=%s, received=%s" % (
    len(value.get("messages") or []), value.get("retainedAnomalyCount", 0),
    value.get("receivedCount", 0)))
PY
)"
  printf '状态: %-10s %s\n' "$FINAL_STATUS" "$COUNTS"
  if [[ "$FINAL_STATUS" == "COMPLETED" || "$FINAL_STATUS" == "CANCELLED" ]]; then
    break
  fi
  sleep "$POLL_SECONDS"
done

if [[ ! -f "$CAPTURE_DIR/latest.json" ]]; then
  fail "等待期间没有取得有效结果；sessionId=$SESSION_ID"
fi

cp -- "$CAPTURE_DIR/latest.json" "$CAPTURE_DIR/result.json"

python3 - "$CAPTURE_DIR/result.json" "$CAPTURE_DIR" <<'PY'
import json
import os
import re
import sys

result_path, output_dir = sys.argv[1], sys.argv[2]
with open(result_path, encoding="utf-8") as source:
    result = json.load(source)

messages = result.get("messages") or []
jsonl_path = os.path.join(output_dir, "samples.jsonl")
with open(jsonl_path, "w", encoding="utf-8", newline="\n") as target:
    for message in messages:
        target.write(json.dumps(message, ensure_ascii=False, separators=(",", ":")) + "\n")

raw_dir = os.path.join(output_dir, "raw")
os.makedirs(raw_dir, exist_ok=True)
for index, message in enumerate(messages, 1):
    event_id = str(message.get("eventId") or "no-event-id")
    safe_event_id = re.sub(r"[^A-Za-z0-9._-]", "_", event_id)[:80]
    kind = "anomaly" if message.get("anomaly") else "normal"
    path = os.path.join(raw_dir, "%03d_%s_%s.log" % (index, kind, safe_event_id))
    raw = message.get("rawMessage")
    with open(path, "w", encoding="utf-8", newline="") as target:
        target.write("" if raw is None else str(raw))

summary = [
    "sessionId=%s" % result.get("sessionId", ""),
    "deviceId=%s" % result.get("deviceId", ""),
    "status=%s" % result.get("status", ""),
    "startedAt=%s" % result.get("startedAt", ""),
    "captureEndsAt=%s" % result.get("captureEndsAt", ""),
    "settleEndsAt=%s" % result.get("settleEndsAt", ""),
    "receivedCount=%s" % result.get("receivedCount", 0),
    "sampleCount=%s" % len(messages),
    "retainedAnomalyCount=%s" % result.get("retainedAnomalyCount", 0),
    "replacementCount=%s" % result.get("replacementCount", 0),
]
with open(os.path.join(output_dir, "summary.txt"), "w", encoding="utf-8", newline="\n") as target:
    target.write("\n".join(summary) + "\n")
PY

if [[ "$FINAL_STATUS" != "COMPLETED" && "$FINAL_STATUS" != "CANCELLED" ]]; then
  printf 'WARN: 等待超时，已保存最后一次快照；会话仍会在服务端自动结束。\n' >&2
  printf '结果目录: %s\n' "$CAPTURE_DIR"
  exit 2
fi

printf '采集完成。\n完整结果: %s/result.json\n原文目录: %s/raw\n' \
  "$CAPTURE_DIR" "$CAPTURE_DIR"