﻿#!/bin/bash
export PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"

LOG_FILE="/var/log/wems-monitor.log"
COMPOSE_FILE="/opt/lnxall/docker-compose.yml"

LOCK_FILE="/tmp/wems-monitor.lock"
STATE_DIR="/tmp/wems-monitor-state"
DOCKER_RESTART_FLAG="/tmp/wems-docker-restarting.flag"

HTTP_TIMEOUT=10
TCP_TIMEOUT=5
FAIL_THRESHOLD=3
STARTUP_GRACE_SECONDS=600
CONTAINER_RESTART_THRESHOLD=2
DOCKER_RESTART_COOLDOWN=600
DOCKER_CMD_TIMEOUT=120

mkdir -p "$(dirname "$LOG_FILE")"
mkdir -p "$STATE_DIR"

exec 200>"$LOCK_FILE"
flock -n 200 || exit 0

BASE_SERVICES=("postgres" "redis" "emqx")

SERVICES=(
  "postgres|"
  "redis|tcp://127.0.0.1:6379"
  "emqx|tcp://127.0.0.1:3883"
  "wems|http://127.0.0.1:8080/blade-wems/config/system/detail"
  "wems-job|http://127.0.0.1:8082/blade-wems/config/system/detail"
  "wems-open-api|http://127.0.0.1:8083/blade-wems/config/system/detail"
  "wems-web|http://127.0.0.1/"
)

log() {
  echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" >> "$LOG_FILE"
}

is_base_service() {
  local name="$1"
  for b in "${BASE_SERVICES[@]}"; do
    [ "$name" = "$b" ] && return 0
  done
  return 1
}

container_exists() {
  timeout 10 docker ps -a --format '{{.Names}}' | grep -q "^$1$"
}

container_running() {
  timeout 10 docker ps --format '{{.Names}}' | grep -q "^$1$"
}

check_http() {
  local url="$1"
  local code
  local err_file
  local err_msg

  err_file=$(mktemp)

  code=$(curl -sS -o /dev/null -w "%{http_code}" \
        --max-time "$HTTP_TIMEOUT" "$url" 2>"$err_file")

  err_msg=$(<"$err_file")
  rm -f "$err_file"

  # 去掉换行/空白
  err_msg=$(echo "$err_msg" | tr -d '\r\n')

  if [[ "$code" =~ ^[0-9]{3}$ ]] && [ "$code" -ge 200 ] && [ "$code" -lt 500 ]; then
    return 0
  fi

  if [ "$code" = "000" ]; then
    log "HTTP连接失败：$url，错误：$err_msg"
  elif [ -n "$err_msg" ]; then
    log "HTTP检查失败：$url，状态码：$code，错误：$err_msg"
  else
    log "HTTP检查失败：$url，状态码：$code"
  fi

  return 1
}

check_tcp() {
  local target="$1"
  local host
  local port

  target="${target#tcp://}"
  host="${target%:*}"
  port="${target##*:}"

  if timeout "$TCP_TIMEOUT" bash -c "</dev/tcp/${host}/${port}" 2>/dev/null; then
    return 0
  fi

  log "TCP检查失败：${host}:${port}"
  return 1
}

