#!/usr/bin/env sh
#
# Proposa AI 服务部署脚本（Linux 服务器）
#
# 功能：
#   0. 检测到 uv 缺失时自动安装（可用 AUTO_INSTALL_UV=0 关闭）
#   1. 使用 ModelScope 将 gpt2-chinese-cluecorpussmall-onnx 下载到 models/ 目录
#      （模型文件已存在时自动跳过，幂等）
#   2. 使用 uv 按 uv.lock 安装整个项目依赖（uv sync --frozen）
#   3. 后台启动 proposa-api，写入 PID/日志文件并做 /health 健康检查
#
# 用法：
#   sh deploy.sh            完整部署并后台启动（可重复执行）
#   sh deploy.sh status     查看服务运行状态与健康检查
#   sh deploy.sh stop       停止后台服务
#   sh deploy.sh restart    停止后重新后台启动
#
# 可用环境变量：
#   MODEL_REPO    ModelScope 模型仓库 ID（默认 Maiteka/gpt2-chinese-cluecorpussmall-onnx）
#   MODEL_DIR     模型保存目录（默认 <仓库根>/models/gpt2-chinese-cluecorpussmall-onnx）
#   UV_CACHE_DIR  uv 缓存目录（默认 <仓库根>/.uv-cache）
#   UV_INDEX_URL  PyPI 镜像源（默认清华 https://pypi.tuna.tsinghua.edu.cn/simple）
#   PROPOSA_API_PORT  健康检查端口（默认读取 .env 的 PROPOSA_API_PORT，缺省 8000）
#   RESTART=1     服务已运行时强制重启（等价于 sh deploy.sh restart）
#   AUTO_INSTALL_UV   1=uv 缺失时自动安装（默认）；0=只提示不安装

set -eu

case "$0" in
  */*) SCRIPT_DIR=${0%/*} ;;
  *) SCRIPT_DIR=. ;;
esac
REPO_ROOT=$(CDPATH= cd "$SCRIPT_DIR" && pwd)
cd "$REPO_ROOT"

MODEL_REPO=${MODEL_REPO:-"Maiteka/gpt2-chinese-cluecorpussmall-onnx"}
MODEL_DIR=${MODEL_DIR:-"$REPO_ROOT/models/gpt2-chinese-cluecorpussmall-onnx"}
UV_CACHE_DIR=${UV_CACHE_DIR:-"$REPO_ROOT/.uv-cache"}
export UV_CACHE_DIR
UV_INDEX_URL=${UV_INDEX_URL:-"https://pypi.tuna.tsinghua.edu.cn/simple"}
export UV_INDEX_URL

ENV_FILE="$REPO_ROOT/.env"
PID_FILE="$REPO_ROOT/output/api.pid"
LOG_FILE="$REPO_ROOT/output/api.log"

usage() {
  cat <<'EOF'
Proposa AI 服务部署脚本（Linux 服务器）

用法：
  sh deploy.sh            完整部署并后台启动（可重复执行）
  sh deploy.sh status     查看服务运行状态与健康检查
  sh deploy.sh stop       停止后台服务
  sh deploy.sh restart    停止后重新后台启动

环境变量：
  MODEL_REPO    ModelScope 模型仓库 ID（默认 Maiteka/gpt2-chinese-cluecorpussmall-onnx）
  MODEL_DIR     模型保存目录（默认 <仓库根>/models/gpt2-chinese-cluecorpussmall-onnx）
  UV_CACHE_DIR  uv 缓存目录（默认 <仓库根>/.uv-cache）
  UV_INDEX_URL  PyPI 镜像源（默认清华 https://pypi.tuna.tsinghua.edu.cn/simple）
  PROPOSA_API_PORT  健康检查端口（默认读取 .env 的 PROPOSA_API_PORT，缺省 8000）
  RESTART=1     服务已运行时强制重启（等价于 sh deploy.sh restart）
  AUTO_INSTALL_UV   1=uv 缺失时自动安装（默认）；0=只提示不安装
EOF
}

# 从 .env 读取 PROPOSA_API_PORT（容忍行首空格），读不到输出空串
env_port() {
  [ -f "$ENV_FILE" ] || return 0
  line=$(grep -E '^[[:space:]]*PROPOSA_API_PORT[[:space:]]*=' "$ENV_FILE" 2>/dev/null | tail -n 1 || true)
  [ -n "$line" ] || return 0
  printf '%s\n' "$line" | sed 's/^[[:space:]]*PROPOSA_API_PORT[[:space:]]*=[[:space:]]*//' | tr -d '[:space:]"'
}

