#!/bin/bash
# Build MCP-Forge container images
#
# Usage: ./build-images.sh [--all|--jupyter]

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"

cd "$PROJECT_ROOT"

build_jupyter() {
    echo "Building mcp-forge/jupyter:latest..."
    podman build -f docker/Dockerfile.jupyter -t mcp-forge/jupyter:latest .
    echo "✓ Built mcp-forge/jupyter:latest"
}

ensure_base_images() {
    echo "Ensuring base Python images exist..."
    
    # Check if images exist
    if ! podman image exists mcp-forge/python:3.11; then
        echo "  Pulling and tagging python:3.11-slim..."
        podman pull python:3.11-slim
        podman tag python:3.11-slim mcp-forge/python:3.11
        echo "  ✓ Created mcp-forge/python:3.11"
    else
        echo "  ✓ mcp-forge/python:3.11 exists"
    fi
    
    if ! podman image exists mcp-forge/python:3.12; then
        echo "  Pulling and tagging python:3.12-slim..."
        podman pull python:3.12-slim
        podman tag python:3.12-slim mcp-forge/python:3.12
        echo "  ✓ Created mcp-forge/python:3.12"
    else
        echo "  ✓ mcp-forge/python:3.12 exists"
    fi
}

verify_images() {
    echo ""
    echo "Verifying images..."
    echo ""
    podman images | grep -E "REPOSITORY|mcp-forge"
}

test_jupyter_image() {
    echo ""
    echo "Testing Jupyter image..."
    
    echo -n "  Testing ipykernel import... "
    VERSION=$(podman run --rm mcp-forge/jupyter:latest python -c "import ipykernel; print(ipykernel.__version__)")
    echo "✓ ipykernel $VERSION"
    
    echo -n "  Testing kernel launcher... "
    podman run --rm mcp-forge/jupyter:latest python -m ipykernel_launcher --help > /dev/null 2>&1
    echo "✓ launcher works"
}

# Parse arguments
ACTION="${1:-jupyter}"

case "$ACTION" in
    --all)
        ensure_base_images
        build_jupyter
        verify_images
        test_jupyter_image
        ;;
    --jupyter)
        build_jupyter
        test_jupyter_image
        ;;
    *)
        echo "Usage: $0 [--all|--jupyter]"
        echo ""
        echo "  --all      Ensure base images and build Jupyter image"
        echo "  --jupyter  Build Jupyter image only (default)"
        exit 1
        ;;
esac

echo ""
echo "✓ Done!"