#!/bin/bash
# 启动本地 MLX 模型服务：chat + embedding 均由 oMLX(18080) 提供
# chat=Qwen3.5-9B-MLX-4bit（4-bit），embedding=BGE-M3（oMLX 原生 BERT 系支持）
# 首次运行自动检测 .venv：缺失/损坏（如来自其他机器）时按 requirements.txt 重建，就绪则跳过安装
# 根据设备内存大小自动调整上下文窗口和最大 token 数限制
# 用法: ./scripts/start_models.sh
set -eu
cd "$(dirname "$0")/.."
ROOT=$(pwd)
VENV="$ROOT/.venv"
LOG_DIR="$HOME/.omlx/logs"
CACHE_DIR="$HOME/.omlx/cache"
SETTINGS_FILE="$HOME/.omlx/settings.json"
mkdir -p "$LOG_DIR" "$CACHE_DIR"

# 根据设备内存大小计算上下文窗口和最大 token 数
calculate_memory_limits() {
    # 获取系统总内存（字节）
    local total_memory
    total_memory=$(sysctl -n hw.memsize 2>/dev/null || echo 0)

    if [ "$total_memory" -eq 0 ]; then
        echo "[warn] 无法检测系统内存，使用默认值" >&2
        echo "32768 32768"
        return
    fi

    local total_gb=$((total_memory / 1024 / 1024 / 1024))
    echo "[info] 检测到系统内存: ${total_gb}GB" >&2

    # 计算可用内存（预留 10% 作为安全缓冲）
    local available_memory=$((total_memory * 90 / 100))
    local available_gb=$((available_memory / 1024 / 1024 / 1024))

    # 根据内存大小计算上下文窗口和最大 token 数
    # 考虑因素：
    # 1. 模型加载：Qwen3.5-9B-4bit 约 5GB
    # 2. KV 缓存：每 token 约 0.5MB（9B 模型）
    # 3. 系统保留：预留 10% 内存
    # 4. 业务场景：合同审计/证物补全
    #    - 知识库+知识图谱预检索过滤后才传递给模型
    #    - 传递给模型的数据量本身不大（ToolResultMaxChars=400，HybridTopK=5）
    #    - 单次查询模式，不需要多轮对话缓存
    #    - 16GB 机器 prefill 预算约 1900 token（2.9GB）

    local context_window=4096
    local max_tokens=2048
    local kv_cache_size="auto"

    if [ "$total_gb" -ge 32 ]; then
        # 32GB+ 内存：适中配置即可
        context_window=4096
        max_tokens=2048
        kv_cache_size="4GB"
        echo "[info] 内存充足，设置上下文窗口=${context_window}，最大 token 数=${max_tokens}，KV 缓存=${kv_cache_size}" >&2
    elif [ "$total_gb" -ge 24 ]; then
        # 24GB 内存：标准配置
        context_window=4096
        max_tokens=2048
        kv_cache_size="3GB"
        echo "[info] 中等内存配置，设置上下文窗口=${context_window}，最大 token 数=${max_tokens}，KV 缓存=${kv_cache_size}" >&2
    elif [ "$total_gb" -ge 16 ]; then
        # 16GB 内存：优化配置（符合 prefill 预算）
        context_window=4096
        max_tokens=2048
        kv_cache_size="2GB"
        echo "[info] 基础内存配置，设置上下文窗口=${context_window}，最大 token 数=${max_tokens}，KV 缓存=${kv_cache_size}" >&2
    else
        # 小于 16GB 内存：最小配置
        context_window=2048
        max_tokens=1024
        kv_cache_size="1GB"
        echo "[info] 内存较小，设置上下文窗口=${context_window}，最大 token 数=${max_tokens}，KV 缓存=${kv_cache_size}" >&2
    fi

    echo "$context_window $max_tokens $kv_cache_size"
}

