#!/bin/bash
# Qwen SpeedLab — Common configuration
# RTX 3090 24GB VRAM optimized
#
# Two inference engines available:
#   1. BeeLlama.cpp (Anbeeld) → DFlash + TurboQuant
#   2. TurboQ-MTP (Indras-Mirror) → TBQ4 fused FA + RotorQuant + MTP + tensor sharing
#
# Two model variants:
#   1. batiai (standard GGUF, no MTP heads) → baseline, DFlash
#   2. ManniX (MTP-enabled GGUF, blk.64.nextn.* preserved) → MTP self-spec

set -euo pipefail

# ─── Engine paths ──────────────────────────────────────────
BEELLAMA_DIR="/workspace/beellama.cpp"
TURBOQ_DIR="/workspace/llama.cpp-turboq-mtp"

# Default to TurboQ-MTP (more features: RotorQuant, tensor sharing, TBQ4 fused FA)
ENGINE="${ENGINE:-turboq}"
if [ "${ENGINE}" = "beellama" ]; then
    LLAMA_SERVER="${BEELLAMA_DIR}/build/bin/llama-server"
    LLAMA_CLI="${BEELLAMA_DIR}/build/bin/llama-cli"
    LLAMA_BENCH="${BEELLAMA_DIR}/build/bin/llama-bench"
else
    LLAMA_SERVER="${TURBOQ_DIR}/build/bin/llama-server"
    LLAMA_CLI="${TURBOQ_DIR}/build/bin/llama-cli"
    LLAMA_BENCH="${TURBOQ_DIR}/build/bin/llama-bench"
fi

# ─── Model files ───────────────────────────────────────────
MODELS_DIR="/workspace/models"

# Standard GGUF (no MTP heads) — for baseline & DFlash
TARGET_MODEL="${MODELS_DIR}/batiai-qwen36-27b/Qwen-Qwen3.6-27B-Q4_K_M.gguf"

# MTP-enabled GGUF (blk.64.nextn.* preserved at Q8_0) — for self-spec
MTP_MODEL="${MODELS_DIR}/mannix-mtp-qwen36-27b/Qwen3.6-27B-Omnimerge-v4-Q4_K_M.gguf"

# DFlash draft model (block diffusion spec decode)
DFLASH_DRAFT="${MODELS_DIR}/dflash/Qwen3.6-27B-DFlash-Q4_K_M.gguf"

# Multimodal projector (optional)
MMPROJ="${MODELS_DIR}/batiai-qwen36-27b/mmproj-Qwen-Qwen3.6-27B-BF16.gguf"

# ─── GPU Configuration ─────────────────────────────────────
GPU_LAYERS=99
BATCH_SIZE=2048
UBATCH_SIZE=512

# ─── KV Cache ──────────────────────────────────────────────
CTX_SIZE=102400

# ─── Server ────────────────────────────────────────────────
HOST="0.0.0.0"
PORT=8082
N_PARALLEL=1

# ─── Logging ───────────────────────────────────────────────
LOG_DIR="/workspace/logs"
mkdir -p "${LOG_DIR}"

log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*" | tee -a "${LOG_DIR}/qwen-speedlab.log"
}

wait_for_server() {
    local port=${1:-$PORT}
    local max_wait=${2:-60}
    local waited=0
    log "Waiting for server on port ${port}..."
    while ! curl -s "http://localhost:${port}/health" > /dev/null 2>&1; do
        sleep 1
        waited=$((waited + 1))
        if [ $waited -ge $max_wait ]; then
            log "ERROR: Server did not start within ${max_wait}s"
            return 1
        fi
    done
    log "Server is ready! (took ${waited}s)"
}