#!/bin/bash
# Validates that all docker-compose image references use SHA256 digests
# Prevents supply chain attacks by ensuring immutable image references

set -euo pipefail

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

# Configuration - images that are allowed without SHA digests
ALLOWED_EXCEPTIONS=(
    "localdeepresearch/local-deep-research:latest"  # Own image, built by CI
)

# Check if an image reference is in the exceptions list
is_exception() {
    local image="$1"
    for exception in "${ALLOWED_EXCEPTIONS[@]}"; do
        if [[ "$image" == "$exception" ]]; then
            return 0
        fi
    done
    return 1
}

# Validate a single docker-compose file
validate_compose_file() {
    local file="$1"
    local violations=0
    local line_num=0

    while IFS= read -r line; do
        line_num=$((line_num + 1))

        # Check if this line contains an image reference
        if [[ "$line" =~ ^[[:space:]]*image:[[:space:]]*(.+)$ ]]; then
            local image="${BASH_REMATCH[1]}"
            image=$(echo "$image" | tr -d '"' | xargs)  # Remove quotes and whitespace

            # Skip if it's an exception
            if is_exception "$image"; then
                echo -e "${YELLOW}  Line $line_num: $image (exception)${NC}"
                continue
            fi

            # Check if image has SHA digest
            if [[ ! "$image" =~ @sha256: ]]; then
                echo -e "${RED}  ❌ Line $line_num: Missing SHA digest${NC}"
                echo -e "${RED}     Image: $image${NC}"
                violations=$((violations + 1))
            else
                echo -e "${GREEN}  ✓ Line $line_num: $image${NC}"
            fi
        fi
    done < "$file"

    return $violations
}

# Main validation logic
main() {
    local total_violations=0
    local files_checked=0

    echo "🔍 Validating docker-compose image pinning..."
    echo ""

    # Find all docker-compose files
    while IFS= read -r compose_file; do
        # Skip cookiecutter templates and examples (documentation only)
        if [[ "$compose_file" =~ cookiecutter-docker/ ]] || [[ "$compose_file" =~ examples/ ]]; then
            echo -e "${YELLOW}⏭  Skipping: $compose_file (template/example)${NC}"
            continue
        fi

        echo "📄 Checking: $compose_file"
        if validate_compose_file "$compose_file"; then
            : # No violations
        else
            violations=$?
            total_violations=$((total_violations + violations))
        fi
        files_checked=$((files_checked + 1))
        echo ""
    done < <(find . -name "docker-compose*.yml" -o -name "docker-compose*.yaml")

    # Summary
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "📊 Summary:"
    echo "   Files checked: $files_checked"
    echo "   Violations: $total_violations"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

    if [ $total_violations -gt 0 ]; then
        echo ""
        echo -e "${RED}❌ Found $total_violations unpinned images in docker-compose files${NC}"
        echo ""
        echo "Images must use SHA256 digests for security and reproducibility."
        echo ""
        echo "To fix:"
        echo "  1. Pull the image: docker pull <image:tag>"
        echo "  2. Get digest: docker inspect <image:tag> | jq -r '.[0].RepoDigests[0]'"
        echo "  3. Update file: image: <image:tag>@sha256:..."
        echo ""
        echo "Example:"
        echo "  # Bad"
        echo "  image: ollama/ollama:latest"
        echo ""
        echo "  # Good"
        echo "  image: ollama/ollama:latest@sha256:8850b8b33936b9fb..."
        exit 1
    fi

    echo ""
    echo -e "${GREEN}✅ All docker-compose images properly pinned${NC}"
    exit 0
}

main "$@"