#!/bin/bash
# Build llama.cpp at patches/BASE_COMMIT with the flint CPU patch (Xeon with AVX512-VNNI + AMX),
# then serve the Q4_1 model on http://127.0.0.1:20000 and point the pi coding agent at it.
#
# Overridable environment variables:
#   MODEL_FILE    (default flint-int4-q4_1.gguf)
#   LLAMA_PORT    (default 20000)
#   MODEL_ALIAS   (default flint-q4_1)
#   CONTEXT_SIZE  (default 32768)
#   LLAMA_THREADS (default nproc)
#   LLAMA_HOST    (default 127.0.0.1)
set -euo pipefail

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

MODEL_FILE="${MODEL_FILE:-flint-int4-q4_1.gguf}"
MODEL_REPO="${MODEL_REPO:-HelloSun/flint_cpu_opt}"
LLAMA_PORT="${LLAMA_PORT:-20000}"
MODEL_ALIAS="${MODEL_ALIAS:-flint-q4_1}"
CONTEXT_SIZE="${CONTEXT_SIZE:-32768}"
LLAMA_THREADS="${LLAMA_THREADS:-$(nproc 2>/dev/null || echo 4)}"
LLAMA_HOST="${LLAMA_HOST:-127.0.0.1}"
MODEL_PATH="$ROOT/$MODEL_FILE"
SERVER_BIN="$ROOT/llama.cpp/build/bin/llama-server"
PIDFILE="$HOME/llama-server.pid"
LOGFILE="$HOME/llama-server.log"
PI_MODELS_JSON="$HOME/.pi/agent/models.json"
PI_PROVIDER_ID="llama-cpp-local"
PI_MODEL_ID="$MODEL_ALIAS"

info() { echo -e "\033[36m[build.sh]\033[0m $*"; }
ok()   { echo -e "\033[32m[build.sh]\033[0m $*"; }
warn() { echo -e "\033[33m[build.sh]\033[0m $*" >&2; }

mkdir -p "$HOME/.local/bin" "$(dirname "$PI_MODELS_JSON")"

# ── 1. Build llama.cpp (skip when the binary is already there) ────────────────
if [[ ! -x "$SERVER_BIN" ]]; then
    if [[ ! -f "$ROOT/patches/BASE_COMMIT" || ! -f "$ROOT/patches/flint-cpu-opt.patch" ]]; then
        warn "缺少 patches/BASE_COMMIT 或 patches/flint-cpu-opt.patch，無法套用 flint patch"
        exit 1
    fi
    if [[ ! -d "$ROOT/llama.cpp/.git" ]]; then
        info "clone llama.cpp"
        git clone https://github.com/ggml-org/llama.cpp "$ROOT/llama.cpp"
    fi
    cd "$ROOT/llama.cpp"
    BASE_COMMIT="$(cat "$ROOT/patches/BASE_COMMIT")"
    info "checkout llama.cpp base commit $BASE_COMMIT"
    git fetch --depth 1 origin "$BASE_COMMIT" || true
    git checkout "$BASE_COMMIT"
    if ! git apply --reverse --check "$ROOT/patches/flint-cpu-opt.patch" 2>/dev/null; then
        git apply "$ROOT/patches/flint-cpu-opt.patch"
    fi
    info "cmake configure (GGML_NATIVE=ON, CUDA off)"
    cmake -B build -DCMAKE_BUILD_TYPE=Release -DGGML_NATIVE=ON -DGGML_CUDA=OFF -DLLAMA_CURL=OFF
    info "cmake build (llama-server, -j$LLAMA_THREADS)"
    cmake --build build -j"$LLAMA_THREADS" --target llama-server
    cd "$ROOT"
else
    ok "llama-server 已存在，跳過編譯：$SERVER_BIN"
fi

