#!/bin/bash

# Copyright Security Onion Solutions LLC and/or licensed to Security Onion Solutions LLC under one
# or more contributor license agreements. Licensed under the Elastic License 2.0; you may not use
# this file except in compliance with the Elastic License 2.0.
{%- import_yaml 'elasticfleet/defaults.yaml' as ELASTICFLEETDEFAULTS %}
{%- set SUPPORTED_PACKAGES = salt['pillar.get']('elasticfleet:packages', default=ELASTICFLEETDEFAULTS.elasticfleet.packages, merge=True) %}

. /usr/sbin/so-elastic-fleet-common

PKG_LOAD_FAILURES=0
PKG_LOAD_FAILURES_NAMES=()
PKG_UPGRADED=0

cleanup_elasticsearch_fleet_transforms() {
    local transforms transform_id attempt
    
    if ! transforms=$(so-elasticsearch-query "_transform/logs-elasticsearch.index_pivot-default-*" --retry 1 --retry-delay 5); then
        return 0
    fi

    while IFS= read -r transform_id; do
        [ -n "$transform_id" ] || continue
        for attempt in {1..3}; do
            if so-elasticsearch-query "_transform/$transform_id?force=true" -XDELETE --fail --retry 1 --retry-delay 5; then
                break
            fi
            sleep 5
        done
    done < <(jq -r '.transforms[]?.id' <<< "$transforms")
}

{%- for PACKAGE in SUPPORTED_PACKAGES %}
if INSTALLED_VERSION=$(elastic_fleet_package_version_check "{{ PACKAGE }}") && LATEST_VERSION=$(elastic_fleet_package_latest_version_check "{{ PACKAGE }}"); then

    if [ "$INSTALLED_VERSION" == "$LATEST_VERSION" ]; then
        echo "{{ PACKAGE }} integration version $INSTALLED_VERSION is already at the reported latest version $LATEST_VERSION, skipping upgrade."
    else
{%- if PACKAGE == 'elasticsearch' %}
        cleanup_elasticsearch_fleet_transforms
{%- endif %}
        echo "Upgrading {{ PACKAGE }} package from $INSTALLED_VERSION to version $LATEST_VERSION..."
        if ! elastic_fleet_package_install "{{ PACKAGE }}" "$LATEST_VERSION"; then
            PKG_LOAD_FAILURES=$((PKG_LOAD_FAILURES + 1))
            PKG_LOAD_FAILURES_NAMES+=("{{ PACKAGE }}")
        # check that package has upgraded to the expected version after install command
        elif ! LATEST_VERSION=$(elastic_fleet_package_latest_version_check "{{ PACKAGE }}"); then
            echo "ERROR: Failed to get latest version information for integration {{ PACKAGE }} after upgrade attempt"
            PKG_LOAD_FAILURES=$((PKG_LOAD_FAILURES + 1))
            PKG_LOAD_FAILURES_NAMES+=("{{ PACKAGE }}")
        elif INSTALLED_VERSION=$(elastic_fleet_package_version_check "{{ PACKAGE }}") && [ "$INSTALLED_VERSION" == "$LATEST_VERSION" ]; then
            echo "{{ PACKAGE }} integration upgraded to version $LATEST_VERSION."
            PKG_UPGRADED=$((PKG_UPGRADED + 1))
        else
            echo "ERROR: {{ PACKAGE }} integration still at ${INSTALLED_VERSION:-unknown}; expected $LATEST_VERSION"
            PKG_LOAD_FAILURES=$((PKG_LOAD_FAILURES + 1))
            PKG_LOAD_FAILURES_NAMES+=("{{ PACKAGE }}")
        fi
    fi
else
    echo "ERROR: Failed to get version information for integration {{ PACKAGE }}"
    PKG_LOAD_FAILURES=$((PKG_LOAD_FAILURES + 1))
    PKG_LOAD_FAILURES_NAMES+=("{{ PACKAGE }}")
fi
{%- endfor %}

if [ $PKG_UPGRADED -gt 0 ]; then
    echo "Elasticsearch template statefiles cleared after $PKG_UPGRADED package upgrade(s), so templates can reload."
    rm -f /opt/so/state/estemplates.txt /opt/so/state/addon_estemplates.txt
fi

if [ $PKG_LOAD_FAILURES -gt 0 ]; then
    echo "ERROR: Failed to upgrade $PKG_LOAD_FAILURES package(s):"
    for PKG in "${PKG_LOAD_FAILURES_NAMES[@]}"; do
        echo " - $PKG"
    done
    # exit 1 on failure to upgrade a default package, allow salt to handle retries
    exit 1
else
    echo "Successfully upgraded all packages."
fi