API_PORT=${PROPOSA_API_PORT:-$(env_port)}
API_PORT=${API_PORT:-8000}
case "$API_PORT" in
  ''|*[!0-9]*)
    echo "错误: 无效的 API 端口 '$API_PORT'，请设置 PROPOSA_API_PORT 或检查 .env" >&2
    exit 2
    ;;
esac
HEALTH_URL="http://127.0.0.1:$API_PORT/health"

require_uv() {
  if command -v uv >/dev/null 2>&1; then
    return 0
  fi
  # uv 已安装（例如自动安装到 ~/.local/bin）但当前 PATH 没有时，补进 PATH
  if [ -x "$HOME/.local/bin/uv" ]; then
    export PATH="$HOME/.local/bin:$PATH"
    return 0
  fi
  if [ "${AUTO_INSTALL_UV:-1}" != "1" ]; then
    echo "错误: 未找到 uv。已设置 AUTO_INSTALL_UV=0 跳过自动安装。" >&2
    echo "FIX: 手动安装 uv（curl -LsSf https://astral.sh/uv/install.sh | sh），或执行 AUTO_INSTALL_UV=1 sh deploy.sh" >&2
    exit 127
  fi
  install_uv
  if ! command -v uv >/dev/null 2>&1; then
    echo "错误: uv 自动安装后仍未在 PATH 中找到。" >&2
    echo "FIX: 手动执行 curl -LsSf https://astral.sh/uv/install.sh | sh，然后重跑 sh deploy.sh" >&2
    exit 127
  fi
}

# 自动安装 uv：优先官方安装脚本，失败时回退 python3 -m pip
install_uv() {
  echo "未找到 uv，开始自动安装……"
  if ! command -v curl >/dev/null 2>&1 && ! command -v wget >/dev/null 2>&1; then
    echo "错误: 自动安装需要 curl 或 wget，但两者都不存在。" >&2
    exit 127
  fi

  if command -v curl >/dev/null 2>&1; then
    curl -LsSf https://astral.sh/uv/install.sh | sh || true
  else
    wget -qO- https://astral.sh/uv/install.sh | sh || true
  fi

  # 官方安装脚本把 uv 装到 ~/.local/bin，补进当前会话 PATH
  if [ -x "$HOME/.local/bin/uv" ]; then
    export PATH="$HOME/.local/bin:$PATH"
  fi

  if ! command -v uv >/dev/null 2>&1 && command -v python3 >/dev/null 2>&1; then
    echo "官方安装脚本未生效，尝试 python3 -m pip install --user uv"
    python3 -m pip install --user uv || true
    if [ -x "$HOME/.local/bin/uv" ]; then
      export PATH="$HOME/.local/bin:$PATH"
    fi
  fi

  if command -v uv >/dev/null 2>&1; then
    echo "uv 已就绪: $(command -v uv)"
    uv --version 2>/dev/null || true
  fi
}

# ---- 1. ModelScope 下载 gpt2 模型到 models/ ----
download_model() {
  if [ -s "$MODEL_DIR/model.onnx" ] && [ -s "$MODEL_DIR/tokenizer.json" ]; then
    echo "[1/3] 模型已存在，跳过下载: $MODEL_DIR"
    return 0
  fi

  echo "[1/3] 使用 ModelScope 下载 $MODEL_REPO -> $MODEL_DIR"
  mkdir -p "$MODEL_DIR"
  # 在临时环境安装 modelscope 并执行下载，不污染项目依赖
  uv run --no-project --with modelscope python - "$MODEL_REPO" "$MODEL_DIR" <<'PY'
import os
import shutil
import sys

from modelscope import snapshot_download

repo = sys.argv[1]
target = sys.argv[2]
os.makedirs(target, exist_ok=True)

try:
    snapshot_download(repo, local_dir=target)
except TypeError:
    # 旧版 modelscope 不支持 local_dir：下载到缓存后复制到目标目录
    cached = snapshot_download(repo)
    for name in os.listdir(cached):
        src = os.path.join(cached, name)
        dst = os.path.join(target, name)
        if os.path.isdir(src):
            shutil.copytree(src, dst, dirs_exist_ok=True)
        else:
            shutil.copy2(src, dst)
PY

  if [ ! -s "$MODEL_DIR/model.onnx" ] || [ ! -s "$MODEL_DIR/tokenizer.json" ]; then
    echo "错误: 模型下载后缺少 model.onnx 或 tokenizer.json: $MODEL_DIR" >&2
    return 1
  fi
  echo "模型就绪: $MODEL_DIR"
}

