#!/bin/bash
#
# Unified service lifecycle controller for saas-search.
# Supports: start / stop / restart / status
#

set -euo pipefail

PROJECT_ROOT="$(cd "$(dirname "$0")/.." && pwd)"
LOG_DIR="${PROJECT_ROOT}/logs"

mkdir -p "${LOG_DIR}"

# shellcheck source=scripts/lib/load_env.sh
source "${PROJECT_ROOT}/scripts/lib/load_env.sh"

CORE_SERVICES=("backend" "indexer" "frontend")
OPTIONAL_SERVICES=("embedding" "translator" "reranker" "tei" "cnclip")

all_services() {
  echo "${CORE_SERVICES[@]} ${OPTIONAL_SERVICES[@]}"
}

get_port() {
  local service="$1"
  case "${service}" in
    backend) echo "${API_PORT:-6002}" ;;
    indexer) echo "${INDEXER_PORT:-6004}" ;;
    frontend) echo "${FRONTEND_PORT:-6003}" ;;
    embedding) echo "${EMBEDDING_PORT:-6005}" ;;
    translator) echo "${TRANSLATION_PORT:-6006}" ;;
    reranker) echo "${RERANKER_PORT:-6007}" ;;
    tei) echo "${TEI_PORT:-8080}" ;;
    cnclip) echo "${CNCLIP_PORT:-51000}" ;;
    *) echo "" ;;
  esac
}

pid_file() {
  local service="$1"
  echo "${LOG_DIR}/${service}.pid"
}

log_file() {
  local service="$1"
  echo "${LOG_DIR}/${service}.log"
}

service_start_cmd() {
  local service="$1"
  case "${service}" in
    backend) echo "./scripts/start_backend.sh" ;;
    indexer) echo "./scripts/start_indexer.sh" ;;
    frontend) echo "./scripts/start_frontend.sh" ;;
    embedding) echo "./scripts/start_embedding_service.sh" ;;
    translator) echo "./scripts/start_translator.sh" ;;
    reranker) echo "./scripts/start_reranker.sh" ;;
    tei) echo "./scripts/start_tei_service.sh" ;;
    cnclip) echo "./scripts/start_cnclip_service.sh" ;;
    *) return 1 ;;
  esac
}

service_exists() {
  local service="$1"
  case "${service}" in
    backend|indexer|frontend|embedding|translator|reranker|tei|cnclip) return 0 ;;
    *) return 1 ;;
  esac
}

validate_targets() {
  local targets="$1"
  for svc in ${targets}; do
    if ! service_exists "${svc}"; then
      echo "[error] unknown service: ${svc}" >&2
      return 1
    fi
  done
}

wait_for_health() {
  local service="$1"
  local max_retries="${2:-30}"
  local interval_sec="${3:-1}"
  local port
  port="$(get_port "${service}")"
  local path="/health"

  case "${service}" in
    backend) path="/health" ;;
    indexer) path="/health" ;;
    frontend) path="/" ;;
    embedding) path="/health" ;;
    translator) path="/health" ;;
    reranker) path="/health" ;;
    tei) path="/health" ;;
    *) return 0 ;;
  esac

  local i=0
  while [ "${i}" -lt "${max_retries}" ]; do
    if curl -sf "http://127.0.0.1:${port}${path}" >/dev/null 2>&1; then
      return 0
    fi
    i=$((i + 1))
    sleep "${interval_sec}"
  done
  return 1
}

wait_for_stable_health() {
  local service="$1"
  local checks="${2:-3}"
  local interval_sec="${3:-1}"
  local port
  port="$(get_port "${service}")"
  local path="/health"

  local i=0
  while [ "${i}" -lt "${checks}" ]; do
    if ! is_running_by_port "${service}"; then
      return 1
    fi
    if ! curl -sf "http://127.0.0.1:${port}${path}" >/dev/null 2>&1; then
      return 1
    fi
    i=$((i + 1))
    sleep "${interval_sec}"
  done
  return 0
}

is_running_by_pid() {
  local service="$1"
  local pf
  pf="$(pid_file "${service}")"
  if [ ! -f "${pf}" ]; then
    return 1
  fi
  local pid
  pid="$(cat "${pf}" 2>/dev/null || true)"
  [ -n "${pid}" ] && kill -0 "${pid}" 2>/dev/null
}

is_running_by_port() {
  local service="$1"
  local port
  port="$(get_port "${service}")"
  [ -n "${port}" ] && lsof -ti:"${port}" >/dev/null 2>&1
}

is_running_tei_container() {
  local tei_name="${TEI_CONTAINER_NAME:-saas-search-tei}"
  local cid
  cid="$(docker ps -q -f name=^/${tei_name}$ 2>/dev/null || true)"
  [ -n "${cid}" ]
}

get_cnclip_flow_device() {
  local flow_file="${PROJECT_ROOT}/third-party/clip-as-service/server/torch-flow-temp.yml"
  if [ ! -f "${flow_file}" ]; then
    return 1
  fi
  sed -n "s/^[[:space:]]*device:[[:space:]]*'\\([^']*\\)'.*/\\1/p" "${flow_file}" | head -n 1
}

