#!/usr/bin/env bash
# ==============================================================================
# LLM Runner: Unified Inference Server Orchestrator for llama-server
# Target Hardware: AMD ROCm / Radeon AI PRO R9700 (~28-30GB VRAM)
# ==============================================================================

set -euo pipefail

# Sanitize environment:
# TEMP/TMP are standard POSIX variables for temporary file directories.
# If set to a sampling temperature (float/number), ROCm/libstdc++ crashes during initialization!
if [ -n "${TEMP:-}" ] && [ ! -d "${TEMP}" ]; then
    TEMPERATURE="${TEMPERATURE:-${TEMP}}"
    unset TEMP
fi
if [ -n "${TMP:-}" ] && [ ! -d "${TMP}" ]; then
    unset TMP
fi

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
LOG_DIR="${LLM_LOG_DIR:-${LOG_DIR:-${SCRIPT_DIR}/logs}}"
mkdir -p "${LOG_DIR}"

# ------------------------------------------------------------------------------
# Helpers & Color Output
# ------------------------------------------------------------------------------
BOLD="\033[1m"
GREEN="\033[32m"
BLUE="\033[34m"
YELLOW="\033[33m"
RED="\033[31m"
CYAN="\033[36m"
RESET="\033[0m"

log_info()    { echo -e "${BLUE}[INFO]${RESET} $*"; }
log_success() { echo -e "${GREEN}[✓]${RESET} $*"; }
log_warn()    { echo -e "${YELLOW}[!]${RESET} $*"; }
log_error()   { echo -e "${RED}[✗]${RESET} $*" >&2; }

# Find llama-server binary
find_llama_server() {
    if command -v llama-server >/dev/null 2>&1; then
        command -v llama-server
    elif [ -x "/run/current-system/sw/bin/llama-server" ]; then
        echo "/run/current-system/sw/bin/llama-server"
    else
        log_error "llama-server executable not found in PATH or system profile!"
        exit 1
    fi
}
LLAMA_SERVER="$(find_llama_server)"

# ------------------------------------------------------------------------------
# Status & Inspection
# ------------------------------------------------------------------------------
check_health() {
    local port="${1:-11434}"
    local host="${2:-127.0.0.1}"
    if curl -s "http://${host}:${port}/health" 2>/dev/null | grep -q '"status":\s*"ok"\|"ok"'; then
        return 0
    fi
    return 1
}

cmd_status() {
    echo -e "${BOLD}=== Local LLM Server Status ===${RESET}"
    local ports=(11434 11435 11436)
    local any_running=0

    for port in "${ports[@]}"; do
        local pid
        pid=$(lsof -ti :"$port" 2>/dev/null || fuser "$port/tcp" 2>/dev/null || true)
        if [ -n "$pid" ]; then
            any_running=1
            local model_info
            model_info=$(curl -s "http://127.0.0.1:${port}/v1/models" 2>/dev/null | jq -r '.data[0].id // "unknown"' 2>/dev/null || echo "initializing...")
            echo -e "  Port ${CYAN}${port}${RESET}: ${GREEN}RUNNING${RESET} (PID: ${pid}) - Model: ${BOLD}${model_info}${RESET}"
        else
            echo -e "  Port ${CYAN}${port}${RESET}: ${YELLOW}STOPPED${RESET}"
        fi
    done

    echo ""
    return 0
}

# ------------------------------------------------------------------------------
# Stop Servers
# ------------------------------------------------------------------------------
cmd_stop() {
    local target="${1:-all}"
    log_info "Stopping server(s): ${target}..."

    if [ "$target" = "all" ]; then
        pkill -f "llama-server" 2>/dev/null || true
        log_success "All llama-server processes terminated."
    else
        fuser -k -9 "$target/tcp" 2>/dev/null || true
        log_success "Killed server on port ${target}."
    fi
}