# ---- 2. uv 安装项目依赖 ----
install_deps() {
  echo "[2/3] 使用 uv 安装项目依赖（uv sync --frozen）"
  uv sync --frozen

  if [ ! -f "$ENV_FILE" ]; then
    cp "$REPO_ROOT/.env.example" "$ENV_FILE"
    echo "注意: .env 不存在，已从 .env.example 复制。"
    echo "请编辑 .env 填写 DEEPSEEK_API_KEY 等配置，然后执行: sh deploy.sh restart"
  fi
}

# ---- 服务进程管理 ----
is_running() {
  [ -f "$PID_FILE" ] || return 1
  pid=$(cat "$PID_FILE" 2>/dev/null || true)
  [ -n "$pid" ] || return 1
  kill -0 "$pid" 2>/dev/null
}

wait_healthy() {
  tries=0
  while [ "$tries" -lt 30 ]; do
    if command -v curl >/dev/null 2>&1; then
      if curl -fsS "$HEALTH_URL" >/dev/null 2>&1; then
        return 0
      fi
    elif is_running; then
      sleep 2
      return 0
    fi
    tries=$((tries + 1))
    sleep 1
  done
  return 1
}

start_service() {
  # restart/后台启动同样需要 uv，避免新 shell 中 uv 不在 PATH 时 nohup 启动失败
  require_uv

  if is_running; then
    echo "服务已在运行（PID $(cat "$PID_FILE")），跳过启动；如需重启请执行: sh deploy.sh restart"
    return 0
  fi

  rm -f "$PID_FILE"
  mkdir -p "$REPO_ROOT/output"
  echo "[3/3] 后台启动 proposa-api（端口 $API_PORT）"
  echo "日志: $LOG_FILE"
  nohup uv run --frozen proposa-api >>"$LOG_FILE" 2>&1 &
  echo $! > "$PID_FILE"

  if wait_healthy; then
    echo "健康检查通过: $HEALTH_URL"
  else
    echo "警告: 30 秒内健康检查未通过，请查看日志: $LOG_FILE" >&2
    tail -n 20 "$LOG_FILE" >&2 || true
    return 1
  fi
}

stop_service() {
  if ! is_running; then
    echo "服务未在运行"
    rm -f "$PID_FILE"
    return 0
  fi

  pid=$(cat "$PID_FILE")
  echo "停止 proposa-api（PID $pid）"
  kill "$pid" 2>/dev/null || true

  tries=0
  while [ "$tries" -lt 10 ] && is_running; do
    sleep 1
    tries=$((tries + 1))
  done
  if is_running; then
    echo "进程未在 10 秒内退出，强制终止" >&2
    kill -9 "$pid" 2>/dev/null || true
  fi
  rm -f "$PID_FILE"
  echo "已停止"
}

status_service() {
  if is_running; then
    echo "运行中: PID $(cat "$PID_FILE")"
    echo "日志: $LOG_FILE"
    echo "健康检查: $HEALTH_URL"
    if command -v curl >/dev/null 2>&1; then
      curl -fsS "$HEALTH_URL" 2>/dev/null || echo "（/health 暂不可达）"
      echo
    fi
  else
    echo "未运行"
  fi
}

deploy() {
  echo "== Proposa AI 部署 =="
  echo "仓库: $REPO_ROOT"
  echo "模型: $MODEL_REPO"
  echo "uv 缓存: $UV_CACHE_DIR"
  echo
  require_uv
  download_model
  install_deps
  if [ "${RESTART:-0}" = "1" ]; then
    stop_service
  fi
  start_service
}

ACTION=${1:-deploy}
case "$ACTION" in
  deploy)
    deploy
    ;;
  status)
    status_service
    ;;
  stop)
    stop_service
    ;;
  restart)
    stop_service
    start_service
    ;;
  -h|--help|help)
    usage
    ;;
  *)
    echo "错误: 未知命令 '$ACTION'" >&2
    usage >&2
    exit 2
    ;;
esac