#!/usr/bin/env bash

# 统一管理脚本：启动/停止/查看 四个 API 服务
# 支持服务：
#  - intent      -> rag2_0.api.intent_recognition_api:app   (port 8001, workers 25)
#  - dify        -> rag2_0.api.DifyQueryRetrieval_api:app   (port 8002, workers 25)
#  - answertype  -> rag2_0.api.AnswerType_api:app           (port 8003, workers 1)
#  - qingdan     -> rag2_0.api.query_dinge_qingdan_api:app  (port 8005, workers 1)

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# 定义服务配置：会话名 与 启动命令
SERVICE_NAMES=(intent dify answertype qingdan)

service_port() {
  case "$1" in
    intent) echo "8001" ;;
    dify) echo "8002" ;;
    answertype) echo "8003" ;;
    qingdan) echo "8005" ;;
    *) echo "" ;;
  esac
}

session_name() {
  case "$1" in
    intent) echo "intent_recognition_api" ;;
    dify) echo "DifyQueryRetrieval_api" ;;
    answertype) echo "AnswerType" ;;
    qingdan) echo "query_dinge_qingdan_api" ;;
    *) echo "" ;;
  esac
}

start_command() {
  case "$1" in
    intent)
      echo "cd \"$SCRIPT_DIR\" && uv run uvicorn rag2_0.api.intent_recognition_api:app --host 0.0.0.0 --port 8001 --workers 4" ;;
    dify)
      echo "cd \"$SCRIPT_DIR\" && uv run uvicorn rag2_0.api.DifyQueryRetrieval_api:app --host 0.0.0.0 --port 8002 --workers 4" ;;
    answertype)
      echo "cd \"$SCRIPT_DIR\" && uv run uvicorn rag2_0.api.AnswerType_api:app --host 0.0.0.0 --port 8003 --workers 1" ;;
    qingdan)
      echo "cd \"$SCRIPT_DIR\" && uv run uvicorn rag2_0.api.query_dinge_qingdan_api:app --host 0.0.0.0 --port 8005 --workers 4" ;;
    *) echo "" ;;
  esac
}

exists_session() {
  # 使用严格匹配，避免误判
  local name="$1"
  if screen -ls 2>/dev/null | grep -q "\\.${name}\\s"; then
    return 0
  fi
  return 1
}

# 按端口获取监听该端口的任意一个PID，优先用 ss，其次 lsof
pids_on_port() {
  local port="$1"
  # 从 ss 提取 pid 列表
  local ss_pids
  ss_pids=$(ss -lptn 2>/dev/null \
    | grep -E ":${port}\\b" \
    | awk '{print $NF}' \
    | sed -n 's/.*pid=\([0-9]\+\),.*/\1/p' \
    | sort -u)
  if [[ -n "$ss_pids" ]]; then
    echo "$ss_pids"
    return 0
  fi
  # 从 lsof 提取 pid 列表
  if command -v lsof >/dev/null 2>&1; then
    local lsof_pids
    lsof_pids=$(lsof -nP -i :"$port" -sTCP:LISTEN -t 2>/dev/null | sort -u)
    if [[ -n "$lsof_pids" ]]; then
      echo "$lsof_pids"
      return 0
    fi
  fi
  return 1
}

# 根据端口优雅终止(TERM)并在必要时强制(KILL)清理进程
kill_by_port() {
  local port="$1"
  local pids
  pids=$(pids_on_port "$port" || true)
  if [[ -z "$pids" ]]; then
    return 0
  fi
  echo "[清理] 端口 $port 仍被占用，发送 SIGTERM 到: $pids"
  kill -TERM $pids 2>/dev/null || true
  sleep 2
  # 再次检查
  local left
  left=$(pids_on_port "$port" || true)
  if [[ -n "$left" ]]; then
    echo "[强制] 端口 $port 仍占用，发送 SIGKILL 到: $left"
    kill -KILL $left 2>/dev/null || true
  fi
}

start_service() {
  local svc="$1"
  local sname
  sname="$(session_name "$svc")"
  if [[ -z "$sname" ]]; then echo "未知服务: $svc"; return 2; fi
  if exists_session "$sname"; then
    echo "[跳过] 会话 '$sname' 已存在"
    return 0
  fi
  local cmd
  cmd="$(start_command "$svc")"
  if [[ -z "$cmd" ]]; then echo "未配置启动命令: $svc"; return 2; fi
  screen -dmS "$sname" bash -c "$cmd"
  echo "[启动] $svc -> screen 会话 '$sname'"
}

