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

detect_cpu_quota_threads() {
  local quota period threads

  if [ -r /sys/fs/cgroup/cpu.max ]; then
    read -r quota period < /sys/fs/cgroup/cpu.max || true
    if [ "${quota:-max}" != "max" ] && [ "${period:-0}" -gt 0 ] 2>/dev/null; then
      threads=$(( (quota + period - 1) / period ))
      [ "$threads" -gt 0 ] && echo "$threads" && return
    fi
  fi

  if [ -r /sys/fs/cgroup/cpu/cpu.cfs_quota_us ] && [ -r /sys/fs/cgroup/cpu/cpu.cfs_period_us ]; then
    quota=$(cat /sys/fs/cgroup/cpu/cpu.cfs_quota_us)
    period=$(cat /sys/fs/cgroup/cpu/cpu.cfs_period_us)
    if [ "${quota:-0}" -gt 0 ] && [ "${period:-0}" -gt 0 ] 2>/dev/null; then
      threads=$(( (quota + period - 1) / period ))
      [ "$threads" -gt 0 ] && echo "$threads" && return
    fi
  fi

  echo 0
}

detect_memory_kb() {
  local value

  if [ -r /sys/fs/cgroup/memory.max ]; then
    value=$(cat /sys/fs/cgroup/memory.max)
    if [ "$value" != "max" ] && [ "${value:-0}" -gt 0 ] 2>/dev/null; then
      echo $(( value / 1024 ))
      return
    fi
  fi

  if [ -r /sys/fs/cgroup/memory/memory.limit_in_bytes ]; then
    value=$(cat /sys/fs/cgroup/memory/memory.limit_in_bytes)
    if [ "${value:-0}" -gt 0 ] && [ "$value" -lt 9223372036854771712 ] 2>/dev/null; then
      echo $(( value / 1024 ))
      return
    fi
  fi

  if [ -r /proc/meminfo ]; then
    awk '/MemTotal/ {print $2}' /proc/meminfo || echo 0
  else
    echo 0
  fi
}