start_one() {
  local service="$1"
  cd "${PROJECT_ROOT}"
  local cmd
  if ! cmd="$(service_start_cmd "${service}")"; then
    echo "[error] unknown service: ${service}" >&2
    return 1
  fi
  local pf lf
  pf="$(pid_file "${service}")"
  lf="$(log_file "${service}")"

  if [ "${service}" != "tei" ]; then
    if is_running_by_pid "${service}" || is_running_by_port "${service}"; then
      if [ "${service}" = "cnclip" ]; then
        local expected_device="${CNCLIP_DEVICE:-cuda}"
        expected_device="$(echo "${expected_device}" | tr '[:upper:]' '[:lower:]')"
        if [[ "${expected_device}" != "cuda" && "${expected_device}" != "cpu" ]]; then
          echo "[error] invalid CNCLIP_DEVICE=${CNCLIP_DEVICE}; use cuda/cpu" >&2
          return 1
        fi
        local actual_device
        actual_device="$(get_cnclip_flow_device 2>/dev/null || true)"
        if [ -n "${actual_device}" ] && [ "${actual_device}" != "${expected_device}" ]; then
          echo "[error] cnclip already running with device=${actual_device}, expected=${expected_device}" >&2
          echo "[error] run: ./scripts/service_ctl.sh stop cnclip && CNCLIP_DEVICE=${expected_device} ./scripts/service_ctl.sh start cnclip" >&2
          return 1
        fi
      fi
      echo "[skip] ${service} already running"
      return 0
    fi
  fi

  case "${service}" in
    cnclip|tei)
      echo "[start] ${service} (managed by native script)"
      if [ "${service}" = "cnclip" ]; then
        if ! CNCLIP_DEVICE="${CNCLIP_DEVICE:-cuda}" bash -lc "${cmd}" >> "${lf}" 2>&1; then
          echo "[error] ${service} start script failed, inspect ${lf}" >&2
          return 1
        fi
      else
        if ! bash -lc "${cmd}" >> "${lf}" 2>&1; then
          echo "[error] ${service} start script failed, inspect ${lf}" >&2
          return 1
        fi
      fi
      if [ "${service}" = "tei" ]; then
        if is_running_tei_container; then
          echo "[ok] ${service} started (log=${lf})"
        else
          echo "[error] ${service} failed to start, inspect ${lf}" >&2
          return 1
        fi
      elif is_running_by_pid "${service}" || is_running_by_port "${service}"; then
        echo "[ok] ${service} started (log=${lf})"
      else
        echo "[error] ${service} failed to start, inspect ${lf}" >&2
        return 1
      fi
      ;;
    backend|indexer|frontend|embedding|translator)
      echo "[start] ${service}"
      nohup bash -lc "${cmd}" > "${lf}" 2>&1 &
      local pid=$!
      echo "${pid}" > "${pf}"
      local retries=30
      if wait_for_health "${service}" "${retries}"; then
        echo "[ok] ${service} healthy (pid=${pid}, log=${lf})"
      else
        echo "[error] ${service} health check timeout, inspect ${lf}" >&2
        return 1
      fi
      ;;
    reranker)
      echo "[start] ${service}"
      # Start reranker directly so pid file points to the script process that
      # will exec uvicorn, avoiding extra shell wrapper lifecycle edge-cases.
      nohup "${cmd}" > "${lf}" 2>&1 &
      local pid=$!
      echo "${pid}" > "${pf}"
      if wait_for_health "${service}" 90; then
        if wait_for_stable_health "${service}" 5 1; then
          echo "[ok] ${service} healthy (pid=${pid}, log=${lf})"
        else
          echo "[error] ${service} became unavailable right after startup, inspect ${lf}" >&2
          return 1
        fi
      else
        echo "[error] ${service} health check timeout, inspect ${lf}" >&2
        return 1
      fi
      ;;
    *)
      echo "[warn] ${service} unsupported start path"
      ;;
  esac
}

cleanup_reranker_orphans() {
  local engine_pids
  engine_pids="$(pgrep -f 'VLLM::EngineCore' 2>/dev/null || true)"
  if [ -z "${engine_pids}" ]; then
    return 0
  fi

  echo "[stop] reranker orphan engines=${engine_pids}"
  for pid in ${engine_pids}; do
    kill -TERM "${pid}" 2>/dev/null || true
  done
  sleep 1
  engine_pids="$(pgrep -f 'VLLM::EngineCore' 2>/dev/null || true)"
  for pid in ${engine_pids}; do
    kill -KILL "${pid}" 2>/dev/null || true
  done
}