stop_service() {
  local svc="$1"
  local sname
  sname="$(session_name "$svc")"
  local port
  port="$(service_port "$svc")"
  if [[ -z "$sname" || -z "$port" ]]; then echo "未知服务: $svc"; return 2; fi

  # 1) 先尝试关闭 screen 会话
  if exists_session "$sname"; then
    screen -S "$sname" -X quit || true
    echo "[停止] $svc -> '$sname'"
  else
    echo "[提示] 未发现 screen 会话: $sname"
  fi

  # 2) 等待释放端口
  sleep 2

  # 3) 如果仍占用，按端口清理
  if ss -lptn 2>/dev/null | grep -E -q ":${port}\\b" || (command -v lsof >/dev/null 2>&1 && lsof -i :"$port" -sTCP:LISTEN >/dev/null 2>&1); then
    kill_by_port "$port"
  fi
}

status_service() {
  local svc="$1"
  local sname
  sname="$(session_name "$svc")"
  if [[ -z "$sname" ]]; then echo "未知服务: $svc"; return 2; fi
  if exists_session "$sname"; then
    echo "[运行中] $svc -> '$sname'"
  else
    echo "[未运行] $svc"
  fi
}

attach_service() {
  local svc="$1"
  local sname
  sname="$(session_name "$svc")"
  if [[ -z "$sname" ]]; then echo "未知服务: $svc"; return 2; fi
  if exists_session "$sname"; then
    echo "附着到会话: $sname (退出: Ctrl+A 然后 D)"
    screen -r "$sname"
  else
    echo "服务未运行: $svc"
    return 1
  fi
}

start_all() {
  for s in "${SERVICE_NAMES[@]}"; do
    start_service "$s"
  done
}

stop_all() {
  for s in "${SERVICE_NAMES[@]}"; do
    stop_service "$s"
  done
}

status_all() {
  for s in "${SERVICE_NAMES[@]}"; do
    status_service "$s"
  done
}

restart_service() {
  local svc="$1"
  stop_service "$svc"
  # 等待会话释放
  sleep 1
  start_service "$svc"
}

usage() {
  cat <<EOF
用法: $0 <command> [service]

command:
  start [svc]     启动指定服务；不指定时启动全部
  stop [svc]      停止指定服务；不指定时停止全部
  restart [svc]   重启指定服务；不指定时重启全部
  status          查看所有服务状态
  attach <svc>    附着到指定服务的 screen 会话
  force-stop [svc] 强制结束进程（按端口终止）；不指定时对全部执行

service 可选值:
  intent | dify | answertype | qingdan
EOF
}

main() {
  local cmd="${1:-}"; shift || true
  case "$cmd" in
    start)
      local svc="${1:-all}"
      if [[ "$svc" == "all" ]]; then start_all; else start_service "$svc"; fi
      ;;
    stop)
      local svc="${1:-all}"
      if [[ "$svc" == "all" ]]; then stop_all; else stop_service "$svc"; fi
      ;;
    restart)
      local svc="${1:-all}"
      if [[ "$svc" == "all" ]]; then
        for s in "${SERVICE_NAMES[@]}"; do restart_service "$s"; done
      else
        restart_service "$svc"
      fi
      ;;
    status)
      status_all
      ;;
    attach)
      local svc="${1:-}"
      if [[ -z "$svc" ]]; then echo "请指定服务"; usage; exit 2; fi
      attach_service "$svc"
      ;;
    force-stop)
      local svc="${1:-all}"
      if [[ "$svc" == "all" ]]; then
        for s in "${SERVICE_NAMES[@]}"; do
          # 仅按端口强制清理
          p="$(service_port "$s")"
          if [[ -n "$p" ]]; then kill_by_port "$p"; fi
        done
      else
        local p
        p="$(service_port "$svc")"
        if [[ -z "$p" ]]; then echo "未知服务: $svc"; exit 2; fi
        kill_by_port "$p"
      fi
      ;;
    ""|-h|--help|help)
      usage
      ;;
    *)
      echo "未知命令: $cmd" >&2
      usage
      exit 2
      ;;
  esac
}

main "$@"