#!/usr/bin/env bash
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SEARCH_DIR="$SCRIPT_DIR"
CONFIG_PATH=""

while true; do
  CANDIDATE="$SEARCH_DIR/.clawdbot/credentials/portainer/config.json"
  if [[ -f "$CANDIDATE" ]]; then
    CONFIG_PATH="$CANDIDATE"
    break
  fi
  PARENT="$(dirname "$SEARCH_DIR")"
  if [[ "$PARENT" == "$SEARCH_DIR" ]]; then
    break
  fi
  SEARCH_DIR="$PARENT"
done

if [[ -z "$CONFIG_PATH" && -f "$HOME/.clawdbot/credentials/portainer/config.json" ]]; then
  CONFIG_PATH="$HOME/.clawdbot/credentials/portainer/config.json"
fi

err() {
  echo "Error: $*" >&2
  exit 1
}

require_cmd() {
  command -v "$1" >/dev/null 2>&1 || err "Required command not found: $1"
}

usage() {
  cat >&2 <<'EOF'
Usage: update-stack.sh "<stack-name>" "<compose-file>" [options]

Updates a Portainer stack by name with a new docker-compose file.
Preserves existing env vars unless --env-file is provided.

Arguments:
  stack-name    Name of the stack to update
  compose-file  Path to the new docker-compose.yml file

Options:
  --env-file <file>     Path to a file with env vars (format: NAME=value per line)
                        If not provided, existing env vars are preserved.
  --pull                Force pull images and redeploy (like docker compose down/pull/up).
  --prune-old           After update, remove images that were used by this stack
                        but are no longer in use (old versions left dangling).

Requires config at:
  workspace .clawdbot/credentials/portainer/config.json or ~/.clawdbot/credentials/portainer/config.json
EOF
  exit 2
}

