#!/bin/bash
set -e

# Configuration
REGISTRY="code.foss.global"
NAMESPACE="host.today"
IMAGE_NAME="ht-docker-ai"
TEST_PORT=11434

# Colors for output
GREEN='\033[0;32m'
RED='\033[0;31m'
BLUE='\033[0;34m'
NC='\033[0m'

cleanup() {
    echo -e "${BLUE}Cleaning up test containers...${NC}"
    docker rm -f test-minicpm-gpu 2>/dev/null || true
    docker rm -f test-minicpm-cpu 2>/dev/null || true
}

trap cleanup EXIT

test_image() {
    local tag=$1
    local container_name=$2
    local extra_args=$3

    echo -e "${BLUE}Testing ${tag}...${NC}"

    # Start container
    docker run -d \
        --name ${container_name} \
        -p ${TEST_PORT}:11434 \
        ${extra_args} \
        ${REGISTRY}/${NAMESPACE}/${IMAGE_NAME}:${tag}

    # Wait for startup
    echo "Waiting for container to start..."
    sleep 10

    # Test API endpoint
    echo "Testing API endpoint..."
    if curl -s -f http://localhost:${TEST_PORT}/api/tags > /dev/null; then
        echo -e "${GREEN}API endpoint responding!${NC}"
    else
        echo -e "${RED}API endpoint not responding!${NC}"
        docker logs ${container_name}
        return 1
    fi

    # Cleanup this container
    docker rm -f ${container_name}

    echo -e "${GREEN}${tag} test passed!${NC}"
    echo ""
}

echo -e "${BLUE}=== Testing ht-docker-ai images ===${NC}"
echo ""

# Test CPU variant (doesn't require GPU)
test_image "minicpm45v-cpu" "test-minicpm-cpu" ""

# Test GPU variant only if NVIDIA runtime is available
if docker info 2>/dev/null | grep -q "nvidia"; then
    test_image "minicpm45v" "test-minicpm-gpu" "--gpus all"
else
    echo -e "${BLUE}Skipping GPU test (NVIDIA runtime not available)${NC}"
fi

echo -e "${GREEN}=== All tests passed! ===${NC}"