# .venv 不可跨机复制（pyvenv.cfg / bin 脚本 shebang 写死本机路径），须在本机重建
if [ ! -x "$VENV/bin/python" ]; then
    if [ -d "$VENV" ]; then
        echo "[warn] $VENV 存在但 python 不可用（可能来自其他机器），删除重建"
        rm -rf "$VENV"
    fi
    if ! command -v python3 >/dev/null 2>&1; then
        echo "[error] 未找到 python3，请先安装 Python >= 3.11" >&2
        exit 1
    fi
    if ! python3 -c 'import sys; raise SystemExit(0 if sys.version_info >= (3, 11) else 1)' 2>/dev/null; then
        echo "[error] 需要 Python >= 3.11（当前 $(python3 --version 2>&1)），omlx 依赖无法安装" >&2
        exit 1
    fi
    echo "[setup] 创建 venv 并安装依赖（首次/新机器运行，见 scripts/requirements.txt）..."
    python3 -m venv "$VENV"
    "$VENV/bin/pip" install -r "$ROOT/scripts/requirements.txt"
elif [ ! -x "$VENV/bin/omlx" ]; then
    echo "[setup] venv 存在但依赖未装全，补装 scripts/requirements.txt ..."
    "$VENV/bin/pip" install -r "$ROOT/scripts/requirements.txt"
else
    echo "[ok] .venv 已就绪，跳过安装"
fi

# 根据硬件信息调整上下文窗口和最大 token 数限制
if [ -f "$SETTINGS_FILE" ]; then
    echo "[info] 检测硬件信息并调整内存限制..."
    read -r context_window max_tokens kv_cache_size <<< "$(calculate_memory_limits)"

    # 更新 settings.json 中的 sampling 参数
    if command -v python3 >/dev/null 2>&1; then
        python3 << EOF
import json
import sys

try:
    with open("$SETTINGS_FILE", 'r') as f:
        settings = json.load(f)

    # 更新 sampling 参数
    if 'sampling' not in settings:
        settings['sampling'] = {}

    old_context = settings['sampling'].get('max_context_window', 32768)
    old_max_tokens = settings['sampling'].get('max_tokens', 32768)

    settings['sampling']['max_context_window'] = $context_window
    settings['sampling']['max_tokens'] = $max_tokens

    # 更新 KV 缓存大小（如果配置了 cache 部分）
    if 'cache' not in settings:
        settings['cache'] = {}

    # 只有当 kv_cache_size 不是 "auto" 时才设置
    if "$kv_cache_size" != "auto":
        settings['cache']['ssd_cache_max_size'] = "$kv_cache_size"

    with open("$SETTINGS_FILE", 'w') as f:
        json.dump(settings, f, indent=2)

    print(f"[info] 已更新内存限制: 上下文窗口 {old_context} -> {$context_window}, 最大 token 数 {old_max_tokens} -> {$max_tokens}, KV 缓存: $kv_cache_size")
except Exception as e:
    print(f"[warn] 更新 settings.json 失败: {e}", file=sys.stderr)
EOF
    else
        echo "[warn] 未找到 python3，跳过自动配置" >&2
    fi
else
    echo "[info] settings.json 不存在，将使用默认配置"
fi

# oMLX 模型目录：chat（Qwen3.5-9B-MLX-4bit）+ embedding（bge-m3）均为实体目录，oMLX 按目录扫描注册模型
OMLX_MODEL_DIR="$ROOT/models/mlx/omlx"

if lsof -iTCP:18080 -sTCP:LISTEN >/dev/null 2>&1; then
    echo "[skip] oMLX 18080 已被占用"
else
    # 内存策略写在 ~/.omlx/settings.json（custom 上限 14GB），
    # CLI 不加 --memory-guard，避免覆盖 settings.json 的 custom 配置
    nohup "$VENV/bin/omlx" serve \
        --model-dir "$OMLX_MODEL_DIR" \
        --host 127.0.0.1 --port 18080 \
        --paged-ssd-cache-dir "$CACHE_DIR/omlx-cache" \
        --max-concurrent-requests 8 \
        --log-level info >> "$LOG_DIR/mlx-chat.log" 2>&1 &
    echo "[start] oMLX chat+embedding -> 127.0.0.1:18080 (pid $!)"
fi

echo "日志: $LOG_DIR/mlx-chat.log"