#!/usr/bin/env bash
# Benchmark setup — ensure toolboxes and test models are ready
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "$SCRIPT_DIR/../../lib/common.sh"
source "$SCRIPT_DIR/../../lib/detect.sh"

TOOLBOXES_REPO="/data/workspace/projects/HomeLab/strix-halo-toolboxes/amd-strix-halo-llamacpp-toolboxes"
MODEL_DIR="$(data_dir models)"

log_header "Benchmark Setup"

# ── 1. Check toolbox containers ──────────────────────────
log_info "Checking toolbox containers..."

# Minimum required: vulkan-radv (most stable)
REQUIRED_TOOLBOXES=("llama-vulkan-radv")
OPTIONAL_TOOLBOXES=("llama-rocm-6.4.4" "llama-rocm-7.2" "llama-vulkan-amdvlk")

existing=$(detect_toolbox_names 2>/dev/null || true)
missing=()

for tb in "${REQUIRED_TOOLBOXES[@]}"; do
    if echo "$existing" | grep -q "^${tb}$"; then
        log_success "Toolbox: $tb"
    else
        missing+=("$tb")
        log_warn "Toolbox missing: $tb"
    fi
done

for tb in "${OPTIONAL_TOOLBOXES[@]}"; do
    if echo "$existing" | grep -q "^${tb}$"; then
        log_success "Toolbox: $tb (optional)"
    else
        log_info "Toolbox not present: $tb (optional)"
    fi
done

if (( ${#missing[@]} > 0 )); then
    log_info "Need to create required toolboxes."
    if [[ -d "$TOOLBOXES_REPO" ]]; then
        log_info "Found toolboxes repo at: $TOOLBOXES_REPO"
        if confirm "Create missing toolboxes using refresh-toolboxes.sh?"; then
            for tb in "${missing[@]}"; do
                log_info "Creating $tb..."
                bash "$TOOLBOXES_REPO/refresh-toolboxes.sh" "$tb"
            done
        fi
    else
        log_error "Toolboxes repo not found at: $TOOLBOXES_REPO"
        log_info "Clone it: git clone https://github.com/kyuz0/amd-strix-halo-toolboxes"
        log_info "Then re-run this setup."
        exit 1
    fi
fi

# ── 2. Verify GPU access inside toolboxes ────────────────
log_info "Verifying GPU access in toolboxes..."
for tb in "${REQUIRED_TOOLBOXES[@]}"; do
    if echo "$existing" | grep -qF "$tb"; then
        if toolbox run -c "$tb" -- llama-cli --list-devices 2>&1 | grep -qi "gpu\|vulkan\|rocm"; then
            log_success "GPU accessible in $tb"
        else
            log_warn "GPU may not be accessible in $tb — check device mappings"
        fi
    fi
done

# ── 3. Check for test models ────────────────────────────
log_info "Checking for test models in $MODEL_DIR..."

model_count=$(find "$MODEL_DIR" -name "*.gguf" 2>/dev/null | wc -l)
if (( model_count > 0 )); then
    log_success "Found $model_count model(s):"
    find "$MODEL_DIR" -name "*.gguf" | while read -r f; do
        size=$(du -h "$f" | cut -f1)
        printf "  %s (%s)\n" "$(basename "$f")" "$size"
    done
else
    log_warn "No GGUF models found in $MODEL_DIR"
    log_info "Download a test model. Example:"
    echo ""
    echo "  # Small (4B, ~3 GB):"
    echo "  huggingface-cli download Qwen/Qwen3-4B-GGUF Qwen3-4B-Q4_K_M.gguf \\"
    echo "    --local-dir $MODEL_DIR"
    echo ""
    echo "  # Medium (14B, ~9 GB):"
    echo "  huggingface-cli download Qwen/Qwen3-14B-GGUF Qwen3-14B-Q4_K_M.gguf \\"
    echo "    --local-dir $MODEL_DIR"
    echo ""

    if is_cmd huggingface-cli; then
        if confirm "Download Qwen3-4B Q4_K_M (~3 GB) as test model?"; then
            huggingface-cli download Qwen/Qwen3-4B-GGUF Qwen3-4B-Q4_K_M.gguf \
                --local-dir "$MODEL_DIR"
            log_success "Model downloaded"
        fi
    else
        log_info "Install huggingface-cli: pip install huggingface_hub[cli]"
    fi
fi

log_header "Setup Complete"
log_info "Run 'make benchmark-baseline' to capture your baseline."