# ------------------------------------------------------------------------------
# Launch Single Server Instance
# ------------------------------------------------------------------------------
launch_instance() {
    local name="$1"
    local model="$2"
    local port="$3"
    local ctx="$4"
    local mmproj="${5:-}"
    local template="${6:-}"
    local temp="${7:-0.6}"
    local top_p="${8:-0.95}"
    local top_k="${9:-40}"
    local min_p="${10:-0.05}"
    local ngl="${11:-99}"
    local load_mode="${12:-none}"
    local flash_attn="${13:-on}"
    local ctk="${14:-q8_0}"
    local ctv="${15:-q8_0}"
    local reasoning="${16:-false}"
    local log_file="${LOG_DIR}/${name}.log"

    if [ ! -f "$model" ]; then
        log_error "Model file not found: ${model}"
        exit 1
    fi

    # Ensure target port is free
    if fuser "$port/tcp" >/dev/null 2>&1; then
        log_warn "Port ${port} is in use. Stopping existing process on port ${port}..."
        fuser -k -9 "$port/tcp" 2>/dev/null || true
        sleep 1
    fi

    local flags=(
        -m "${model}"
        --host "${LLM_HOST:-${HOST:-0.0.0.0}}"
        --port "${port}"
        --ctx-size "${ctx}"
        -np "${LLM_PARALLEL:-${PARALLEL:-1}}"
        --n-gpu-layers "${ngl}"
        --load-mode "${load_mode}"
        --temp "${temp}"
        --top-p "${top_p}"
        --top-k "${top_k}"
        --min-p "${min_p}"
        -ctk "${ctk}"
        -ctv "${ctv}"
        -fa "${flash_attn}"
        -b "${LLM_BATCH_SIZE:-${BATCH_SIZE:-2048}}"
        -ub "${LLM_UBATCH_SIZE:-${UBATCH_SIZE:-512}}"
    )

    if [ -n "${mmproj}" ] && [ -f "${mmproj}" ]; then
        flags+=(--mmproj "${mmproj}" --image-min-tokens "${LLM_IMAGE_MIN_TOKENS:-${IMAGE_MIN_TOKENS:-1024}}")
    fi

    if [ -n "${template}" ] && [ -f "${template}" ]; then
        flags+=(--chat-template-file "${template}" --jinja)
    fi

    if [ "${reasoning}" = "true" ]; then
        flags+=(--reasoning-preserve)
    fi

    log_info "Starting ${BOLD}${name}${RESET} on port ${CYAN}${port}${RESET} (CTX: ${ctx})..."
    "${LLAMA_SERVER}" "${flags[@]}" > "${log_file}" 2>&1 &
    local s_pid=$!

    # Wait for initialization
    log_info "Waiting for ${name} to initialize (PID: ${s_pid})..."
    while true; do
        if ! kill -0 "${s_pid}" 2>/dev/null; then
            log_error "${name} failed to start! Last 20 log lines from ${log_file}:"
            tail -n 20 "${log_file}" >&2
            exit 1
        fi
        if check_health "${port}"; then
            break
        fi
        sleep 1
    done

    log_success "${name} is READY at http://localhost:${port} (PID: ${s_pid})"
    echo "${s_pid}"
}