# Parse arguments
[[ $# -ge 2 ]] || usage
STACK_NAME="$1"
COMPOSE_FILE="$2"
shift 2

ENV_FILE=""
PRUNE_OLD=false
PULL_IMAGE=false
while [[ $# -gt 0 ]]; do
  case "$1" in
    --env-file)
      [[ $# -ge 2 ]] || err "--env-file requires a value"
      ENV_FILE="$2"
      shift 2
      ;;
    --pull)
      PULL_IMAGE=true
      shift
      ;;
    --prune-old)
      PRUNE_OLD=true
      shift
      ;;
    *)
      err "Unknown option: $1"
      ;;
  esac
done

require_cmd curl
require_cmd jq

[[ -f "$CONFIG_PATH" ]] || err "Missing config file: $CONFIG_PATH"
[[ -f "$COMPOSE_FILE" ]] || err "Compose file not found: $COMPOSE_FILE"

BASE_URL="$(jq -r '.base_url // empty' "$CONFIG_PATH")"
API_KEY="$(jq -r '.api_key // empty' "$CONFIG_PATH")"

[[ -n "$BASE_URL" ]] || err "config.base_url is missing"
[[ -n "$API_KEY" ]] || err "config.api_key is missing"

BASE_URL="${BASE_URL%/}"

# Resolve stack ID from name
STACK_ID="$(bash "$SCRIPT_DIR/get-stack-id.sh" "$STACK_NAME")"
[[ -n "$STACK_ID" ]] || err "Failed to resolve stack ID for: $STACK_NAME"

# Get current stack info for EndpointId and existing env
STACK_INFO="$(curl -sS -w $'\n%{http_code}' \
  -H "X-API-Key: $API_KEY" \
  "$BASE_URL/api/stacks/$STACK_ID")"

http_code="$(printf '%s' "$STACK_INFO" | tail -n1)"
body="$(printf '%s' "$STACK_INFO" | sed '$d')"

if [[ "$http_code" -lt 200 || "$http_code" -ge 300 ]]; then
  msg="$(printf '%s' "$body" | jq -r '.message // empty' 2>/dev/null || true)"
  [[ -n "$msg" ]] || msg="$body"
  err "Failed to fetch stack info (HTTP $http_code): $msg"
fi

ENDPOINT_ID="$(printf '%s' "$body" | jq -r '.EndpointId')"
[[ -n "$ENDPOINT_ID" && "$ENDPOINT_ID" != "null" ]] || err "Could not determine endpoint ID for stack"

# Capture old image IDs before update (if --prune-old)
OLD_IMAGE_IDS=""
if [[ "$PRUNE_OLD" == true ]]; then
  # Get containers for this stack using the compose project label
  # The label format is: com.docker.compose.project=<stack-name>
  FILTERS="$(jq -n --arg project "$STACK_NAME" '{"label": ["com.docker.compose.project=\($project)"]}')"
  
  CONTAINERS_RESPONSE="$(curl -sS -w $'\n%{http_code}' \
    -H "X-API-Key: $API_KEY" \
    -G \
    --data-urlencode "all=1" \
    --data-urlencode "filters=$FILTERS" \
    "$BASE_URL/api/endpoints/$ENDPOINT_ID/docker/containers/json")"
  
  containers_http_code="$(printf '%s' "$CONTAINERS_RESPONSE" | tail -n1)"
  containers_body="$(printf '%s' "$CONTAINERS_RESPONSE" | sed '$d')"
  
  if [[ "$containers_http_code" -ge 200 && "$containers_http_code" -lt 300 ]]; then
    # Extract image IDs from containers
    OLD_IMAGE_IDS="$(printf '%s' "$containers_body" | jq -r '.[].ImageID // empty' | sort -u | tr '\n' ' ' | sed 's/ $//')"
    if [[ -n "$OLD_IMAGE_IDS" ]]; then
      echo "Captured old image IDs for pruning: $OLD_IMAGE_IDS"
    else
      echo "No existing container images found to track for pruning"
      PRUNE_OLD=false
    fi
  else
    echo "Warning: Could not fetch containers for prune tracking, skipping prune"
    PRUNE_OLD=false
  fi
fi

# Determine env vars to use
if [[ -n "$ENV_FILE" ]]; then
  # Read env file and convert to Portainer format
  [[ -f "$ENV_FILE" ]] || err "Env file not found: $ENV_FILE"
  ENV_JSON="$(grep -v '^#' "$ENV_FILE" | grep -v '^$' | while IFS='=' read -r name value; do
    [[ -n "$name" ]] && printf '{"name":"%s","value":"%s"},' "$name" "$value"
  done | sed 's/,$//')"
  ENV_PAYLOAD="[$ENV_JSON]"
else
  # Preserve existing env vars
  ENV_PAYLOAD="$(printf '%s' "$body" | jq '.Env // []')"
fi

# Read new compose file content
COMPOSE_CONTENT="$(cat "$COMPOSE_FILE")"

# Build JSON payload (include pullImage if --pull flag is set)
if [[ "$PULL_IMAGE" == true ]]; then
  PAYLOAD="$(jq -n \
    --argjson env "$ENV_PAYLOAD" \
    --arg compose "$COMPOSE_CONTENT" \
    '{Env: $env, StackFileContent: $compose, pullImage: true}')"
else
  PAYLOAD="$(jq -n \
    --argjson env "$ENV_PAYLOAD" \
    --arg compose "$COMPOSE_CONTENT" \
    '{Env: $env, StackFileContent: $compose}')"
fi

# Build URL
UPDATE_URL="$BASE_URL/api/stacks/$STACK_ID?endpointId=$ENDPOINT_ID"

# Update the stack
response="$(curl -sS -w $'\n%{http_code}' -X PUT \
  -H "X-API-Key: $API_KEY" \
  -H "Content-Type: application/json" \
  -d "$PAYLOAD" \
  "$UPDATE_URL")"

http_code="$(printf '%s' "$response" | tail -n1)"
body="$(printf '%s' "$response" | sed '$d')"

if [[ "$http_code" -lt 200 || "$http_code" -ge 300 ]]; then
  msg="$(printf '%s' "$body" | jq -r '.message // empty' 2>/dev/null || true)"
  [[ -n "$msg" ]] || msg="$body"
  err "Failed to update stack (HTTP $http_code): $msg"
fi

echo "Stack '$STACK_NAME' (ID: $STACK_ID) updated successfully"

# Prune old images if requested and we captured image IDs
if [[ "$PRUNE_OLD" == true && -n "$OLD_IMAGE_IDS" ]]; then
  echo "Pruning old images no longer in use..."
  
  for IMAGE_ID in $OLD_IMAGE_IDS; do
    # Check if image is still in use by any container
    IMAGE_FILTERS="$(jq -n --arg imageId "$IMAGE_ID" '{"ancestor": [$imageId]}')"
    
    USAGE_RESPONSE="$(curl -sS -w $'\n%{http_code}' \
      -H "X-API-Key: $API_KEY" \
      -G \
      --data-urlencode "all=1" \
      --data-urlencode "filters=$IMAGE_FILTERS" \
      "$BASE_URL/api/endpoints/$ENDPOINT_ID/docker/containers/json")"
    
    usage_http_code="$(printf '%s' "$USAGE_RESPONSE" | tail -n1)"
    usage_body="$(printf '%s' "$USAGE_RESPONSE" | sed '$d')"
    
    if [[ "$usage_http_code" -ge 200 && "$usage_http_code" -lt 300 ]]; then
      container_count="$(printf '%s' "$usage_body" | jq 'length')"
      
      if [[ "$container_count" -eq 0 ]]; then
        # Image not in use, safe to remove
        DELETE_RESPONSE="$(curl -sS -w $'\n%{http_code}' -X DELETE \
          -H "X-API-Key: $API_KEY" \
          "$BASE_URL/api/endpoints/$ENDPOINT_ID/docker/images/$IMAGE_ID?force=false")"
        
        delete_http_code="$(printf '%s' "$DELETE_RESPONSE" | tail -n1)"
        delete_body="$(printf '%s' "$DELETE_RESPONSE" | sed '$d')"
        
        if [[ "$delete_http_code" -ge 200 && "$delete_http_code" -lt 300 ]]; then
          echo "Removed unused image: $IMAGE_ID"
        else
          msg="$(printf '%s' "$delete_body" | jq -r '.message // empty' 2>/dev/null || true)"
          [[ -n "$msg" ]] || msg="$delete_body"
          echo "Warning: Could not remove image $IMAGE_ID: $msg"
        fi
      else
        echo "Image $IMAGE_ID still in use by $container_count container(s), skipping"
      fi
    fi
  done
  
  echo "Pruning complete"
fi