start.sh

#!/usr/bin/env bash
# =============================================================================
# start.sh — Model selector & dashboard launcher for llama.cpp Docker wrapper
#
# Usage:
#   ./start.sh              # Interactive menu + monitoring dashboard
#   ./start.sh glm-flash-q4  # Direct launch by section ID
#   ./start.sh --list       # List available models
#   ./start.sh --no-dashboard  # Launch without dashboard (raw logs)
# =============================================================================

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
CONF="$SCRIPT_DIR/models.conf"
MODELS_DIR="$SCRIPT_DIR/models"
ENV_FILE="$SCRIPT_DIR/.env"
COMPOSE_FILE="$SCRIPT_DIR/docker-compose.yml"
DASHBOARD="$SCRIPT_DIR/dashboard.py"

# --- INI parser -----------------------------------------------------------
# Reads models.conf into parallel arrays. No eval, no external tools.

declare -a SECTION_IDS=()
declare -a SECTION_NAMES=()
declare -A CONFIG=()  # CONFIG[section.KEY]=value

parse_conf() {
    local section=""
    while IFS= read -r line || [[ -n "$line" ]]; do
        # Strip leading/trailing whitespace
        line="${line#"${line%%[![:space:]]*}"}"
        line="${line%"${line##*[![:space:]]}"}"

        # Skip empty lines and comments
        [[ -z "$line" || "$line" == \#* ]] && continue

        # Section header
        if [[ "$line" =~ ^\[([a-zA-Z0-9_-]+)\]$ ]]; then
            section="${BASH_REMATCH[1]}"
            SECTION_IDS+=("$section")
            continue
        fi

        # Key=Value
        if [[ -n "$section" && "$line" == *=* ]]; then
            local key="${line%%=*}"
            local value="${line#*=}"
            CONFIG["${section}.${key}"]="$value"
            if [[ "$key" == "NAME" ]]; then
                SECTION_NAMES+=("$value")
            fi
        fi
    done < "$CONF"
}

# --- Helpers ---------------------------------------------------------------

get() { echo "${CONFIG["${1}.${2}"]:-}"; }

die() { echo "Error: $*" >&2; exit 1; }

ctx_label() {
    local ctx="$1"
    if (( ctx >= 262144 )); then echo "256K ctx"
    elif (( ctx >= 131072 )); then echo "128K ctx"
    elif (( ctx >= 65536 )); then echo "64K ctx"
    else echo "${ctx} ctx"
    fi
}

list_models() {
    local prod_ids=() bench_ids=()
    for id in "${SECTION_IDS[@]}"; do
        if [[ "$id" == bench-* ]]; then
            bench_ids+=("$id")
        else
            prod_ids+=("$id")
        fi
    done

    echo "Production models:"
    echo ""
    for id in "${prod_ids[@]}"; do
        local name; name=$(get "$id" NAME)
        local ctx; ctx=$(get "$id" CTX_SIZE)
        local speed; speed=$(get "$id" SPEED)
        local desc; desc=$(get "$id" DESCRIPTION)
        printf "  %-20s  %s" "$id" "$name"
        [[ -n "$speed" ]] && printf "  %s" "$speed"
        [[ -n "$ctx" ]] && printf "  (%s)" "$(ctx_label "$ctx")"
        printf "\n"
        [[ -n "$desc" ]] && printf "  %-20s  %s\n" "" "$desc"
    done

    if [[ ${#bench_ids[@]} -gt 0 ]]; then
        echo ""
        echo "Benchmark profiles:"
        echo ""
        for id in "${bench_ids[@]}"; do
            local name; name=$(get "$id" NAME)
            local ctx; ctx=$(get "$id" CTX_SIZE)
            printf "  %-20s  %s" "$id" "$name"
            [[ -n "$ctx" ]] && printf "  (%s)" "$(ctx_label "$ctx")"
            printf "\n"
        done
    fi
}

# --- Prerequisite checks ---------------------------------------------------

check_docker() {
    if ! command -v docker &>/dev/null; then
        die "docker not found. Install Docker first."
    fi
    if ! docker info &>/dev/null 2>&1; then
        die "Docker daemon is not running."
    fi
}

check_existing_container() {
    local state
    state=$(docker inspect --format='{{.State.Status}}' llama-server 2>/dev/null || true)
    if [[ "$state" == "running" ]]; then
        echo "Container 'llama-server' is already running."
        read -rp "Stop it and continue? [y/N] " answer
        if [[ "${answer,,}" == "y" ]]; then
            docker compose -f "$COMPOSE_FILE" down
        else
            echo "Aborted."
            exit 0
        fi
    fi
}

# --- .env generation -------------------------------------------------------

generate_env() {
    local id="$1"
    local model; model=$(get "$id" MODEL)
    local ctx; ctx=$(get "$id" CTX_SIZE)
    local ngl; ngl=$(get "$id" N_GPU_LAYERS)
    local fit; fit=$(get "$id" FIT)
    local fit_target; fit_target=$(get "$id" FIT_TARGET)
    local extra; extra=$(get "$id" EXTRA_ARGS)

    {
        echo "# Generated by start.sh — $(get "$id" NAME)"
        echo "# Section: [$id] from models.conf"
        echo ""
        [[ -n "$model" ]]      && echo "MODEL=$model"
        [[ -n "$ctx" ]]        && echo "CTX_SIZE=$ctx"
        [[ -n "$ngl" ]]        && echo "N_GPU_LAYERS=$ngl"
        [[ -n "$fit" ]]        && echo "FIT=$fit"
        [[ -n "$fit_target" ]] && echo "FIT_TARGET=$fit_target"
        [[ -n "$extra" ]]      && echo "EXTRA_ARGS=$extra"
    } > "$ENV_FILE"
}

# --- Model file check ------------------------------------------------------

check_model_file() {
    local id="$1"
    local model; model=$(get "$id" MODEL)
    local full_path="$MODELS_DIR/$model"
    if [[ ! -f "$full_path" ]]; then
        die "Model file not found: $full_path
Download it first, then try again."
    fi
}

# --- Menu -------------------------------------------------------------------

show_menu() {
    # Split into production and benchmark profiles
    local prod_ids=() bench_ids=()
    for id in "${SECTION_IDS[@]}"; do
        if [[ "$id" == bench-* ]]; then
            bench_ids+=("$id")
        else
            prod_ids+=("$id")
        fi
    done

    echo "" >&2
    echo "llama.cpp Model Selector" >&2
    echo "========================" >&2
    echo "" >&2

    for i in "${!prod_ids[@]}"; do
        local id="${prod_ids[$i]}"
        local name; name=$(get "$id" NAME)
        local ctx; ctx=$(get "$id" CTX_SIZE)
        local speed; speed=$(get "$id" SPEED)
        local desc; desc=$(get "$id" DESCRIPTION)
        local label=""
        [[ -n "$ctx" ]] && label=$(ctx_label "$ctx")

        # Right-align speed + context
        local right=""
        [[ -n "$speed" ]] && right="${speed}  ${label}" || right="$label"
        printf "  %d) %-44s %s\n" "$((i + 1))" "$name" "$right" >&2

        # Description on next line (indented)
        [[ -n "$desc" ]] && printf "     %-44s\n" "$desc" >&2
    done

    echo "" >&2
    [[ ${#bench_ids[@]} -gt 0 ]] && echo "  b) Benchmarks >" >&2
    echo "  q) Quit" >&2
    echo "" >&2

    local count=${#prod_ids[@]}
    while true; do
        read -rp "Select model [1-${count}, b, q]: " choice
        if [[ "$choice" == "q" || "$choice" == "Q" ]]; then
            echo "Aborted." >&2
            exit 0
        fi
        if [[ "$choice" == "b" || "$choice" == "B" ]] && [[ ${#bench_ids[@]} -gt 0 ]]; then
            local bench_result
            bench_result=$(show_bench_menu "${bench_ids[@]}")
            if [[ "$bench_result" == "__back__" ]]; then
                # Return to main menu — re-display
                show_menu
                return
            fi
            echo "$bench_result"
            return
        fi
        if [[ "$choice" =~ ^[0-9]+$ ]] && (( choice >= 1 && choice <= count )); then
            echo "${prod_ids[$((choice - 1))]}"
            return
        fi
        echo "Invalid choice. Enter 1-${count}, b, or q to quit." >&2
    done
}

show_bench_menu() {
    local bench_ids=("$@")

    echo "" >&2
    echo "Benchmark Profiles (10K context, optimized GPU layers)" >&2
    echo "======================================================" >&2
    echo "" >&2

    for i in "${!bench_ids[@]}"; do
        local id="${bench_ids[$i]}"
        local name; name=$(get "$id" NAME)
        local ctx; ctx=$(get "$id" CTX_SIZE)
        local label=""
        [[ -n "$ctx" ]] && label=$(ctx_label "$ctx")
        printf "  %d) %-48s %s\n" "$((i + 1))" "$name" "$label" >&2
    done

    echo "" >&2
    echo "  r) Return to main menu" >&2
    echo "  q) Quit" >&2
    echo "" >&2

    local count=${#bench_ids[@]}
    while true; do
        read -rp "Select benchmark [1-${count}, r, q]: " choice
        if [[ "$choice" == "q" || "$choice" == "Q" ]]; then
            echo "Aborted." >&2
            exit 0
        fi
        if [[ "$choice" == "r" || "$choice" == "R" ]]; then
            echo "__back__"
            return
        fi
        if [[ "$choice" =~ ^[0-9]+$ ]] && (( choice >= 1 && choice <= count )); then
            echo "${bench_ids[$((choice - 1))]}"
            return
        fi
        echo "Invalid choice. Enter 1-${count}, r, or q to quit." >&2
    done
}

# --- Health check polling ---------------------------------------------------

wait_for_health() {
    local url="http://localhost:8080/health"
    local timeout=300  # 5 minutes
    local elapsed=0
    local spin=('|' '/' '-' '\')
    local i=0

    echo -n "Waiting for server to be ready "

    while (( elapsed < timeout )); do
        # Check if container is still running
        local state
        state=$(docker inspect --format='{{.State.Status}}' llama-server 2>/dev/null || echo "missing")
        if [[ "$state" != "running" ]]; then
            echo ""
            echo ""
            echo "Container stopped unexpectedly. Showing recent logs:"
            echo ""
            docker compose -f "$COMPOSE_FILE" logs --tail=30 2>/dev/null || true
            return 1
        fi

        if curl -sf "$url" &>/dev/null; then
            echo -e " \033[1;32mready!\033[0m"
            return 0
        fi

        printf "\rWaiting for server to be ready %s " "${spin[$((i % 4))]}"
        sleep 2
        elapsed=$((elapsed + 2))
        i=$((i + 1))
    done

    echo ""
    echo "Timeout waiting for server health endpoint after ${timeout}s."
    echo "The server may still be loading. Showing recent logs:"
    echo ""
    docker compose -f "$COMPOSE_FILE" logs --tail=20 2>/dev/null || true
    return 1
}

# --- Cleanup ----------------------------------------------------------------

cleanup() {
    docker compose -f "$COMPOSE_FILE" down 2>/dev/null || true
}

# --- Dashboard launch -------------------------------------------------------

run_with_dashboard() {
    local selected="$1"
    local model_name
    model_name=$(get "$selected" NAME)

    echo "Starting container in background..."
    docker compose -f "$COMPOSE_FILE" up -d

    echo ""
    if ! wait_for_health; then
        echo ""
        read -rp "Server may not be ready. Launch dashboard anyway? [Y/n] " answer
        if [[ "${answer,,}" == "n" ]]; then
            docker compose -f "$COMPOSE_FILE" down
            return 1
        fi
    fi

    echo "Launching dashboard..."
    echo ""

    # Set cleanup trap — ensures docker compose down on unexpected exit
    trap cleanup EXIT HUP TERM INT

    # Run dashboard (blocks until user presses q or r)
    local rc=0
    python3 "$DASHBOARD" \
        --compose-file "$COMPOSE_FILE" \
        --model-name "$model_name" \
        --models-conf "$CONF" \
        --current-profile "$selected" \
    || rc=$?

    # Remove trap — we handle cleanup ourselves from here
    trap - EXIT HUP TERM INT

    case "$rc" in
        0)
            # q pressed: stop & exit
            echo ""
            echo "Stopping container..."
            docker compose -f "$COMPOSE_FILE" down
            echo "Done."
            return 0
            ;;
        2)
            # r pressed: stop & return to menu
            echo ""
            echo "Stopping container..."
            docker compose -f "$COMPOSE_FILE" down
            echo ""
            return 2
            ;;
        *)
            # Unexpected exit
            echo ""
            echo "Dashboard exited with code $rc. Stopping container..."
            docker compose -f "$COMPOSE_FILE" down
            return 0
            ;;
    esac
}

# =============================================================================
# Main
# =============================================================================

[[ ! -f "$CONF" ]] && die "Config file not found: $CONF"

parse_conf

if [[ ${#SECTION_IDS[@]} -eq 0 ]]; then
    die "No models defined in $CONF"
fi

# Handle arguments
selected=""
no_dashboard=false

while [[ $# -gt 0 ]]; do
    case "$1" in
        --list|-l)
            list_models
            exit 0
            ;;
        --help|-h)
            echo "Usage: $0 [model-id | --list | --help | --no-dashboard]"
            echo ""
            echo "Options:"
            echo "  --list, -l         List available models"
            echo "  --no-dashboard     Skip dashboard, show raw docker compose logs"
            echo "  --help, -h         Show this help"
            echo ""
            list_models
            exit 0
            ;;
        --no-dashboard)
            no_dashboard=true
            shift
            ;;
        -*)
            die "Unknown option: $1"
            ;;
        *)
            # Direct model ID
            found=false
            for id in "${SECTION_IDS[@]}"; do
                if [[ "$id" == "$1" ]]; then
                    selected="$1"
                    found=true
                    break
                fi
            done
            if [[ "$found" == false ]]; then
                echo "Error: Unknown model '$1'" >&2
                echo ""
                list_models >&2
                exit 1
            fi
            shift
            ;;
    esac
done

check_docker

# --- Main loop (supports return-to-menu via 'r' key) ---

while true; do
    # Interactive menu if no model specified
    if [[ -z "$selected" ]]; then
        selected=$(show_menu)
    fi

    echo ""
    echo "Selected: $(get "$selected" NAME)"

    check_model_file "$selected"
    check_existing_container
    generate_env "$selected"

    echo "Generated .env for [$selected]"

    # Decide launch mode
    if [[ "$no_dashboard" == true ]] || [[ ! -f "$DASHBOARD" ]]; then
        echo "Starting docker compose..."
        echo ""
        exec docker compose -f "$COMPOSE_FILE" up
    fi

    # Launch with dashboard
    rc=0
    run_with_dashboard "$selected" || rc=$?

    if [[ $rc -eq 2 ]]; then
        # Return to menu
        selected=""
        continue
    fi

    # Normal exit
    break
done