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

APP_DIR="/opt/vmqos"
PYTHON="python3"
DOWNLOAD_URL="https://pan.halocloud.net/?f=/vmqos/vmqos.zip"  # TODO: 如需更换地址请修改此处
TMP_ZIP="/tmp/vmqos.zip"

if [[ $EUID -ne 0 ]]; then
  echo "请用 root 运行此脚本"
  exit 1
fi
echo "===> 安装系统依赖"
apt-get update -y
apt-get install -y python3.13-venv redis-server unzip

echo "===> 准备目录 $APP_DIR"
mkdir -p "$APP_DIR"
if [[ -n "$(ls -A "$APP_DIR")" ]]; then
  read -r -p "$APP_DIR 已存在文件，继续可能覆盖，是否继续? [y/N] " go
  if [[ "${go,,}" != "y" ]]; then
    echo "已取消"
    exit 1
  fi
fi

echo "===> 下载程序包（如失败请替换 DOWNLOAD_URL）"
curl -fL "$DOWNLOAD_URL" -o "$TMP_ZIP"

echo "===> 输入解压密码（必填）"
read -r -p "ZIP password: " ZIP_PW
if [[ -z "$ZIP_PW" ]]; then
  echo "ZIP 密码必填"
  exit 1
fi

rm -rf "${APP_DIR:?}/"*
unzip -o -P "$ZIP_PW" "$TMP_ZIP" -d "$APP_DIR"

cd "$APP_DIR"

echo "===> 创建虚拟环境"
if [[ ! -d ".venv" ]]; then
  "$PYTHON" -m venv .venv
fi
source .venv/bin/activate
pip install -U pip

echo "===> 安装 Python 依赖"
if [[ -f requirements.txt ]]; then
  pip install -r requirements.txt
else
  pip install redis pymysql
fi

echo "===> 配置数据库连接（全部必填，无默认值）"
read -r -p "DB host: " DB_HOST
read -r -p "DB port: " DB_PORT
read -r -p "DB user: " DB_USER
read -r -p "DB password: " DB_PASS
read -r -p "DB name: " DB_NAME

if [[ -z "$DB_HOST" || -z "$DB_PORT" || -z "$DB_USER" || -z "$DB_PASS" || -z "$DB_NAME" ]]; then
  echo "DB 配置项均为必填，不能留空"
  exit 1
fi
if ! [[ "$DB_PORT" =~ ^[0-9]+$ ]]; then
  echo "DB port 必须为数字"
  exit 1
fi

DB_HOST="$DB_HOST" DB_PORT="$DB_PORT" DB_USER="$DB_USER" DB_PASS="$DB_PASS" DB_NAME="$DB_NAME" APP_DIR="$APP_DIR" python3 - <<'PY'
from pathlib import Path
import os
import re

cfg = Path(os.environ["APP_DIR"]) / "config.py"
text = cfg.read_text()

match = re.search(r'DB\s*=\s*{.*?}', text, flags=re.S)
if not match:
    raise SystemExit("未找到 DB 配置块")

db_block = (
    "DB = {\n"
    f'    "host": "{os.environ["DB_HOST"]}",\n'
    f'    "port": {int(os.environ["DB_PORT"])},\n'
    f'    "user": "{os.environ["DB_USER"]}",\n'
    f'    "password": "{os.environ["DB_PASS"]}",\n'
    f'    "name": "{os.environ["DB_NAME"]}",\n'
    '    "clusters": [],            # 动态填充：根据 vmbr0 IP 查询 ms_server_cluster\n'
    '    "service_status": 1,       # 仅处理 status = 1 的服务\n'
    "}\n"
)

text = text[:match.start()] + db_block + text[match.end():]
cfg.write_text(text)
PY

echo "===> 写入 systemd 服务"
cat >/etc/systemd/system/vmqos.service <<'EOF'
[Unit]
Description=HaloCloud VMQOS Dynamic Bandwidth Management
After=network-online.target
Wants=network-online.target

[Service]
WorkingDirectory=/opt/vmqos
ExecStart=/opt/vmqos/.venv/bin/python /opt/vmqos/vmqos.py
Restart=always
RestartSec=3
User=root
LimitNOFILE=65535

[Install]
WantedBy=multi-user.target
EOF

echo "===> 重新加载并启用服务"
systemctl daemon-reload
systemctl enable --now  redis-server
systemctl restart redis-server
systemctl enable --now  vmqos
systemctl restart vmqos

echo "完成！可用 'systemctl status vmqos' 查看运行状态，'journalctl -fu vmqos' 查看日志。"