# ── 2. Make sure the Q4_1 model is present ───────────────────────────────────
if [[ ! -s "$MODEL_PATH" ]]; then
    info "下載模型 $MODEL_REPO/$MODEL_FILE（約 24 GiB，支援續傳）"
    URL="https://huggingface.co/$MODEL_REPO/resolve/main/$MODEL_FILE?download=true"
    # 直接用 curl 續傳（hf CLI 在部分映像會因 typer/click 版本衝突而崩潰）。
    # HF_TOKEN 有設定時帶上 Authorization（私有 repo 或避免匿名限速）。
    CURL_AUTH=()
    if [[ -n "${HF_TOKEN:-}" ]]; then
        CURL_AUTH=(-H "Authorization: Bearer $HF_TOKEN")
    fi
    curl -fL -C - --retry 3 --retry-delay 5 "${CURL_AUTH[@]}" "$URL" -o "$MODEL_PATH" \
        || warn "curl 下載失敗（exit=$?）"
fi
if [[ ! -s "$MODEL_PATH" ]]; then
    warn "模型不存在或下載失敗：$MODEL_PATH"
    exit 1
fi
ok "模型：$MODEL_PATH"

# ── 3. Install pi if needed, and configure it to use this llama-server ───────
export PATH="$HOME/.local/bin:$HOME/.pi/bin:${XDG_DATA_HOME:-$HOME/.local/share}/pi-node/current/bin:$PATH"
# pi 1.x 起要求 Node >= 22.19；Space 映像是 Node 20，npm 的 latest 會裝到啟動即崩潰的版本
# （SyntaxError: enableCompileCache）。依目前 Node 版本選擇：新 Node 用 latest，
# 舊 Node 用 npm dist-tag legacy-node20（0.74.2）。並且安裝後實際執行 pi --version 驗證。
PI_PKG="@earendil-works/pi-coding-agent"
pi_runs() { command -v pi >/dev/null 2>&1 && pi --version >/dev/null 2>&1; }
if ! pi_runs; then
    if command -v npm >/dev/null 2>&1 && command -v node >/dev/null 2>&1; then
        if node -e 'const [a,b]=process.versions.node.split(".").map(Number); process.exit(a>22||(a===22&&b>=19)?0:1)'; then
            PI_SPEC="$PI_PKG@latest"
        else
            PI_SPEC="$PI_PKG@legacy-node20"
        fi
        info "安裝 pi coding agent（$PI_SPEC）到 $HOME/.local（Node $(node -v)）"
        # 舊版殘留的 bin 連結會讓 npm 拒絕覆寫（EEXIST），先移除。
        rm -f "$HOME/.local/bin/pi"
        npm install -g --ignore-scripts --no-fund --no-audit --progress=false \
            --prefix "$HOME/.local" "$PI_SPEC" || warn "pi 安裝失敗"
        hash -r
        if pi_runs; then
            ok "pi 可執行：$(pi --version 2>&1 | head -n1)"
        else
            warn "pi 安裝後仍無法執行（pi --version 失敗）"
        fi
    else
        warn "找不到 npm/node，略過 pi 安裝（僅寫入設定）"
    fi
else
    ok "pi 已可執行：$(pi --version 2>&1 | head -n1)"
fi

info "設定 pi 使用本地 llama-server：http://127.0.0.1:$LLAMA_PORT/v1"
python3 - "$PI_MODELS_JSON" "$LLAMA_PORT" "$PI_MODEL_ID" "$PI_PROVIDER_ID" "$CONTEXT_SIZE" "$MODEL_ALIAS" <<'PY'
import json, os, sys
path, port, model_id, provider, ctx_size, alias = sys.argv[1:7]
ctx_size = int(ctx_size)
os.makedirs(os.path.dirname(path), exist_ok=True)
cfg = {"providers": {}}
if os.path.exists(path):
    try:
        with open(path) as f:
            loaded = json.load(f)
        if isinstance(loaded, dict):
            cfg = loaded
            cfg.setdefault("providers", {})
    except Exception as exc:
        print(f"[WARN] 無法解析舊 models.json，將覆蓋：{exc}")