# ------------------------------------------------------------------------------
# Presets
# ------------------------------------------------------------------------------
cmd_cluster() {
    log_info "Starting Multi-Agent Cluster (SuperQwen 27B + 2x Qwen 4B)..."
    pkill -f "llama-server" 2>/dev/null || true
    sleep 1

    # Trap to gracefully stop all background instances on Ctrl+C
    cleanup_cluster() {
        log_warn "Stopping cluster instances..."
        pkill -P $$ 2>/dev/null || true
        pkill -f "llama-server" 2>/dev/null || true
        exit 0
    }
    trap cleanup_cluster SIGINT SIGTERM EXIT

    # SuperQwen 3.8 27B (Planner)
    local planner_model="${LLM_PLANNER_MODEL:-${PLANNER_MODEL:-${LLM_MODEL:-${MODEL:-${SCRIPT_DIR}/models/Qwen3.8/Qwen3.8-27B-UD-Q4_K_XL.gguf}}}}"
    if [ ! -f "${planner_model}" ] && [ -f "${SCRIPT_DIR}/models/Qwen3.8/SuperQwen3.8-27b-abliterated-Q4_K_M.gguf" ]; then
        planner_model="${SCRIPT_DIR}/models/Qwen3.8/SuperQwen3.8-27b-abliterated-Q4_K_M.gguf"
    fi
    local planner_mmproj="${LLM_PLANNER_MMPROJ:-${PLANNER_MMPROJ:-${SCRIPT_DIR}/models/Qwen3.8/mmproj-SuperQwen3.8-27b-abliterated-Q8_0.gguf}}"
    if [ "${planner_mmproj}" = "none" ]; then planner_mmproj=""; fi
    local planner_tmpl="${LLM_PLANNER_TEMPLATE:-${PLANNER_TEMPLATE:-${SCRIPT_DIR}/templates/Qwen3.8/chat_template.jinja}}"

    # Qwen 3.5 4B (Workers)
    local worker_model="${LLM_WORKER_MODEL:-${WORKER_MODEL:-${SCRIPT_DIR}/models/Qwen3.5/Qwen3.5-4B-Q5_K_M.gguf}}"
    local worker_mmproj="${LLM_WORKER_MMPROJ:-${WORKER_MMPROJ:-${SCRIPT_DIR}/models/Qwen3.5/mmproj-BF16.gguf}}"
    if [ "${worker_mmproj}" = "none" ]; then worker_mmproj=""; fi
    local worker_tmpl="${LLM_WORKER_TEMPLATE:-${WORKER_TEMPLATE:-${SCRIPT_DIR}/templates/Qwen3.5/chat_template.jinja}}"

    # 1. Planner on 11434
    launch_instance "planner_27b" "$planner_model" "${LLM_PLANNER_PORT:-${PLANNER_PORT:-11434}}" "${LLM_PLANNER_CTX:-${PLANNER_CTX:-65536}}" \
        "$planner_mmproj" "$planner_tmpl" "${LLM_PLANNER_TEMP:-${PLANNER_TEMP:-0.6}}" "${LLM_PLANNER_TOP_P:-${PLANNER_TOP_P:-0.95}}" "${LLM_PLANNER_TOP_K:-${PLANNER_TOP_K:-40}}" \
        "${LLM_PLANNER_MIN_P:-${PLANNER_MIN_P:-0.05}}" "${LLM_PLANNER_NGL:-${PLANNER_NGL:-99}}" "${LLM_PLANNER_LOAD_MODE:-${PLANNER_LOAD_MODE:-none}}" "${LLM_FLASH_ATTN:-${FLASH_ATTN:-on}}" \
        "${LLM_PLANNER_CTK:-${PLANNER_CTK:-q8_0}}" "${LLM_PLANNER_CTV:-${PLANNER_CTV:-q8_0}}" "false"

    # 2. Worker 1 on 11435
    launch_instance "worker_4b_1" "$worker_model" "${LLM_WORKER1_PORT:-${WORKER1_PORT:-11435}}" "${LLM_WORKER1_CTX:-${WORKER1_CTX:-16384}}" \
        "$worker_mmproj" "$worker_tmpl" "${LLM_WORKER1_TEMP:-${WORKER1_TEMP:-0.7}}" "${LLM_WORKER1_TOP_P:-${WORKER1_TOP_P:-0.8}}" "${LLM_WORKER1_TOP_K:-${WORKER1_TOP_K:-20}}" \
        "${LLM_WORKER1_MIN_P:-${WORKER1_MIN_P:-0.05}}" "${LLM_WORKER1_NGL:-${WORKER1_NGL:-99}}" "${LLM_LOAD_MODE:-${LOAD_MODE:-mmap}}" "${LLM_FLASH_ATTN:-${FLASH_ATTN:-on}}" \
        "${LLM_WORKER1_CTK:-${WORKER1_CTK:-q8_0}}" "${LLM_WORKER1_CTV:-${WORKER1_CTV:-q8_0}}" "false"

    # 3. Worker 2 on 11436
    launch_instance "worker_4b_2" "$worker_model" "${LLM_WORKER2_PORT:-${WORKER2_PORT:-11436}}" "${LLM_WORKER2_CTX:-${WORKER2_CTX:-16384}}" \
        "$worker_mmproj" "$worker_tmpl" "${LLM_WORKER2_TEMP:-${WORKER2_TEMP:-0.7}}" "${LLM_WORKER2_TOP_P:-${WORKER2_TOP_P:-0.8}}" "${LLM_WORKER2_TOP_K:-${WORKER2_TOP_K:-20}}" \
        "${LLM_WORKER2_MIN_P:-${WORKER2_MIN_P:-0.05}}" "${LLM_WORKER2_NGL:-${WORKER2_NGL:-99}}" "${LLM_LOAD_MODE:-${LOAD_MODE:-mmap}}" "${LLM_FLASH_ATTN:-${FLASH_ATTN:-on}}" \
        "${LLM_WORKER2_CTK:-${WORKER2_CTK:-q8_0}}" "${LLM_WORKER2_CTV:-${WORKER2_CTV:-q8_0}}" "false"

    echo ""
    log_success "All 3 cluster servers are operational:"
    echo "  [Planner]     SuperQwen 3.8 27B -> http://localhost:${LLM_PLANNER_PORT:-${PLANNER_PORT:-11434}}"
    echo "  [Implementer] Qwen 3.5 4B #1    -> http://localhost:${LLM_WORKER1_PORT:-${WORKER1_PORT:-11435}}"
    echo "  [Implementer] Qwen 3.5 4B #2    -> http://localhost:${LLM_WORKER2_PORT:-${WORKER2_PORT:-11436}}"
    echo ""
    echo "Logs directory: ${LOG_DIR}/"
    echo "Press Ctrl+C to stop all instances gracefully."
    wait
}

