#!/bin/bash
set -e

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

# 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
    docker rm -f test-paddleocr-gpu 2>/dev/null || true
    docker rm -f test-paddleocr-cpu 2>/dev/null || true
}

trap cleanup EXIT

test_minicpm_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 ${MINICPM_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:${MINICPM_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 ""
}

test_paddleocr_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 ${PADDLEOCR_PORT}:5000 \
        ${extra_args} \
        ${REGISTRY}/${NAMESPACE}/${IMAGE_NAME}:${tag}

    # Wait for startup (PaddleOCR takes longer to initialize)
    echo "Waiting for container to start..."
    sleep 30

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

    # Test OCR endpoint with a minimal base64 image (1x1 white pixel PNG)
    echo "Testing OCR endpoint..."
    local test_image="iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
    local response=$(curl -s -X POST http://localhost:${PADDLEOCR_PORT}/ocr \
        -H "Content-Type: application/json" \
        -d "{\"image\": \"${test_image}\"}")

    if echo "$response" | grep -q '"success"'; then
        echo -e "${GREEN}OCR endpoint responding!${NC}"
    else
        echo -e "${RED}OCR endpoint not responding correctly!${NC}"
        echo "Response: $response"
        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 ""

echo -e "${BLUE}--- MiniCPM-V Tests ---${NC}"
echo ""

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

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

echo ""
echo -e "${BLUE}--- PaddleOCR Tests ---${NC}"
echo ""

# Test PaddleOCR CPU variant (doesn't require GPU)
test_paddleocr_image "paddleocr-cpu" "test-paddleocr-cpu" ""

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

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