cfg["providers"][provider] = {
    "baseUrl": f"http://127.0.0.1:{port}/v1",
    "api": "openai-completions",
    "apiKey": "not-needed",
    "compat": {"supportsDeveloperRole": True, "supportsReasoningEffort": False},
    "models": [{
        "id": model_id,
        "name": f"Flint Q4_1 (Local {alias})",
        "reasoning": True,
        "input": ["text"],
        "contextWindow": ctx_size,
        "maxTokens": ctx_size,
        "cost": {"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0},
    }],
}
with open(path, "w") as f:
    json.dump(cfg, f, indent=2)
print(f"[build.sh] 已寫入 {path}")
PY

# 讓 pi 預設就用這個 llama-server（不只知道 provider，還要設成啟動預設）
PI_SETTINGS_JSON="$HOME/.pi/agent/settings.json"
info "設定 pi 預設模型：$PI_PROVIDER_ID/$PI_MODEL_ID"
python3 - "$PI_SETTINGS_JSON" "$PI_PROVIDER_ID" "$PI_MODEL_ID" <<'PY'
import json, os, sys
path, provider, model = sys.argv[1:4]
os.makedirs(os.path.dirname(path), exist_ok=True)
settings = {}
if os.path.exists(path):
    try:
        with open(path) as f:
            loaded = json.load(f)
        if isinstance(loaded, dict):
            settings = loaded
    except Exception as exc:
        print(f"[WARN] 無法解析舊 settings.json，將覆蓋：{exc}")
settings["defaultProvider"] = provider
settings["defaultModel"] = model
enabled = settings.get("enabledModels")
if not isinstance(enabled, list) or not enabled:
    settings["enabledModels"] = [f"{provider}/{model}"]
with open(path, "w") as f:
    json.dump(settings, f, indent=2)
print(f"[build.sh] 已寫入 {path}")
PY

WRAPPER="$HOME/.local/bin/pi-local"
cat > "$WRAPPER" <<EOS
#!/usr/bin/env bash
exec pi --provider $PI_PROVIDER_ID --model $PI_MODEL_ID "\$@"
EOS
chmod +x "$WRAPPER"
ok "pi-local wrapper：$WRAPPER"

# ── 4. Start llama-server on port 20000 ──────────────────────────────────────
if curl -sf "http://127.0.0.1:$LLAMA_PORT/health" >/dev/null 2>&1; then
    ok "llama-server 已在 127.0.0.1:$LLAMA_PORT 執行，沿用"
else
    if [[ -f "$PIDFILE" ]]; then
        OLD_PID="$(cat "$PIDFILE" 2>/dev/null || true)"
        if [[ -n "$OLD_PID" ]] && kill -0 "$OLD_PID" 2>/dev/null; then
            info "停止舊 llama-server (PID $OLD_PID)"
            kill "$OLD_PID" 2>/dev/null || true
            sleep 2
            kill -9 "$OLD_PID" 2>/dev/null || true
        fi
    fi
    info "啟動 llama-server：127.0.0.1:$LLAMA_PORT（model=$MODEL_FILE, alias=$MODEL_ALIAS, ctx=$CONTEXT_SIZE, threads=$LLAMA_THREADS）"
    nohup "$SERVER_BIN" \
        -m "$MODEL_PATH" \
        --alias "$MODEL_ALIAS" \
        -t "$LLAMA_THREADS" \
        -fa on \
        --jinja \
        -c "$CONTEXT_SIZE" \
        --host "$LLAMA_HOST" \
        --port "$LLAMA_PORT" \
        >"$LOGFILE" 2>&1 &
    echo $! > "$PIDFILE"
    info "llama-server PID: $(cat "$PIDFILE")，log: $LOGFILE"

    READY=0
    for _ in $(seq 1 120); do
        if curl -sf "http://127.0.0.1:$LLAMA_PORT/health" >/dev/null 2>&1; then
            READY=1
            break
        fi
        if ! kill -0 "$(cat "$PIDFILE")" 2>/dev/null; then
            warn "llama-server 程序已結束，最後 40 行 log："
            tail -n 40 "$LOGFILE" >&2 || true
            break
        fi
        sleep 5
    done
    if [[ "$READY" -ne 1 ]]; then
        warn "llama-server 未就緒（見 $LOGFILE）"
        exit 1
    fi
fi

ok "完成：llama-server http://127.0.0.1:$LLAMA_PORT/v1（model=$PI_MODEL_ID, provider=$PI_PROVIDER_ID）"
ok "pi 用法：pi --provider $PI_PROVIDER_ID --model $PI_MODEL_ID  或  pi-local"