cmd_muse() {
    log_info "Starting Muse-Glimmer 30B Solo Server (Max Context: 128K)..."
    pkill -f "llama-server" 2>/dev/null || true
    sleep 1
    local model="${LLM_MODEL:-${MODEL:-${SCRIPT_DIR}/models/muse-glimmer/Muse-Glimmer-30B-UD-Q4_K_XL.gguf}}"
    local mmproj="${LLM_MMPROJ:-${MMPROJ:-${SCRIPT_DIR}/models/muse-glimmer/mmproj-Muse-Glimmer-30B-Q8_0.gguf}}"
    if [ "${mmproj}" = "none" ]; then mmproj=""; fi
    local tmpl="${LLM_TEMPLATE:-${TEMPLATE:-${SCRIPT_DIR}/templates/muse-glimmer/chat_template.jinja}}"

    launch_instance "muse_solo" "$model" "${LLM_PORT:-${PORT:-11434}}" "${LLM_CTX:-${CTX:-131072}}" \
        "$mmproj" "$tmpl" "${LLM_TEMPERATURE:-${TEMPERATURE:-${SAMPLING_TEMP:-0.6}}}" \
        "${LLM_TOP_P:-${TOP_P:-0.95}}" "${LLM_TOP_K:-${TOP_K:-40}}" \
        "${LLM_MIN_P:-${MIN_P:-0.05}}" "${LLM_NGL:-${NGL:-99}}" \
        "${LLM_LOAD_MODE:-${LOAD_MODE:-none}}" "${LLM_FLASH_ATTN:-${FLASH_ATTN:-on}}" \
        "${LLM_CTK:-${CTK:-q8_0}}" "${LLM_CTV:-${CTV:-q8_0}}" \
        "${LLM_REASONING_PRESERVE:-${REASONING_PRESERVE:-true}}"

    echo ""
    echo "Streaming logs: tail -f ${LOG_DIR}/muse_solo.log"
    echo "Press Ctrl+C to stop."
    wait
}

