#!/usr/bin/env bash
# loosecanvas single-container entrypoint: native llama-server on loopback +
# uvicorn on the one public port. Health-gates the app on the server and
# fail-fasts if the server dies on boot (so HF restarts instead of zombieing).
set -euo pipefail

cd /home/user/app

VENV_BIN="${VENV_BIN:-/opt/venv/bin}"
MODEL_DIR="${MODEL_DIR:-/home/user/app/models}"
MODEL_FILE="${GEMMA_MODEL_FILENAME:?GEMMA_MODEL_FILENAME must be set}"
MODEL_PATH="${MODEL_DIR}/${MODEL_FILE}"
MTP_FILE="${MTP_MODEL_FILE:-mtp-gemma-4-26B-A4B-it.gguf}"
MTP_PATH="${MODEL_DIR}/${MTP_FILE}"
LLAMA_PORT="${LLAMA_CPP_PORT:-8080}"
CTX="${LLAMA_CPP_CTX:-32768}"
APP_PORT="${APP_PORT:-7860}"
USE_MTP="${LLAMA_CPP_MTP:-1}"
N_MAX="${SPEC_DRAFT_N_MAX:-2}"
KV_TYPE="${LLAMA_CPP_KV_TYPE:-f16}"  # q8_0 ≈ lossless, halves KV-cache VRAM → ~2x ctx; f16 = validated default
HEALTH_TIMEOUT="${LLAMA_HEALTH_TIMEOUT:-900}"  # model load on a cold GPU is slow

# Keep the port the server binds identical to the port the app dials.
export LLAMA_CPP_PORT="$LLAMA_PORT"

LLAMA_BIN="$(command -v llama-server || echo /app/llama-server)"
[ -x "$LLAMA_BIN" ] || { echo "[entrypoint] FATAL: llama-server not found ($LLAMA_BIN)"; exit 1; }

# Runtime download if the model wasn't baked (BAKE_MODEL=0). Needs a public repo
# (default) or HF_TOKEN for a gated one. Guard inputs so a missing var fails
# loudly instead of an opaque KeyError under set -e.
if [ ! -f "$MODEL_PATH" ]; then
  : "${MODEL_REPO:?MODEL_REPO required to download the model when not baked}"
  echo "[entrypoint] model missing at $MODEL_PATH — downloading from $MODEL_REPO"
  "${VENV_BIN}/python" - <<'PY'
import os, sys
from huggingface_hub import hf_hub_download
d = os.environ["MODEL_DIR"]
tok = os.environ.get("HF_TOKEN")
try:
    hf_hub_download(os.environ["MODEL_REPO"], os.environ["GEMMA_MODEL_FILENAME"], local_dir=d, token=tok)
    if os.environ.get("LLAMA_CPP_MTP", "1") != "0":
        hf_hub_download(os.environ.get("MTP_REPO", os.environ["MODEL_REPO"]),
                        os.environ["MTP_MODEL_FILE"], local_dir=d, token=tok)
except Exception as e:  # noqa: BLE001
    print(f"[entrypoint] FATAL: model download failed: {e}", file=sys.stderr)
    sys.exit(1)
PY
fi

# Mirror the validated docker-compose flags (structured output is build-sensitive).
# --reasoning off is confirmed valid on this image (the validated local server
# runs with it). MTP flags per plan/mtp-gemma-investigation.md.
LLAMA_ARGS=(
  -m "$MODEL_PATH"
  --host 127.0.0.1 --port "$LLAMA_PORT"
  -c "$CTX"
  -ngl all --flash-attn on --cont-batching --cache-prompt
  --parallel 1 --jinja --reasoning off --no-mmap --metrics
)

# KV-cache quantization: q8_0 ≈ lossless and halves KV-cache VRAM, roughly doubling
# the context that fits (~128K on a 24 GB L4). Requires flash-attn (set above).
# Default f16 = the validated config; set LLAMA_CPP_KV_TYPE=q8_0 to enable.
if [ "$KV_TYPE" != "f16" ]; then
  echo "[entrypoint] KV-cache type=$KV_TYPE (quantized; flash-attn required)"
  LLAMA_ARGS+=(--cache-type-k "$KV_TYPE" --cache-type-v "$KV_TYPE")
fi

# MTP speculative decoding (~1.4-2.2x decode, +~2 GB VRAM). On for the 24 GB L4;
# set LLAMA_CPP_MTP=0 on a 16 GB GPU. draft-mtp is the current arg (not "mtp").
if [ "$USE_MTP" != "0" ] && [ -f "$MTP_PATH" ]; then
  echo "[entrypoint] MTP on — draft head $MTP_PATH (spec-draft-n-max=$N_MAX)"
  LLAMA_ARGS+=(--model-draft "$MTP_PATH" --spec-type draft-mtp --spec-draft-n-max "$N_MAX")
elif [ "$USE_MTP" != "0" ]; then
  echo "[entrypoint] MTP requested but draft head missing ($MTP_PATH); running without MTP"
fi

echo "[entrypoint] starting llama-server on 127.0.0.1:${LLAMA_PORT} (ctx=${CTX})"
"$LLAMA_BIN" "${LLAMA_ARGS[@]}" &
LLAMA_PID=$!

echo "[entrypoint] waiting up to ${HEALTH_TIMEOUT}s for llama-server /health..."
healthy=0
elapsed=0
while [ "$elapsed" -lt "$HEALTH_TIMEOUT" ]; do
  if ! kill -0 "$LLAMA_PID" 2>/dev/null; then
    echo "[entrypoint] FATAL: llama-server exited during startup (check flags / VRAM / model)"
    wait "$LLAMA_PID" || true
    exit 1
  fi
  if curl -fsS --max-time 5 "http://127.0.0.1:${LLAMA_PORT}/health" >/dev/null 2>&1; then
    healthy=1; break
  fi
  sleep 1
  elapsed=$((elapsed + 1))
done
if [ "$healthy" != "1" ]; then
  echo "[entrypoint] FATAL: llama-server not healthy within ${HEALTH_TIMEOUT}s"
  kill "$LLAMA_PID" 2>/dev/null || true
  exit 1
fi
echo "[entrypoint] llama-server healthy — starting uvicorn on 0.0.0.0:${APP_PORT}"

exec "${VENV_BIN}/uvicorn" loosecanvas.main:app --host 0.0.0.0 --port "$APP_PORT"