#!/bin/bash
#
# Snowflake Snapshot Viewer
#
# This script mounts a Snowflake snapshot into a JupyterLab container for inspection.
# The snapshot will be available at /data in the JupyterLab environment.
#
# Prerequisites:
#   - Snowflake CLI (snow): https://docs.snowflake.com/en/developer-guide/snowflake-cli/index
#   - Docker: https://docs.docker.com/get-started/get-docker/
#   - jq: https://jqlang.github.io/jq/download/
#

set -e

# Check for required commands
MISSING_COMMANDS=()

if ! command -v snow >/dev/null 2>&1; then
  MISSING_COMMANDS+=("snow (Snowflake CLI)")
fi

if ! command -v docker >/dev/null 2>&1; then
  MISSING_COMMANDS+=("docker")
fi

if ! command -v jq >/dev/null 2>&1; then
  MISSING_COMMANDS+=("jq")
fi

if [ ${#MISSING_COMMANDS[@]} -gt 0 ]; then
  echo "Error: Missing required commands:"
  for cmd in "${MISSING_COMMANDS[@]}"; do
    echo "  - $cmd"
  done
  echo ""
  echo "Please install the missing prerequisites:"
  echo "  - Snowflake CLI: https://docs.snowflake.com/en/developer-guide/snowflake-cli/index"
  echo "  - Docker: https://docs.docker.com/get-started/get-docker/"
  echo "  - jq: https://jqlang.github.io/jq/download/"
  exit 1
fi

SNAPSHOT_VIEWER_DATABASE=SNAPSHOT_VIEWER
SNAPSHOT_VIEWER_COMPUTE_POOL=${SNAPSHOT_VIEWER_DATABASE}
SNAPSHOT_VIEWER_SCHEMA=${SNAPSHOT_VIEWER_DATABASE}.DATA
SNAPSHOT_VIEWER_IMAGE_REPOSITORY=${SNAPSHOT_VIEWER_SCHEMA}.IMAGE_REPOSITORY
SNAPSHOT_VIEWER_SERVICE=${SNAPSHOT_VIEWER_SCHEMA}.SNAPSHOT_VIEWER

# Cleanup function to ensure resources are always deleted
cleanup() {
  if [ -n "$SNOWFLAKE_CONNECTION_NAME" ]; then
    echo ""
    echo "==================================================================="
    echo "Cleaning up resources..."
    echo "==================================================================="

    # Stop services and delete all created resources
    cat <<EOF | snow sql -i --connection "$SNOWFLAKE_CONNECTION_NAME" 2>/dev/null || true
-- Stop all services in the compute pool
ALTER COMPUTE POOL IF EXISTS $SNAPSHOT_VIEWER_COMPUTE_POOL STOP ALL;

-- Delete the compute pool
DROP COMPUTE POOL IF EXISTS $SNAPSHOT_VIEWER_COMPUTE_POOL;

-- Delete the image repository
DROP IMAGE REPOSITORY IF EXISTS $SNAPSHOT_VIEWER_IMAGE_REPOSITORY;
EOF

    echo "Cleanup complete. All temporary resources have been deleted."
    echo "Your original snapshot remains unchanged."
  fi
}

# Set trap to ensure cleanup runs on exit
trap cleanup EXIT INT TERM

echo "==================================================================="
echo "Snapshot Viewer Settings:"
echo "==================================================================="
echo "  Database:          $SNAPSHOT_VIEWER_DATABASE"
echo "  Compute Pool:      $SNAPSHOT_VIEWER_COMPUTE_POOL"
echo "  Schema:            $SNAPSHOT_VIEWER_SCHEMA"
echo "  Image Repository:  $SNAPSHOT_VIEWER_IMAGE_REPOSITORY"
echo "  Service:           $SNAPSHOT_VIEWER_SERVICE"
echo ""

# Collect required configuration from user
echo "Fetching Snowflake connections..."
CONNECTIONS_JSON=$(snow connection list --format=JSON)
connections=($(echo "$CONNECTIONS_JSON" | jq -r '.[].connection_name'))

if [ ${#connections[@]} -eq 0 ]; then
  echo "Error: No Snowflake connections found. Please configure a connection first."
  exit 1
fi

echo ""
echo "Select Snowflake connection:"
select SNOWFLAKE_CONNECTION_NAME in "${connections[@]}"; do
  [[ -n $SNOWFLAKE_CONNECTION_NAME ]] && break
  echo "Invalid selection"
done

echo ""
echo "Selected connection: $SNOWFLAKE_CONNECTION_NAME"

# Extract account from selected connection and construct registry URL
SNOWFLAKE_ACCOUNT=$(echo "$CONNECTIONS_JSON" | jq -r --arg conn "$SNOWFLAKE_CONNECTION_NAME" '.[] | select(.connection_name == $conn) | .parameters.account')
SNOWFLAKE_REGISTRY="$SNOWFLAKE_ACCOUNT.registry.snowflakecomputing.com"

echo "Snowflake registry: $SNOWFLAKE_REGISTRY"
echo ""

read -p "Enter Snowflake database to view snapshots from: " SNOWFLAKE_DATABASE
if [ -z "$SNOWFLAKE_DATABASE" ]; then
  echo "Error: Snowflake database is required"
  exit 1
fi

echo ""
echo "Fetching snapshots from database $SNOWFLAKE_DATABASE..."
SNAPSHOTS_JSON=$(snow sql -q "show snapshots in database \"$SNOWFLAKE_DATABASE\"" --connection "$SNOWFLAKE_CONNECTION_NAME" --format=JSON)

# Create array of snapshot display names with size and date
snapshot_options=()
snapshot_names=()
snapshot_sizes=()
snapshot_databases=()
snapshot_schemas=()

while IFS= read -r line; do
  name=$(echo "$line" | jq -r '.name')
  size=$(echo "$line" | jq -r '.size')
  created=$(echo "$line" | jq -r '.created_on' | cut -d'T' -f1)
  database=$(echo "$line" | jq -r '.database_name')
  schema=$(echo "$line" | jq -r '.schema_name')

  snapshot_options+=("$name ($size GiB, created: $created)")
  snapshot_names+=("$name")
  snapshot_sizes+=("$size")
  snapshot_databases+=("$database")
  snapshot_schemas+=("$schema")
done < <(echo "$SNAPSHOTS_JSON" | jq -c '.[]')

if [ ${#snapshot_options[@]} -eq 0 ]; then
  echo "Error: No snapshots found in database $SNOWFLAKE_DATABASE"
  exit 1
fi

echo ""
echo "Select snapshot:"
PS3="Enter selection number: "
select snapshot_display in "${snapshot_options[@]}"; do
  if [[ -n $snapshot_display ]]; then
    # Get the index of the selected option
    selected_index=$((REPLY - 1))
    SNOWFLAKE_SNAPSHOT_NAME="${snapshot_names[$selected_index]}"
    SNOWFLAKE_SNAPSHOT_SIZE="${snapshot_sizes[$selected_index]}"
    SNOWFLAKE_SNAPSHOT_DATABASE="${snapshot_databases[$selected_index]}"
    SNOWFLAKE_SNAPSHOT_SCHEMA="${snapshot_schemas[$selected_index]}"
    SNOWFLAKE_SNAPSHOT="\"$SNOWFLAKE_SNAPSHOT_DATABASE\".\"$SNOWFLAKE_SNAPSHOT_SCHEMA\".\"$SNOWFLAKE_SNAPSHOT_NAME\""
    break
  fi
  echo "Invalid selection"
done

echo ""
echo "Selected snapshot: $SNOWFLAKE_SNAPSHOT"
echo "Snapshot size: ${SNOWFLAKE_SNAPSHOT_SIZE} GiB"

# Transform SNAPSHOT_VIEWER_IMAGE_REPOSITORY: convert to lowercase and replace . with /
SNAPSHOT_VIEWER_IMAGE_PATH=$(echo "$SNAPSHOT_VIEWER_IMAGE_REPOSITORY" | tr '[:upper:]' '[:lower:]' | tr '.' '/')
SNAPSHOT_VIEWER_FULL_IMAGE="$SNOWFLAKE_REGISTRY/$SNAPSHOT_VIEWER_IMAGE_PATH/jupyterlab:latest"

echo ""
echo "Configuration Summary:"
echo "  Connection:  $SNOWFLAKE_CONNECTION_NAME"
echo "  Account:     $SNOWFLAKE_ACCOUNT"
echo "  Registry:    $SNOWFLAKE_REGISTRY"
echo "  Database:    $SNOWFLAKE_DATABASE"
echo "  Snapshot:    $SNOWFLAKE_SNAPSHOT"
echo "  Size:        ${SNOWFLAKE_SNAPSHOT_SIZE} GiB"
echo "  Image:       $SNAPSHOT_VIEWER_FULL_IMAGE"
echo ""
read -p "Proceed with these settings? (y/N): " confirm
if [[ ! "$confirm" =~ ^[Yy]$ ]]; then
  echo "Aborted"
  exit 0
fi

echo ""
echo "==================================================================="
echo "Step 1: Creating Snowflake resources"
echo "==================================================================="

# Create database, schemas, image repository, and compute pool
cat <<EOF | snow sql -i --connection "$SNOWFLAKE_CONNECTION_NAME"
-- Verify current Snowflake context
SELECT CURRENT_ROLE(), CURRENT_ORGANIZATION_NAME(), CURRENT_ACCOUNT_NAME();

-- Create database and schemas for snapshot viewer resources
CREATE DATABASE IF NOT EXISTS $SNAPSHOT_VIEWER_DATABASE;
CREATE SCHEMA IF NOT EXISTS $SNAPSHOT_VIEWER_SCHEMA;

-- Create image repository to store the JupyterLab container
CREATE IMAGE REPOSITORY IF NOT EXISTS $SNAPSHOT_VIEWER_IMAGE_REPOSITORY;

-- Create a small compute pool for running the snapshot viewer service
CREATE COMPUTE POOL IF NOT EXISTS $SNAPSHOT_VIEWER_COMPUTE_POOL
    MIN_NODES=1
    MAX_NODES=1
    INSTANCE_FAMILY=CPU_X64_XS
    AUTO_SUSPEND_SECS=60;
EOF

echo ""
echo "==================================================================="
echo "Step 2: Building and pushing JupyterLab container"
echo "==================================================================="

# Log into Snowflake container registry
snow spcs image-registry login --connection "$SNOWFLAKE_CONNECTION_NAME"

# Build a minimal JupyterLab container
cat <<EOF | docker build -t "$SNAPSHOT_VIEWER_FULL_IMAGE" -
FROM jupyter/base-notebook
USER root
ENV JUPYTER_ENABLE_LAB=yes
WORKDIR /
EOF

# Push the container to Snowflake's registry
docker push "$SNAPSHOT_VIEWER_FULL_IMAGE"

echo ""
echo "==================================================================="
echo "Step 3: Creating JupyterLab service with snapshot mounted"
echo "==================================================================="

# Deploy JupyterLab service with the snapshot mounted at /data
cat <<EOF | snow sql -i --connection "$SNOWFLAKE_CONNECTION_NAME"
CREATE SERVICE IF NOT EXISTS $SNAPSHOT_VIEWER_SERVICE
    IN COMPUTE POOL $SNAPSHOT_VIEWER_COMPUTE_POOL
    MIN_INSTANCES=1
    MAX_INSTANCES=1
    FROM SPECIFICATION \$\$
spec:
  containers:
  - name: "jupyterlab"
    image: "/$SNAPSHOT_VIEWER_IMAGE_PATH/jupyterlab:latest"
    env:
      GRANT_SUDO: "yes"
    args:
     - "jupyter"
     - "lab"
     - "--ip=0.0.0.0"
     - "--no-browser"
     - "--allow-root"
     - "--NotebookApp.token=''"
     - "--ServerApp.root_dir=/"
    readinessProbe:
      port: 8888
      path: "/api"
    volumeMounts:
     - name: "snapshot"
       mountPath: "/data"
  endpoints:
  - name: "http"
    port: 8888
    public: true
    protocol: HTTP
  volumes:
  - name: "snapshot"
    source: block
    size: "${SNOWFLAKE_SNAPSHOT_SIZE}Gi"
    blockConfig:
      initialContents:
        fromSnapshot: '$SNOWFLAKE_SNAPSHOT'
    \$\$;

SELECT SYSTEM\$WAIT_FOR_SERVICES(300, '$SNAPSHOT_VIEWER_SERVICE');
EOF

echo ""
echo "==================================================================="
echo "Step 4: Waiting for service endpoint"
echo "==================================================================="

# Poll for the service endpoint (may take up to a minute after jupyterlab is running)
echo "Waiting for service endpoint to become available..."
INGRESS_URL=""
for i in {1..12}; do
  echo "Checking endpoints (attempt $i/12)..."

  ENDPOINTS_JSON=$(snow sql -q "show endpoints in service $SNAPSHOT_VIEWER_SERVICE" --connection "$SNOWFLAKE_CONNECTION_NAME" --format=JSON)
  INGRESS_URL=$(echo "$ENDPOINTS_JSON" | jq -r '.[0].ingress_url // empty')

  # Check if we have a valid URL (not the provisioning message)
  if [[ -n "$INGRESS_URL" && "$INGRESS_URL" != "Endpoints provisioning in progress"* ]]; then
    echo "Endpoint is ready!"
    break
  fi

  if [ $i -lt 12 ]; then
    echo "Endpoint still provisioning... waiting 10 seconds"
    sleep 10
  fi
done

if [[ -z "$INGRESS_URL" || "$INGRESS_URL" == "Endpoints provisioning in progress"* ]]; then
  echo ""
  echo "Warning: Endpoint may not be fully provisioned yet."
  echo "Current status: $INGRESS_URL"
  echo "You may need to wait a few more minutes and check manually."
fi

echo ""
echo "==================================================================="
echo "JupyterLab is ready!"
echo "==================================================================="
echo ""
JUPYTERLAB_URL="https://$INGRESS_URL"
echo "JupyterLab URL: $JUPYTERLAB_URL"
echo ""
echo "Opening JupyterLab in your browser..."

# Open browser based on platform
if [[ "$OSTYPE" == "darwin"* ]]; then
  # macOS
  open "$JUPYTERLAB_URL"
elif command -v xdg-open >/dev/null 2>&1; then
  # Linux with xdg-open
  xdg-open "$JUPYTERLAB_URL"
else
  echo "Could not auto-open browser. Please manually open: $JUPYTERLAB_URL"
fi

echo ""
echo "Access Instructions:"
echo "  - Your snapshot data is available at /data in JupyterLab"
echo "  - You can create a tarball and download by running:"
echo "    tar -czvf data.tar.gz /data"
echo ""
echo "Note: Authentication is disabled for convenience (no token required)"
echo ""
echo "When finished, press [Enter] to clean up all resources..."
read