cmd_superqwen() {
    log_info "Starting SuperQwen 3.8 27B Solo Server (Max Context: 256K)..."
    pkill -f "llama-server" 2>/dev/null || true
    sleep 1
    local model="${LLM_MODEL:-${MODEL:-${SCRIPT_DIR}/models/Qwen3.8/SuperQwen3.8-27b-abliterated-Q4_K_M.gguf}}"
    if [ ! -f "${model}" ] && [ -f "${SCRIPT_DIR}/models/Qwen3.8/Qwen3.8-27B-UD-Q4_K_XL.gguf" ]; then
        model="${SCRIPT_DIR}/models/Qwen3.8/Qwen3.8-27B-UD-Q4_K_XL.gguf"
    fi
    local mmproj="${LLM_MMPROJ:-${MMPROJ:-${SCRIPT_DIR}/models/Qwen3.8/mmproj-SuperQwen3.8-27b-abliterated-Q8_0.gguf}}"
    if [ "${mmproj}" = "none" ]; then mmproj=""; fi
    local tmpl="${LLM_TEMPLATE:-${TEMPLATE:-${SCRIPT_DIR}/templates/Qwen3.8/chat_template.jinja}}"

    launch_instance "superqwen_solo" "$model" "${LLM_PORT:-${PORT:-11434}}" "${LLM_CTX:-${CTX:-262144}}" \
        "$mmproj" "$tmpl" "${LLM_TEMPERATURE:-${TEMPERATURE:-${SAMPLING_TEMP:-0.6}}}" \
        "${LLM_TOP_P:-${TOP_P:-0.95}}" "${LLM_TOP_K:-${TOP_K:-40}}" \
        "${LLM_MIN_P:-${MIN_P:-0.05}}" "${LLM_NGL:-${NGL:-99}}" \
        "${LLM_LOAD_MODE:-${LOAD_MODE:-none}}" "${LLM_FLASH_ATTN:-${FLASH_ATTN:-on}}" \
        "${LLM_CTK:-${CTK:-q8_0}}" "${LLM_CTV:-${CTV:-q8_0}}" \
        "${LLM_REASONING_PRESERVE:-${REASONING_PRESERVE:-true}}"

    echo ""
    echo "Streaming logs: tail -f ${LOG_DIR}/superqwen_solo.log"
    echo "Press Ctrl+C to stop."
    wait
}

cmd_qwen4b() {
    log_info "Starting Qwen 3.5 4B Solo Server..."
    pkill -f "llama-server" 2>/dev/null || true
    sleep 1
    local model="${LLM_MODEL:-${MODEL:-${SCRIPT_DIR}/models/Qwen3.5/Qwen3.5-4B-Q5_K_M.gguf}}"
    local mmproj="${LLM_MMPROJ:-${MMPROJ:-${SCRIPT_DIR}/models/Qwen3.5/mmproj-BF16.gguf}}"
    if [ "${mmproj}" = "none" ]; then mmproj=""; fi
    local tmpl="${LLM_TEMPLATE:-${TEMPLATE:-${SCRIPT_DIR}/templates/Qwen3.5/chat_template.jinja}}"

    launch_instance "qwen4b_solo" "$model" "${LLM_PORT:-${PORT:-11434}}" "${LLM_CTX:-${CTX:-32768}}" \
        "$mmproj" "$tmpl" "${LLM_TEMPERATURE:-${TEMPERATURE:-${SAMPLING_TEMP:-0.7}}}" \
        "${LLM_TOP_P:-${TOP_P:-0.8}}" "${LLM_TOP_K:-${TOP_K:-20}}" \
        "${LLM_MIN_P:-${MIN_P:-0.05}}" "${LLM_NGL:-${NGL:-99}}" \
        "${LLM_LOAD_MODE:-${LOAD_MODE:-mmap}}" "${LLM_FLASH_ATTN:-${FLASH_ATTN:-on}}" \
        "${LLM_CTK:-${CTK:-q8_0}}" "${LLM_CTV:-${CTV:-q8_0}}" "false"

    echo ""
    echo "Streaming logs: tail -f ${LOG_DIR}/qwen4b_solo.log"
    echo "Press Ctrl+C to stop."
    wait
}