normalize_arg_value() {
  local value="$1"

  value="${value//$'\r'/}"
  value="${value//$'\n'/}"
  value="${value#"${value%%[![:space:]]*}"}"
  value="${value%"${value##*[![:space:]]}"}"

  # Hugging Face env values are plain values. A copied shell-style "$q4_0"
  # should be treated as "q4_0" instead of crashing llama-server.
  value="${value#\$}"

  printf '%s' "$value"
}

HOST_THREADS="$(nproc)"
QUOTA_THREADS="$(detect_cpu_quota_threads)"
if [ -z "${THREADS:-}" ]; then
  if [ "$QUOTA_THREADS" -gt 0 ] && [ "$QUOTA_THREADS" -lt "$HOST_THREADS" ]; then
    THREADS="$QUOTA_THREADS"
  else
    THREADS="$HOST_THREADS"
  fi
fi

export OMP_NUM_THREADS="$THREADS"
export OPENBLAS_NUM_THREADS="$THREADS"
export OMP_PROC_BIND="${OMP_PROC_BIND:-FALSE}"
export OMP_PLACES="${OMP_PLACES:-cores}"
export OMP_WAIT_POLICY="${OMP_WAIT_POLICY:-PASSIVE}"
export MKL_NUM_THREADS="$THREADS"
export VECLIB_MAXIMUM_THREADS="$THREADS"
export NUMEXPR_NUM_THREADS="$THREADS"
export HF_HOME="${HF_HOME:-/data/hf-cache}"
export MODEL_DIR="${MODEL_DIR:-/data/models}"
export MODEL_PATH="${MODEL_PATH:-/data/models/model.gguf}"
REQUESTED_CTX_SIZE="${CONTEXT_LENGTH:-${CTX_SIZE:-}}"
export QUANT_PREFERENCE="${QUANT_PREFERENCE:-q4_k_m,q4_k_s,q4_k,q4_0,q4_1}"
export PERF_PROFILE="${PERF_PROFILE:-balanced}"
REQUESTED_CACHE_TYPE_K="${CACHE_TYPE_K:-}"
REQUESTED_CACHE_TYPE_V="${CACHE_TYPE_V:-}"
export REASONING="${REASONING:-auto}"
export TOOLS="${TOOLS:-}"
export ENABLE_TOOLS="${ENABLE_TOOLS:-0}"
export HTTP_THREADS="${HTTP_THREADS:-1}"
export LOG_VERBOSITY="${LOG_VERBOSITY:-3}"
export MMAP="${MMAP:-0}"
export FLASH_ATTN="${FLASH_ATTN:-1}"
export NO_WARMUP="${NO_WARMUP:-1}"
export LANGSEARCH_API_KEY="${LANGSEARCH_API_KEY:-}"

# Ensure data directories exist and are writable. Run as root in container startup.
echo "preparing storage: $HF_HOME and $MODEL_DIR"
install -d -m 0777 "$HF_HOME" "$MODEL_DIR" || mkdir -p "$HF_HOME" "$MODEL_DIR"
chmod 0777 "$HF_HOME" "$MODEL_DIR" || true

# Ensure HF cache logging directories are writable for xet
install -d -m 0777 "$HF_HOME/xet/logs" || true
chmod 0777 "$HF_HOME/xet" "$HF_HOME/xet/logs" 2>/dev/null || true
chown -R 65532:65532 "$HF_HOME/xet" 2>/dev/null || true

# Run downloader as the non-root `appuser` using the venv python when available.
echo "starting model downloader"
if [ -x "/opt/venv/bin/python" ]; then
  su -s /bin/sh appuser -c "/opt/venv/bin/python -u /app/download_model.py"
else
  su -s /bin/sh appuser -c "python3 -u /app/download_model.py"
fi

# If downloader wrote a repo marker with selected filename, update MODEL_PATH accordingly
if [ -f "$MODEL_DIR/model.repo" ]; then
  marker=$(cat "$MODEL_DIR/model.repo" || echo "")
  if [[ "$marker" == *"|"* ]]; then
    IFS='|' read -r _ selected_file _ <<< "$marker"
    if [ -n "$selected_file" ]; then
      MODEL_PATH="$MODEL_DIR/$selected_file"
      echo "using downloaded model path: $MODEL_PATH"
    fi
  fi
fi

LLAMA_SERVER_BIN="${LLAMA_SERVER_BIN:-/usr/local/bin/llama-server}"

# set server profile defaults
case "$PERF_PROFILE" in
  low_latency)
    DEFAULT_BATCH=64
    DEFAULT_UBATCH=64
    DEFAULT_THREADS_BATCH=${THREADS}
    ;;
  balanced)
    DEFAULT_BATCH=512
    DEFAULT_UBATCH=512
    DEFAULT_THREADS_BATCH=${THREADS}
    ;;
  throughput)
    DEFAULT_BATCH=512
    DEFAULT_UBATCH=512
    DEFAULT_THREADS_BATCH=${THREADS}
    ;;
  *)
    DEFAULT_BATCH=128
    DEFAULT_UBATCH=128
    DEFAULT_THREADS_BATCH=${THREADS}
    ;;
esac

# Adjust defaults for low-memory environments (Spaces free tier: ~16GB).
# Use cgroup limits first because /proc/meminfo can report the host machine.
MEM_KB="$(detect_memory_kb)"
echo "detected effective memory (KB): $MEM_KB"
if [ "$MEM_KB" -gt 0 ] && [ "$MEM_KB" -le 18000000 ]; then
  echo "low-memory profile detected: using compact defaults"
  DEFAULT_CTX_SIZE=2048
else
  DEFAULT_CTX_SIZE=4096
fi

