#!/bin/bash

# usage: etc/final_deployment.sh [-e|--docker-export]
#
# This script handles post-installation deployment tasks for Prole.
# 1st use case: -e|--docker-export
#   Runs docker image export for each image required to deploy Prole.
#   Artifacts are written to DOCKER_IMPORT_DIR.

usage() {
    echo "Usage: $0 [options]"
    echo ""
    echo "Options:"
    echo "  -e, --docker-export    Export required docker images to DOCKER_IMPORT_DIR"
    echo ""
}

if [[ $# -eq 0 ]]; then
    usage
    exit 1
fi

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

# Load environment if env.sh exists
if [[ -f "${PROJECT_ROOT}/env.sh" ]]; then
    source "${PROJECT_ROOT}/env.sh"
fi

DOCKER_IMPORT_DIR="${DOCKER_IMPORT_DIR:-${PROJECT_ROOT}/data/docker-import}"

export_docker_images() {
    echo "Exporting Docker images to: ${DOCKER_IMPORT_DIR}"
    mkdir -p "${DOCKER_IMPORT_DIR}"

    # Collect images from K8s manifests
    # Robust extraction of image names from YAML files
    find "${PROJECT_ROOT}/k8s/prole" "${PROJECT_ROOT}/k8s/openbao" "${PROJECT_ROOT}/k8s/opentofu" -name "*.yaml" -exec grep -h "image:" {} + | awk -F'image:' '{print $2}' | awk '{print $1}' | sed "s/['\"]//g" > /tmp/prole_images.txt

    # Also check Supabase compose if it exists
    if [[ -f "${PROJECT_ROOT}/supabase/docker/docker-compose.yml" ]]; then
        grep -h "image:" "${PROJECT_ROOT}/supabase/docker/docker-compose.yml" | awk -F'image:' '{print $2}' | awk '{print $1}' | sed "s/['\"]//g" >> /tmp/prole_images.txt
    fi

    # Add Kerberos proxy image if defined
    KRB_IMG="${KRB5_AD_PROXY_IMAGE:-alpine/socat}"
    if [[ -n "$KRB_IMG" ]]; then
        echo "$KRB_IMG" >> /tmp/prole_images.txt
    fi

    # Add Kerberos test image (used by init_kerberos_test.sh)
    KRB_TEST_IMG="${KRB5_TEST_IMAGE:-${PROLE_KRB_TEST_IMAGE:-ubuntu:24.04}}"
    if [[ -n "$KRB_TEST_IMG" ]]; then
        echo "$KRB_TEST_IMG" >> /tmp/prole_images.txt
    fi

    # Unique images
    sort -u /tmp/prole_images.txt > /tmp/prole_images_unique.txt

    while read -r image; do
        if [[ -z "$image" ]]; then continue; fi
        
        # Ensure we have the image locally
        echo "Checking image: $image"
        if ! docker image inspect "$image" >/dev/null 2>&1; then
            echo "Pulling $image..."
            docker pull "$image"
        fi

        safe_name=$(echo "$image" | sed 's/\//_/g' | sed 's/:/_/g')
        tar_path="${DOCKER_IMPORT_DIR}/${safe_name}.tar"
        
        echo "Exporting $image to $tar_path..."
        docker save "$image" -o "$tar_path"
        if [[ $? -eq 0 ]]; then
            echo "[OK] Exported $image"
        else
            echo "[ERROR] Failed to export $image"
        fi
    done < /tmp/prole_images_unique.txt

    rm -f /tmp/prole_images.txt /tmp/prole_images_unique.txt
    echo "Docker export complete."
}

case "$1" in
    -e|--docker-export)
        export_docker_images
        ;;
    *)
        usage
        exit 1
        ;;
esac