#!/bin/bash
#
# ATCR End-to-End Test Script
# Tests single-arch and multi-arch image push/pull flows
#
# Usage:
#   ./test-e2e.sh                    # Run full test suite
#   ./test-e2e.sh --multi            # Skip to multi-arch test only
#   REGISTRY=localhost:5000 ./test-e2e.sh       # Custom registry
#   NAMESPACE=myuser ./test-e2e.sh              # Custom namespace
#   VERBOSE=0 ./test-e2e.sh                     # Hide docker output
#
# To see bash command execution, edit this file and uncomment 'set -x' below
#

set -e

# Verbose mode - shows docker command output
# Set VERBOSE=0 to hide docker output: VERBOSE=0 ./test-e2e.sh
VERBOSE="${VERBOSE:-1}"

# Bash trace mode - shows every bash command as it executes
# Uncomment the line below to see bash commands as they execute
# set -x

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Configuration
REGISTRY="${REGISTRY:-127.0.0.1:5000}"
NAMESPACE="${NAMESPACE:-evan.jarrett.net}"
IMAGE_NAME="test-image"
TAG="latest"
MULTI_ARCH_TAG="multiarch"

# Full image references
SINGLE_ARCH_IMAGE="${REGISTRY}/${NAMESPACE}/${IMAGE_NAME}:${TAG}"
MULTI_ARCH_IMAGE="${REGISTRY}/${NAMESPACE}/${IMAGE_NAME}:${MULTI_ARCH_TAG}"
AMD64_IMAGE="${REGISTRY}/${NAMESPACE}/${IMAGE_NAME}:${TAG}-amd64"
ARM64_IMAGE="${REGISTRY}/${NAMESPACE}/${IMAGE_NAME}:${TAG}-arm64"

# Temporary directory for test images
BUILD_DIR=$(mktemp -d)

# Cleanup function
cleanup() {
    log_info "Cleaning up..."
    rm -rf ${BUILD_DIR}

    # Clean up any dangling manifest lists
    docker manifest rm ${MULTI_ARCH_IMAGE} 2>/dev/null || true
}

trap cleanup EXIT

# Logging functions
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_step() {
    echo -e "\n${GREEN}==>${NC} $1"
}

log_cmd() {
    echo -e "${BLUE}[CMD]${NC} $*"
    if [ "$VERBOSE" = "1" ]; then
        "$@"
    else
        "$@" > /dev/null 2>&1
    fi
}

# Check if docker compose services are running
check_compose_services() {
    log_step "Checking docker compose services..."

    if ! docker compose ps | grep -q "Up"; then
        log_warn "Some services may not be running. Starting services..."
        docker compose up -d
        sleep 5
    fi

    # Check for errors in logs
    log_info "Checking for errors in service logs..."
    if docker compose logs --tail=50 | grep -i "error" | grep -v "level=error msg=\"error processing" | grep -v "test"; then
        log_warn "Found some errors in logs (may be normal)"
    else
        log_info "No critical errors found in recent logs"
    fi
}

# Build a simple test image
build_single_arch_image() {
    log_step "Building single-arch test image..."

    cat > ${BUILD_DIR}/Dockerfile <<EOF
FROM alpine:latest

# Add some content to make the image non-trivial
RUN apk add --no-cache curl bash

# Create a test file with some content
RUN echo "This is a test image created at \$(date)" > /test.txt
RUN echo "Architecture: \$(uname -m)" >> /test.txt

# Add a simple script
RUN echo '#!/bin/sh' > /test.sh && \\
    echo 'echo "Test image running successfully!"' >> /test.sh && \\
    echo 'cat /test.txt' >> /test.sh && \\
    chmod +x /test.sh

CMD ["/test.sh"]
EOF

    log_cmd docker build -t ${SINGLE_ARCH_IMAGE} ${BUILD_DIR}
    log_info "Built single-arch image: ${SINGLE_ARCH_IMAGE}"
}