cmd_custom() {
    log_info "Starting Custom Server Instance..."
    pkill -f "llama-server" 2>/dev/null || true
    sleep 1
    local model="${LLM_MODEL:-${MODEL:-}}"
    if [ -z "${model}" ] || [ ! -f "${model}" ]; then
        log_error "No valid MODEL specified! Pass LLM_MODEL=/path/to/model.gguf"
        exit 1
    fi
    local mmproj="${LLM_MMPROJ:-${MMPROJ:-}}"
    if [ "${mmproj}" = "none" ]; then mmproj=""; fi
    local tmpl="${LLM_TEMPLATE:-${TEMPLATE:-}}"
    local name="custom_solo"

    launch_instance "${name}" "$model" "${LLM_PORT:-${PORT:-11434}}" "${LLM_CTX:-${CTX:-65536}}" \
        "$mmproj" "$tmpl" "${LLM_TEMPERATURE:-${TEMPERATURE:-${SAMPLING_TEMP:-0.6}}}" \
        "${LLM_TOP_P:-${TOP_P:-0.95}}" "${LLM_TOP_K:-${TOP_K:-40}}" \
        "${LLM_MIN_P:-${MIN_P:-0.05}}" "${LLM_NGL:-${NGL:-99}}" \
        "${LLM_LOAD_MODE:-${LOAD_MODE:-none}}" "${LLM_FLASH_ATTN:-${FLASH_ATTN:-on}}" \
        "${LLM_CTK:-${CTK:-q8_0}}" "${LLM_CTV:-${CTV:-q8_0}}" \
        "${LLM_REASONING_PRESERVE:-${REASONING_PRESERVE:-false}}"

    echo ""
    echo "Streaming logs: tail -f ${LOG_DIR}/${name}.log"
    echo "Press Ctrl+C to stop."
    wait
}

# ------------------------------------------------------------------------------
# Help / Usage
# ------------------------------------------------------------------------------
cmd_help() {
    cat <<EOF
Usage: $(basename "$0") <command> [options]

Commands:
  cluster       Launch multi-agent cluster (SuperQwen 27B on 11434, 2x Qwen 4B on 11435/11436)
  muse          Launch Muse-Glimmer 30B Solo (128K context, port 11434)
  superqwen     Launch SuperQwen 3.8 27B Solo (256K context, port 11434)
  qwen4b        Launch Qwen 3.5 4B Solo (port 11434)
  custom        Launch custom model instance (requires LLM_MODEL=/path/to/model.gguf)
  status        Show current server health and active ports
  stop [target] Stop servers (target: "all" or port number like 11434)
  help          Display this help message

Environment Variables:
  LLM_MODEL         Path to GGUF model file (fallback: MODEL)
  LLM_MMPROJ        Path to mmproj file, or "none" (fallback: MMPROJ)
  LLM_TEMPLATE      Path to chat template file (fallback: TEMPLATE)
  LLM_HOST          Bind address (default: 0.0.0.0)
  LLM_PORT          Main port override (default: 11434)
  LLM_CTX           Context size override (default: preset-specific)
  LLM_TEMPERATURE   Sampling temperature (default: preset-specific)
  LLM_PARALLEL      Parallel slot count (default: 1)
  LLM_FLASH_ATTN    Enable Flash Attention: "on" or "off" (default: on)
  LLM_LOG_DIR       Path to write logs (default: ./logs)

EOF
}

# ------------------------------------------------------------------------------
# Dispatch
# ------------------------------------------------------------------------------
ACTION="${1:-help}"
shift || true

case "${ACTION}" in
    cluster)    cmd_cluster "$@" ;;
    muse)       cmd_muse "$@" ;;
    superqwen)  cmd_superqwen "$@" ;;
    qwen4b)     cmd_qwen4b "$@" ;;
    custom)     cmd_custom "$@" ;;
    status)     cmd_status "$@" ;;
    stop)       cmd_stop "$@" ;;
    help|--help|-h) cmd_help ;;
    *)
        log_error "Unknown command: ${ACTION}"
        cmd_help
        exit 1
        ;;
esac