CTX_SIZE="${REQUESTED_CTX_SIZE:-$DEFAULT_CTX_SIZE}"
CACHE_TYPE_K="$(normalize_arg_value "${REQUESTED_CACHE_TYPE_K:-f16}")"
CACHE_TYPE_V="$(normalize_arg_value "${REQUESTED_CACHE_TYPE_V:-f16}")"
PARALLEL=1

echo "effective threads: $THREADS (host=$HOST_THREADS, quota=$QUOTA_THREADS)"
echo "effective llama profile: perf=$PERF_PROFILE single_request=1 ctx=$CTX_SIZE batch=${BATCH_SIZE:-$DEFAULT_BATCH} ubatch=${UBATCH_SIZE:-$DEFAULT_UBATCH} http_threads=$HTTP_THREADS log_verbosity=$LOG_VERBOSITY mmap=$MMAP flash_attn=$FLASH_ATTN"

export CACHE_TYPE_K CACHE_TYPE_V CTX_SIZE PARALLEL

server_args=(
  --model "$MODEL_PATH"
  --host 0.0.0.0
  --port "${PORT:-7860}"
  --api-key "$API_PASSWORD"
  -ngl 0
  -t "$THREADS"
  --threads-batch "${THREADS_BATCH:-$DEFAULT_THREADS_BATCH}"
  --threads-http "$HTTP_THREADS"
  --batch-size "${BATCH_SIZE:-$DEFAULT_BATCH}"
  --ubatch-size "${UBATCH_SIZE:-$DEFAULT_UBATCH}"
  --ctx-size "${CTX_SIZE:-4096}"
  --cache-type-k "${CACHE_TYPE_K}"
  --cache-type-v "${CACHE_TYPE_V}"
  --metrics
  -lv "$LOG_VERBOSITY"
)

if [ "$MMAP" = "1" ]; then
  server_args+=(--mmap)
else
  server_args+=(--no-mmap)
fi

if [ "$FLASH_ATTN" = "1" ]; then
  server_args+=(--flash-attn on)
fi

if [ "${NO_WARMUP}" = "1" ]; then
  server_args+=(--no-warmup)
fi

case "${REASONING,,}" in
  true|1|on|yes)
    server_args+=(--reasoning on)
    ;;
  false|0|off|no)
    server_args+=(--reasoning off)
    ;;
  auto|"")
    server_args+=(--reasoning auto)
    ;;
  *)
    echo "invalid REASONING value: $REASONING (expected True/False/auto)"
    exit 1
    ;;
esac

if [ "$ENABLE_TOOLS" = "1" ] && [ -n "$LANGSEARCH_API_KEY" ] && [ -z "$TOOLS" ]; then
  TOOLS="exec_shell_command"
fi

if [[ -n "${TOOLS}" ]]; then
  server_args+=(--tools "$TOOLS")
  echo "enabled tools: $TOOLS"
fi

server_args+=(--parallel "$PARALLEL")

# mmproj (vision projector) support
MMPROJ_PATH=""
if [ -n "${MMPROJ_FILE:-}" ]; then
  MMPROJ_PATH="$MODEL_DIR/$MMPROJ_FILE"
elif [ -f "$MODEL_DIR/model.mmproj" ]; then
  mmproj_name=$(cat "$MODEL_DIR/model.mmproj" || echo "")
  if [ -n "$mmproj_name" ]; then
    MMPROJ_PATH="$MODEL_DIR/$mmproj_name"
  fi
fi

if [ -n "$MMPROJ_PATH" ] && [ -f "$MMPROJ_PATH" ]; then
  server_args+=(--mmproj "$MMPROJ_PATH")
  echo "vision support enabled: $MMPROJ_PATH"
else
  echo "no mmproj found; running text-only"
fi

server_cmd=("$LLAMA_SERVER_BIN" "${server_args[@]}")
echo "server command: $LLAMA_SERVER_BIN ${server_args[*]}"
echo "starting llama-server"
exec su -s /bin/sh appuser -c "$(printf '%q ' "${server_cmd[@]}")"