# Build multi-arch images using docker manifest create (old school!)
build_multi_arch_images() {
    log_step "Building multi-arch images (amd64 and arm64) using docker manifest..."

    # Create Dockerfile for multi-arch builds
    cat > ${BUILD_DIR}/Dockerfile.multiarch <<EOF
FROM alpine:latest

ARG TARGETARCH=unknown
ARG TARGETOS=linux

# Add some content to make the image non-trivial
RUN apk add --no-cache curl bash

# Create a test file with arch info
RUN echo "This is a multi-arch test image" > /test.txt
RUN echo "Target OS: \${TARGETOS}" >> /test.txt
RUN echo "Target Architecture: \${TARGETARCH}" >> /test.txt
RUN echo "Built at: \$(date)" >> /test.txt

# Add a simple script
RUN echo '#!/bin/sh' > /test.sh && \\
    echo 'echo "Multi-arch test image running!"' >> /test.sh && \\
    echo 'cat /test.txt' >> /test.sh && \\
    chmod +x /test.sh

CMD ["/test.sh"]
EOF

    # Build amd64 image
    log_info "Building amd64 image..."
    log_cmd docker build \
        --platform linux/amd64 \
        --build-arg TARGETARCH=amd64 \
        --build-arg TARGETOS=linux \
        -t ${AMD64_IMAGE} \
        -f ${BUILD_DIR}/Dockerfile.multiarch \
        ${BUILD_DIR}

    # Push amd64 image
    log_info "Pushing amd64 image..."
    echo -e "${BLUE}[CMD]${NC} docker push ${AMD64_IMAGE}"
    if ! docker push ${AMD64_IMAGE}; then
        log_error "Failed to push amd64 image"
        return 1
    fi

    # Build arm64 image
    log_info "Building arm64 image..."
    log_cmd docker build \
        --platform linux/arm64 \
        --build-arg TARGETARCH=arm64 \
        --build-arg TARGETOS=linux \
        -t ${ARM64_IMAGE} \
        -f ${BUILD_DIR}/Dockerfile.multiarch \
        ${BUILD_DIR}

    # Push arm64 image
    log_info "Pushing arm64 image..."
    echo -e "${BLUE}[CMD]${NC} docker push ${ARM64_IMAGE}"
    if ! docker push ${ARM64_IMAGE}; then
        log_error "Failed to push arm64 image"
        return 1
    fi

    # Create manifest list
    log_info "Creating multi-arch manifest list..."
    echo -e "${BLUE}[CMD]${NC} docker manifest create --insecure ${MULTI_ARCH_IMAGE} ${AMD64_IMAGE} ${ARM64_IMAGE}"
    docker manifest create --insecure ${MULTI_ARCH_IMAGE} \
        ${AMD64_IMAGE} \
        ${ARM64_IMAGE}

    # Annotate manifests with platform info
    log_info "Annotating manifest with platform info..."
    docker manifest annotate ${MULTI_ARCH_IMAGE} ${AMD64_IMAGE} \
        --os linux --arch amd64
    docker manifest annotate ${MULTI_ARCH_IMAGE} ${ARM64_IMAGE} \
        --os linux --arch arm64

    # Push the manifest list
    log_info "Pushing multi-arch manifest list..."
    echo -e "${BLUE}[CMD]${NC} docker manifest push --insecure ${MULTI_ARCH_IMAGE}"
    docker manifest push --insecure ${MULTI_ARCH_IMAGE}

    log_info "Multi-arch manifest created and pushed: ${MULTI_ARCH_IMAGE}"
}

# Push single-arch image and verify digest
push_single_arch_image() {
    log_step "Pushing single-arch image..."

    echo -e "${BLUE}[CMD]${NC} docker push ${SINGLE_ARCH_IMAGE}"
    local push_output
    push_output=$(docker push ${SINGLE_ARCH_IMAGE} 2>&1 | tee /dev/tty)

    # Extract and verify digest
    local digest
    digest=$(echo "$push_output" | grep -oP 'digest: \K[a-z0-9:]+' | tail -1)

    if [ -z "$digest" ]; then
        log_error "Failed to get digest from push output"
        return 1
    fi

    log_info "Pushed with digest: ${digest}"

    # Verify we can reference by digest
    log_info "Verifying digest reference..."
    echo -e "${BLUE}[CMD]${NC} docker manifest inspect --insecure ${REGISTRY}/${NAMESPACE}/${IMAGE_NAME}@${digest}"
    docker manifest inspect --insecure "${REGISTRY}/${NAMESPACE}/${IMAGE_NAME}@${digest}"
    log_info "Digest verification successful!"
}

# Verify multi-arch manifest
verify_multi_arch_manifest() {
    log_step "Verifying multi-arch manifest..."

    echo -e "${BLUE}[CMD]${NC} docker manifest inspect --insecure ${MULTI_ARCH_IMAGE}"
    local manifest
    manifest=$(docker manifest inspect --insecure ${MULTI_ARCH_IMAGE} | tee /dev/tty)

    # Check for both architectures (check for "amd64" and "arm64" in platform.architecture)
    if echo "$manifest" | grep -q '"architecture": "amd64"'; then
        log_info "Found linux/amd64 manifest"
    else
        log_error "Missing linux/amd64 manifest"
        return 1
    fi

    if echo "$manifest" | grep -q '"architecture": "arm64"'; then
        log_info "Found linux/arm64 manifest"
    else
        log_error "Missing linux/arm64 manifest"
        return 1
    fi

    # Extract digest
    local digest
    digest=$(echo "$manifest" | jq -r '.manifests[0].digest' 2>/dev/null || echo "$manifest" | grep -oP 'sha256:[a-f0-9]+' | head -1)

    if [ -n "$digest" ]; then
        log_info "Multi-arch manifest digest: ${digest}"
    fi
}

