#!/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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
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"\|"status":\s*"loading model"\|"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
        local pid
        pid=$(lsof -ti :"$target" 2>/dev/null || fuser "$target/tcp" 2>/dev/null || true)
        if [ -n "$pid" ]; then
            kill "$pid" 2>/dev/null || true
            log_success "Killed server on port ${target} (PID: ${pid})."
        else
            log_warn "No active server found on port ${target}."
        fi
    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

    # Check if port is already in use
    local existing_pid
    existing_pid=$(lsof -ti :"$port" 2>/dev/null || fuser "$port/tcp" 2>/dev/null || true)
    if [ -n "$existing_pid" ]; then
        log_warn "Port ${port} is already in use by PID ${existing_pid}. Terminating existing process..."
        kill -9 "$existing_pid" 2>/dev/null || true
        sleep 1
    fi

    local flags=(
        -m "${model}"
        --host "${HOST:-0.0.0.0}"
        --port "${port}"
        --ctx-size "${ctx}"
        --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 "${BATCH_SIZE:-2048}"
        -ub "${UBATCH_SIZE:-512}"
    )

    if [ -n "${mmproj}" ] && [ -f "${mmproj}" ]; then
        flags+=(--mmproj "${mmproj}" --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)..."

    # 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="${PLANNER_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="${PLANNER_MMPROJ:-${SCRIPT_DIR}/models/Qwen3.8/mmproj-SuperQwen3.8-27b-abliterated-Q8_0.gguf}"
    local planner_tmpl="${PLANNER_TEMPLATE:-${SCRIPT_DIR}/templates/Qwen3.8/chat_template.jinja}"

    # Qwen 3.5 4B (Workers)
    local worker_model="${WORKER_MODEL:-${SCRIPT_DIR}/models/Qwen3.5/Qwen3.5-4B-Q5_K_M.gguf}"
    local worker_mmproj="${WORKER_MMPROJ:-${SCRIPT_DIR}/models/Qwen3.5/mmproj-BF16.gguf}"
    local worker_tmpl="${WORKER_TEMPLATE:-${SCRIPT_DIR}/templates/Qwen3.5/chat_template.jinja}"

    # 1. Planner on 11434
    launch_instance "planner_27b" "$planner_model" "${PLANNER_PORT:-11434}" "${PLANNER_CTX:-65536}" \
        "$planner_mmproj" "$planner_tmpl" "${PLANNER_TEMP:-0.6}" "${PLANNER_TOP_P:-0.95}" "${PLANNER_TOP_K:-40}" \
        "${PLANNER_MIN_P:-0.05}" "${PLANNER_NGL:-99}" "${PLANNER_LOAD_MODE:-none}" "${FLASH_ATTN:-on}" \
        "${PLANNER_CTK:-q8_0}" "${PLANNER_CTV:-q8_0}" "false"

    # 2. Worker 1 on 11435
    launch_instance "worker_4b_1" "$worker_model" "${WORKER1_PORT:-11435}" "${WORKER1_CTX:-16384}" \
        "$worker_mmproj" "$worker_tmpl" "${WORKER1_TEMP:-0.7}" "${WORKER1_TOP_P:-0.8}" "${WORKER1_TOP_K:-20}" \
        "${WORKER1_MIN_P:-0.05}" "${WORKER1_NGL:-99}" "${LOAD_MODE:-mmap}" "${FLASH_ATTN:-on}" \
        "${WORKER1_CTK:-q8_0}" "${WORKER1_CTV:-q8_0}" "false"

    # 3. Worker 2 on 11436
    launch_instance "worker_4b_2" "$worker_model" "${WORKER2_PORT:-11436}" "${WORKER2_CTX:-16384}" \
        "$worker_mmproj" "$worker_tmpl" "${WORKER2_TEMP:-0.7}" "${WORKER2_TOP_P:-0.8}" "${WORKER2_TOP_K:-20}" \
        "${WORKER2_MIN_P:-0.05}" "${WORKER2_NGL:-99}" "${LOAD_MODE:-mmap}" "${FLASH_ATTN:-on}" \
        "${WORKER2_CTK:-q8_0}" "${WORKER2_CTV:-q8_0}" "false"

    echo ""
    log_success "All 3 cluster servers are operational:"
    echo "  [Planner]     SuperQwen 3.8 27B -> http://localhost:${PLANNER_PORT:-11434}"
    echo "  [Implementer] Qwen 3.5 4B #1    -> http://localhost:${WORKER1_PORT:-11435}"
    echo "  [Implementer] Qwen 3.5 4B #2    -> http://localhost:${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)..."
    local model="${MODEL:-${SCRIPT_DIR}/models/muse-glimmer/Muse-Glimmer-30B-UD-Q4_K_XL.gguf}"
    local mmproj="${MMPROJ:-${SCRIPT_DIR}/models/muse-glimmer/mmproj-Muse-Glimmer-30B-Q8_0.gguf}"
    local tmpl="${TEMPLATE:-${SCRIPT_DIR}/templates/muse-glimmer/chat_template.jinja}"

    launch_instance "muse_solo" "$model" "${PORT:-11434}" "${CTX:-131072}" \
        "$mmproj" "$tmpl" "${TEMP:-0.6}" "${TOP_P:-0.95}" "${TOP_K:-40}" \
        "${MIN_P:-0.05}" "${NGL:-99}" "${LOAD_MODE:-none}" "${FLASH_ATTN:-on}" \
        "${CTK:-q8_0}" "${CTV:-q8_0}" "${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)..."
    local 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="${MMPROJ:-${SCRIPT_DIR}/models/Qwen3.8/mmproj-SuperQwen3.8-27b-abliterated-Q8_0.gguf}"
    local tmpl="${TEMPLATE:-${SCRIPT_DIR}/templates/Qwen3.8/chat_template.jinja}"

    launch_instance "superqwen_solo" "$model" "${PORT:-11434}" "${CTX:-262144}" \
        "$mmproj" "$tmpl" "${TEMP:-0.6}" "${TOP_P:-0.95}" "${TOP_K:-40}" \
        "${MIN_P:-0.05}" "${NGL:-99}" "${LOAD_MODE:-none}" "${FLASH_ATTN:-on}" \
        "${CTK:-q8_0}" "${CTV:-q8_0}" "${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..."
    local model="${MODEL:-${SCRIPT_DIR}/models/Qwen3.5/Qwen3.5-4B-Q5_K_M.gguf}"
    local mmproj="${MMPROJ:-${SCRIPT_DIR}/models/Qwen3.5/mmproj-BF16.gguf}"
    local tmpl="${TEMPLATE:-${SCRIPT_DIR}/templates/Qwen3.5/chat_template.jinja}"

    launch_instance "qwen4b_solo" "$model" "${PORT:-11434}" "${CTX:-32768}" \
        "$mmproj" "$tmpl" "${TEMP:-0.7}" "${TOP_P:-0.8}" "${TOP_K:-20}" \
        "${MIN_P:-0.05}" "${NGL:-99}" "${LOAD_MODE:-mmap}" "${FLASH_ATTN:-on}" \
        "${CTK:-q8_0}" "${CTV:-q8_0}" "false"

    echo ""
    echo "Streaming logs: tail -f ${LOG_DIR}/qwen4b_solo.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)
  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:
  HOST          Bind address (default: 0.0.0.0)
  PORT          Main port override (default: 11434)
  CTX           Context size override
  FLASH_ATTN    Enable Flash Attention: "on" or "off" (default: on)
  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 "$@" ;;
    status)     cmd_status "$@" ;;
    stop)       cmd_stop "$@" ;;
    help|--help|-h) cmd_help ;;
    *)
        log_error "Unknown command: ${ACTION}"
        cmd_help
        exit 1
        ;;
esac