stop_one() {
  local service="$1"
  cd "${PROJECT_ROOT}"
  if [ "${service}" = "cnclip" ]; then
    echo "[stop] cnclip (managed by native script)"
    bash -lc "./scripts/stop_cnclip_service.sh" || true
    return 0
  fi
  if [ "${service}" = "tei" ]; then
    echo "[stop] tei (managed by native script)"
    bash -lc "./scripts/stop_tei_service.sh" || true
    return 0
  fi

  local pf
  pf="$(pid_file "${service}")"

  if [ -f "${pf}" ]; then
    local pid
    pid="$(cat "${pf}" 2>/dev/null || true)"
    if [ -n "${pid}" ] && kill -0 "${pid}" 2>/dev/null; then
      echo "[stop] ${service} pid=${pid}"
      kill -TERM "${pid}" 2>/dev/null || true
      sleep 1
      if kill -0 "${pid}" 2>/dev/null; then
        kill -KILL "${pid}" 2>/dev/null || true
      fi
    fi
    rm -f "${pf}"
  fi

  local port
  port="$(get_port "${service}")"
  if [ -n "${port}" ]; then
    local pids
    pids="$(lsof -ti:${port} 2>/dev/null || true)"
    if [ -n "${pids}" ]; then
      echo "[stop] ${service} port=${port} pids=${pids}"
      for pid in ${pids}; do
        kill -TERM "${pid}" 2>/dev/null || true
      done
      sleep 1
      pids="$(lsof -ti:${port} 2>/dev/null || true)"
      for pid in ${pids}; do
        kill -KILL "${pid}" 2>/dev/null || true
      done
    fi
  fi

  if [ "${service}" = "reranker" ]; then
    cleanup_reranker_orphans
  fi
}

status_one() {
  local service="$1"
  local port
  port="$(get_port "${service}")"
  local running="no"
  local pid_info="-"

  if [ "${service}" = "tei" ]; then
    local cid
    local tei_name="${TEI_CONTAINER_NAME:-saas-search-tei}"
    cid="$(docker ps -q -f name=^/${tei_name}$ 2>/dev/null || true)"
    if [ -n "${cid}" ]; then
      running="yes"
      pid_info="${cid:0:12}"
    fi
    printf "%-10s running=%-3s port=%-6s pid=%s\n" "${service}" "${running}" "${port:--}" "${pid_info}"
    return
  fi

  if is_running_by_pid "${service}"; then
    running="yes"
    pid_info="$(cat "$(pid_file "${service}")" 2>/dev/null || echo "-")"
  elif is_running_by_port "${service}"; then
    running="yes"
    pid_info="$(lsof -ti:${port} 2>/dev/null | tr '\n' ',' | sed 's/,$//' || echo "-")"
  fi

  printf "%-10s running=%-3s port=%-6s pid=%s\n" "${service}" "${running}" "${port:--}" "${pid_info}"
}

resolve_targets() {
  local scope="$1"
  shift || true

  if [ "$#" -gt 0 ]; then
    echo "$*"
    return
  fi

  case "${scope}" in
    start)
      echo "${CORE_SERVICES[@]}"
      ;;
    stop|status)
      echo "$(all_services)"
      ;;
    restart)
      # Restart with no explicit services uses the same explicit start target order.
      echo "$(resolve_targets start)"
      ;;
    *)
      echo ""
      ;;
  esac
}

usage() {
  cat <<'EOF'
Usage:
  ./scripts/service_ctl.sh start [service...]
  ./scripts/service_ctl.sh stop [service...]
  ./scripts/service_ctl.sh restart [service...]
  ./scripts/service_ctl.sh status [service...]

Default target set (when no service provided):
  start   -> backend indexer frontend
  stop    -> all known services
  restart -> stop all known services, then start with start targets
  status  -> all known services

Optional service startup:
  ./scripts/service_ctl.sh start embedding translator reranker tei cnclip
  TEI_DEVICE=cuda|cpu ./scripts/service_ctl.sh start tei
  CNCLIP_DEVICE=cuda|cpu ./scripts/service_ctl.sh start cnclip
EOF
}

main() {
  if [ "$#" -lt 1 ]; then
    usage
    exit 1
  fi

  local action="$1"
  shift || true

  load_env_file "${PROJECT_ROOT}/.env"
  local stop_targets=""
  local targets
  # For restart without explicit services, stop everything first, then start
  # with the default start target order.
  if [ "${action}" = "restart" ] && [ "$#" -eq 0 ]; then
    stop_targets="$(resolve_targets stop)"
  fi
  targets="$(resolve_targets "${action}" "$@")"
  if [ -z "${targets}" ]; then
    usage
    exit 1
  fi
  validate_targets "${targets}"

  case "${action}" in
    start)
      for svc in ${targets}; do
        start_one "${svc}"
      done
      ;;
    stop)
      for svc in ${targets}; do
        stop_one "${svc}"
      done
      ;;
    restart)
      for svc in ${stop_targets:-${targets}}; do
        stop_one "${svc}"
      done
      for svc in ${targets}; do
        start_one "${svc}"
      done
      ;;
    status)
      for svc in ${targets}; do
        status_one "${svc}"
      done
      ;;
    *)
      usage
      exit 1
      ;;
  esac
}

main "$@"