# Remove local images
cleanup_local_images() {
    log_step "Removing local images..."

    echo -e "${BLUE}[CMD]${NC} docker rmi ${SINGLE_ARCH_IMAGE}"
    docker rmi ${SINGLE_ARCH_IMAGE} || true

    echo -e "${BLUE}[CMD]${NC} docker rmi ${AMD64_IMAGE}"
    docker rmi ${AMD64_IMAGE} || true

    echo -e "${BLUE}[CMD]${NC} docker rmi ${ARM64_IMAGE}"
    docker rmi ${ARM64_IMAGE} || true

    echo -e "${BLUE}[CMD]${NC} docker rmi ${MULTI_ARCH_IMAGE}"
    docker rmi ${MULTI_ARCH_IMAGE} || true

    # Clean up manifest list
    docker manifest rm ${MULTI_ARCH_IMAGE} 2>/dev/null || true

    # Also remove any cached layers
    log_info "Pruning dangling images..."
    log_cmd docker image prune -f

    log_info "Local images removed"
}

# Pull images back
pull_images() {
    log_step "Pulling images back from registry..."

    # Pull single-arch image
    log_info "Pulling single-arch image..."
    echo -e "${BLUE}[CMD]${NC} docker pull ${SINGLE_ARCH_IMAGE}"
    if docker pull ${SINGLE_ARCH_IMAGE}; then
        log_info "Successfully pulled: ${SINGLE_ARCH_IMAGE}"
    else
        log_error "Failed to pull single-arch image"
        return 1
    fi

    # Pull multi-arch image
    log_info "Pulling multi-arch image..."
    echo -e "${BLUE}[CMD]${NC} docker pull ${MULTI_ARCH_IMAGE}"
    if docker pull ${MULTI_ARCH_IMAGE}; then
        log_info "Successfully pulled: ${MULTI_ARCH_IMAGE}"
    else
        log_error "Failed to pull multi-arch image"
        return 1
    fi
}

# Test running the images
test_images() {
    log_step "Testing pulled images..."

    # Test single-arch image
    log_info "Running single-arch image..."
    log_cmd docker run --rm ${SINGLE_ARCH_IMAGE}

    # Test multi-arch image
    log_info "Running multi-arch image..."
    log_cmd docker run --rm ${MULTI_ARCH_IMAGE}

    log_info "All images ran successfully!"
}

# Check for errors in compose logs after operations
check_compose_logs() {
    log_step "Checking compose logs for errors..."

    local error_count
    error_count=$(docker compose logs --tail=100 | grep -i "error" | grep -v "level=error msg=\"error processing" | grep -v "test" | wc -l)

    if [ "$error_count" -gt 0 ]; then
        log_warn "Found ${error_count} error messages in logs:"
        docker compose logs --tail=100 | grep -i "error" | grep -v "level=error msg=\"error processing" | grep -v "test" | tail -10
    else
        log_info "No errors found in compose logs"
    fi
}

# Multi-arch only test flow
multi_arch_only() {
    log_step "Starting ATCR multi-arch test (skipping single-arch)"
    log_info "Registry: ${REGISTRY}"
    log_info "Namespace: ${NAMESPACE}"
    log_info "Build directory: ${BUILD_DIR}"
    log_info "Verbose mode: ${VERBOSE} (set VERBOSE=0 to hide docker output)"

    # Check prerequisites
    if ! command -v docker &> /dev/null; then
        log_error "Docker is not installed"
        exit 1
    fi

    if ! command -v jq &> /dev/null; then
        log_warn "jq is not installed - some checks may be limited"
    fi

    # Run multi-arch test steps only
    check_compose_services
    build_multi_arch_images
    verify_multi_arch_manifest
    cleanup_local_images
    pull_images
    log_info "Running multi-arch image..."
    log_cmd docker run --rm ${MULTI_ARCH_IMAGE}
    check_compose_logs

    log_step "Multi-arch test completed successfully!"
    echo -e "${GREEN}========================================${NC}"
    echo -e "${GREEN}Multi-arch test passed!${NC}"
    echo -e "${GREEN}========================================${NC}"
}

# Main test flow
main() {
    log_step "Starting ATCR end-to-end test"
    log_info "Registry: ${REGISTRY}"
    log_info "Namespace: ${NAMESPACE}"
    log_info "Build directory: ${BUILD_DIR}"
    log_info "Verbose mode: ${VERBOSE} (set VERBOSE=0 to hide docker output)"

    # Check prerequisites
    if ! command -v docker &> /dev/null; then
        log_error "Docker is not installed"
        exit 1
    fi

    if ! command -v jq &> /dev/null; then
        log_warn "jq is not installed - some checks may be limited"
    fi

    # Run test steps
    check_compose_services
    build_single_arch_image
    push_single_arch_image
    build_multi_arch_images
    verify_multi_arch_manifest
    cleanup_local_images
    pull_images
    test_images
    check_compose_logs

    log_step "End-to-end test completed successfully!"
    echo -e "${GREEN}========================================${NC}"
    echo -e "${GREEN}All tests passed!${NC}"
    echo -e "${GREEN}========================================${NC}"
}

# Parse arguments
if [ "$1" = "--multi" ]; then
    multi_arch_only
else
    main "$@"
fi