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

########################################
# Defaults
########################################

PORT="${PORT:-7860}"
HOST="${HOST:-0.0.0.0}"

CTX_SIZE="${CTX_SIZE:-2048}"
THREADS="${THREADS:-2}"
N_GPU_LAYERS="${N_GPU_LAYERS:-0}"

MODEL_REPO_ID="${MODEL_REPO_ID:-bldeaw/m1_planner_q4_k_m}"
MODEL_FILE="${MODEL_FILE:-m1_planner_q4_k_m.gguf}"
MODEL_REVISION="${MODEL_REVISION:-main}"

HF_TOKEN="${HF_TOKEN:-}"

# persistent path (HF Storage Bucket mount)
MODEL_DIR="${MODEL_DIR:-/data/models}"

# optional checksum
MODEL_SHA256="${MODEL_SHA256:-}"

########################################
# Validate config
########################################

if [[ -z "$MODEL_REPO_ID" ]]; then
  echo "ERROR: MODEL_REPO_ID missing"
  exit 1
fi

if [[ -z "$MODEL_FILE" ]]; then
  echo "ERROR: MODEL_FILE missing"
  exit 1
fi

mkdir -p "$MODEL_DIR"

MODEL_PATH="$MODEL_DIR/$MODEL_FILE"

########################################
# Download model if missing
########################################

download_model() {
  URL="https://huggingface.co/${MODEL_REPO_ID}/resolve/${MODEL_REVISION}/${MODEL_FILE}"

  echo "Downloading model..."
  echo "Repo      : $MODEL_REPO_ID"
  echo "File      : $MODEL_FILE"
  echo "Revision  : $MODEL_REVISION"

  if [[ -n "$HF_TOKEN" ]]; then
    curl -L \
      --retry 5 \
      --retry-delay 5 \
      -H "Authorization: Bearer ${HF_TOKEN}" \
      "$URL" \
      -o "$MODEL_PATH"
  else
    curl -L \
      --retry 5 \
      --retry-delay 5 \
      "$URL" \
      -o "$MODEL_PATH"
  fi
}

if [[ ! -f "$MODEL_PATH" ]]; then
  download_model
else
  echo "Model already exists: $MODEL_PATH"
fi

########################################
# Verify checksum (optional)
########################################

if [[ -n "$MODEL_SHA256" ]]; then
  echo "Verifying checksum..."
  ACTUAL=$(sha256sum "$MODEL_PATH" | awk '{print $1}')

  if [[ "$ACTUAL" != "$MODEL_SHA256" ]]; then
    echo "ERROR: checksum mismatch"
    echo "Expected: $MODEL_SHA256"
    echo "Actual  : $ACTUAL"
    exit 1
  fi
fi

########################################
# Info
########################################

echo "===================================="
echo "Model Path   : $MODEL_PATH"
echo "Context Size : $CTX_SIZE"
echo "Threads      : $THREADS"
echo "GPU Layers   : $N_GPU_LAYERS"
echo "Port         : $PORT"
echo "===================================="

########################################
# Start server
########################################

exec /app/llama-server \
  --model "$MODEL_PATH" \
  --host "$HOST" \
  --port "$PORT" \
  --ctx-size "$CTX_SIZE" \
  --threads "$THREADS" \
  --n-gpu-layers "$N_GPU_LAYERS"