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

# If the first argument is not an option, treat it as a full command override.
if [[ $# -gt 0 && "${1#-}" == "$1" ]]; then
    exec "$@"
fi

# Allow bypassing defaults for fully manual llama-server control.
if [[ "${LLAMA_DISABLE_DEFAULT_ARGS:-0}" == "1" ]]; then
    exec llama-server "$@"
fi

host="${LLAMA_HOST:-0.0.0.0}"
port="${LLAMA_PORT:-8080}"
preset="${LLAMA_PRESET:-latency}"
ngl="${LLAMA_NGL:-all}"
flash_attn="${LLAMA_FLASH_ATTN:-auto}"
cache_type_k="${LLAMA_CACHE_TYPE_K:-q8_0}"
cache_type_v="${LLAMA_CACHE_TYPE_V:-q8_0}"
reasoning_format="${LLAMA_REASONING_FORMAT:-deepseek}"

case "${preset}" in
    latency)
        parallel="${LLAMA_PARALLEL:-2}"
        ctx_size="${LLAMA_CTX_SIZE:-131072}"
        batch_size="${LLAMA_BATCH_SIZE:-2048}"
        ubatch_size="${LLAMA_UBATCH_SIZE:-1024}"
        ;;
    throughput)
        parallel="${LLAMA_PARALLEL:-4}"
        ctx_size="${LLAMA_CTX_SIZE:-262144}"
        batch_size="${LLAMA_BATCH_SIZE:-4096}"
        ubatch_size="${LLAMA_UBATCH_SIZE:-2048}"
        ;;
    *)
        echo "Unknown LLAMA_PRESET='${preset}'. Use 'latency' or 'throughput'." >&2
        exit 2
        ;;
esac

default_args=(
    --host "${host}"
    --port "${port}"
    --swa-full
    -ngl "${ngl}"
    --flash-attn "${flash_attn}"
    --parallel "${parallel}"
    --ctx-size "${ctx_size}"
    -b "${batch_size}"
    -ub "${ubatch_size}"
    -ctk "${cache_type_k}"
    -ctv "${cache_type_v}"
    --kv-unified
    --jinja
    --reasoning-format "${reasoning_format}"
    --no-context-shift
    --mlock
    --metrics
)

# Optional thread controls (recommended to set explicitly per host if needed).
if [[ -n "${LLAMA_THREADS:-}" ]]; then
    default_args+=(-t "${LLAMA_THREADS}")
fi
if [[ -n "${LLAMA_THREADS_BATCH:-}" ]]; then
    default_args+=(--threads-batch "${LLAMA_THREADS_BATCH}")
fi
if [[ -n "${LLAMA_THREADS_HTTP:-}" ]]; then
    default_args+=(--threads-http "${LLAMA_THREADS_HTTP}")
fi

# Optional compatibility switch for memory mapping behavior.
if [[ "${LLAMA_NO_MMAP:-0}" == "1" ]]; then
    default_args+=(--no-mmap)
fi

# User-provided args come last so they can override defaults.
exec llama-server "${default_args[@]}" "$@"