check_service() {
  local name="$1"
  local url="$2"

  if ! container_exists "$name"; then
    log "容器不存在：$name"
    return 1
  fi

  if ! container_running "$name"; then
    log "容器未运行：$name"
    return 1
  fi

  if [ -n "$url" ]; then
    if [[ "$url" == tcp://* ]]; then
      check_tcp "$url" || return 1
    else
      check_http "$url" || return 1
    fi
  fi

  return 0
}

get_count() {
  [ -f "$1" ] && cat "$1" || echo 0
}

set_count() {
  echo "$2" > "$1"
}

mark_state() {
  local name="$1"
  local new_state="$2"
  local state_file="$STATE_DIR/${name}.state"
  local old_state="unknown"

  [ -f "$state_file" ] && old_state=$(cat "$state_file")

  if [ "$old_state" != "$new_state" ]; then
    echo "$new_state" > "$state_file"

    if [ "$new_state" = "ok" ]; then
      log "$name 恢复正常"
    elif [ "$new_state" = "fail" ]; then
      log "$name 出现异常"
    fi
  fi
}

is_in_grace() {
  local name="$1"
  local file="$STATE_DIR/${name}.grace"

  [ ! -f "$file" ] && return 1

  local now
  local start
  local age

  now=$(date +%s)
  start=$(cat "$file")
  age=$((now - start))

  if [ "$age" -lt "$STARTUP_GRACE_SECONDS" ]; then
    return 0
  fi

  rm -f "$file"
  return 1
}

enter_grace() {
  date +%s > "$STATE_DIR/${1}.grace"
}

restart_one_container() {
  local name="$1"

  log "开始重启单个容器：$name"

  timeout "$DOCKER_CMD_TIMEOUT" docker restart "$name" >> "$LOG_FILE" 2>&1
  local rc=$?

  if [ "$rc" -eq 0 ]; then
    log "单个容器重启命令执行完成：$name"
  elif [ "$rc" -eq 124 ]; then
    log "警告：单个容器重启命令超时：$name，进入启动保护期"
  else
    log "错误：单个容器重启命令失败：$name，返回码：$rc"
  fi

  enter_grace "$name"
}

compose_up() {
  if [ ! -f "$COMPOSE_FILE" ]; then
    log "错误：compose 文件不存在：$COMPOSE_FILE"
    return 1
  fi

  if timeout 10 docker compose version >/dev/null 2>&1; then
    log "使用 docker compose 启动服务：$COMPOSE_FILE"
    timeout "$DOCKER_CMD_TIMEOUT" docker compose -f "$COMPOSE_FILE" up -d --no-build >> "$LOG_FILE" 2>&1
    return $?
  fi

  if command -v docker-compose >/dev/null 2>&1; then
    log "使用 docker-compose 启动服务：$COMPOSE_FILE"
    timeout "$DOCKER_CMD_TIMEOUT" docker-compose -f "$COMPOSE_FILE" up -d --no-build >> "$LOG_FILE" 2>&1
    return $?
  fi

  if [ -x "/usr/local/bin/docker-compose" ]; then
    log "使用 /usr/local/bin/docker-compose 启动服务：$COMPOSE_FILE"
    timeout "$DOCKER_CMD_TIMEOUT" /usr/local/bin/docker-compose -f "$COMPOSE_FILE" up -d --no-build >> "$LOG_FILE" 2>&1
    return $?
  fi

  if [ -x "/usr/bin/docker-compose" ]; then
    log "使用 /usr/bin/docker-compose 启动服务：$COMPOSE_FILE"
    timeout "$DOCKER_CMD_TIMEOUT" /usr/bin/docker-compose -f "$COMPOSE_FILE" up -d --no-build >> "$LOG_FILE" 2>&1
    return $?
  fi

  log "错误：未找到 docker compose / docker-compose 命令"
  return 1
}

log_top_memory_usage() {
  log "当前宿主机内存占用最高进程："
  ps aux --sort=-%mem | head -10 >> "$LOG_FILE" 2>&1

  log "当前容器资源占用情况："
  timeout 10 docker stats --no-stream --format "table {{.Name}}\t{{.MemUsage}}\t{{.MemPerc}}\t{{.CPUPerc}}" >> "$LOG_FILE" 2>&1
}
restart_docker_and_all_services() {
  if [ -f "$DOCKER_RESTART_FLAG" ]; then
    local now
    local flag_time
    local age

    now=$(date +%s)
    flag_time=$(stat -c %Y "$DOCKER_RESTART_FLAG")
    age=$((now - flag_time))

    if [ "$age" -lt "$DOCKER_RESTART_COOLDOWN" ]; then
      log "Docker 重启冷却期内，已等待 ${age}s，跳过本次 Docker 重启"
      return 0
    fi

    log "Docker 重启标志已过期，清理后继续"
    rm -f "$DOCKER_RESTART_FLAG"
  fi

  touch "$DOCKER_RESTART_FLAG"

  log "开始重启 Docker 服务"
  log_top_memory_usage

  timeout "$DOCKER_CMD_TIMEOUT" systemctl restart docker >> "$LOG_FILE" 2>&1
  local rc=$?

  if [ "$rc" -eq 124 ]; then
    log "错误：Docker 重启命令超时"
    rm -f "$DOCKER_RESTART_FLAG"
    return 1
  fi

  sleep 30

  if ! systemctl is-active --quiet docker || ! timeout 10 docker ps >/dev/null 2>&1; then
    log "错误：Docker 状态异常"
    rm -f "$DOCKER_RESTART_FLAG"
    return 1
  fi

  log "Docker 服务已恢复，开始使用 compose 启动整套服务"

  if ! compose_up; then
    log "错误：compose 启动失败"
    rm -f "$DOCKER_RESTART_FLAG"
    return 1
  fi

  log "compose 启动命令执行完成"

  for item in "${SERVICES[@]}"; do
    local svc="${item%%|*}"
    enter_grace "$svc"
  done

  return 0
}

handle_failure() {
  local name="$1"

  if is_in_grace "$name"; then
    return 0
  fi

  local fail_file="$STATE_DIR/${name}.fail"
  local restart_file="$STATE_DIR/${name}.restart"
  local fail_count
  local restart_count

  fail_count=$(get_count "$fail_file")
  fail_count=$((fail_count + 1))
  set_count "$fail_file" "$fail_count"

  if [ "$fail_count" -lt "$FAIL_THRESHOLD" ]; then
    log "$name 连续失败次数：$fail_count，未达到重启阈值"
    return 0
  fi

  set_count "$fail_file" 0

  restart_count=$(get_count "$restart_file")
  restart_count=$((restart_count + 1))
  set_count "$restart_file" "$restart_count"

  log "$name 达到失败阈值，开始重启单个容器；累计重启次数：$restart_count"

  restart_one_container "$name"

  if [ "$restart_count" -ge "$CONTAINER_RESTART_THRESHOLD" ]; then
    log "$name 单容器多次重启后仍异常，触发 Docker 级恢复"
    set_count "$restart_file" 0

    restart_docker_and_all_services
  fi
}

clear_state_if_ok() {
  local name="$1"

  rm -f "$STATE_DIR/${name}.fail"
  rm -f "$STATE_DIR/${name}.restart"
  rm -f "$STATE_DIR/${name}.grace"
}

main() {
  local base_ok=1
  local skipped_business=0

  for item in "${SERVICES[@]}"; do
    name="${item%%|*}"
    url="${item#*|}"

    if [ "$base_ok" -eq 0 ] && ! is_base_service "$name"; then
      if [ "$skipped_business" -eq 0 ]; then
        log "基础服务异常，本轮跳过业务服务检查"
        skipped_business=1
      fi
      continue
    fi

    if check_service "$name" "$url"; then
      mark_state "$name" "ok"
      clear_state_if_ok "$name"
    else
      mark_state "$name" "fail"
      handle_failure "$name"

      if is_base_service "$name"; then
        base_ok=0
      fi
    fi
  done
}

main