#!/usr/bin/env bash
# vim: set ft=sh:

OR_DIR="/var/lib/openreplay"
APP_NS="${APP_NS:-app}"
DB_NS="${DB_NS:-db}"
OR_REPO="${OR_REPO:-'https://github.com/openreplay/openreplay'}"
# For example HELM_OPTIONS="--set dbMigrationUpstreamBranch=dev"
#HELM_OPTIONS=""
# If you want to install the dev version. It can be any branch or tag.
#OR_VERSION="dev"
# Only upgrade Openreplay
# UPGRADE_OR_ONLY=1 openreplay -u
# Pass override file to the installation
# OVERRIDE_FILE=override.yaml

[[ -d $OR_DIR ]] || {
    sudo mkdir $OR_DIR
}
export PATH=/var/lib/openreplay:$PATH

# Check arch
architecture=$(uname -m)
pkg_type="amd64"
if [[ "$architecture" == "x86_64" ]]; then
    echo "This machine is AMD (x86_64) architecture."
elif [[ "$architecture" == "aarch64" ]]; then
    pkg_type="arm64"
    echo "This machine is ARM architecture."
else
    echo "This machine has another architecture: $architecture"
fi

[[ $(awk '/enterpriseEditionLicense/{print $2}' <"/var/lib/openreplay/vars.yaml") != "" ]] && EE=true

tools=(
    zyedidia/eget
    stern/stern
    hidetatz/kubecolor
)

# Ref: https://stackoverflow.com/questions/5947742/how-to-change-the-output-color-of-echo-in-linux
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BWHITE='\033[1;37m'
NC='\033[0m' # No Color

# Checking whether the app exists or we do have to upgade.
function exists() {
    which "${1}" &>/dev/null
    return $?
}

function err_cd() {
    if ! cd "$1" &>/dev/null; then
        log err not able to cd to "$1"
        exit 100
    fi
}

function log() {
    case "$1" in
    info)
        shift
        echo -e "${GREEN}[INFO]" "$@" "${NC}"
        return
        ;;
    warn)
        shift
        echo -e "${YELLOW}[WARN]" "$@" "${NC}"
        return
        ;;
    debug)
        shift
        echo -e "${YELLOW}[DEBUG]" "$@" "${NC}"
        return
        ;;
    title)
        shift
        echo -e "\n${BWHITE}-" "$@" "${NC}"
        return
        ;;
    err)
        shift
        echo -e "${RED}[ERROR]" "$@" "${NC}"
        exit 100
        ;;
    *)
        echo "Not supported log format"
        ;;
    esac
    echo "[Error]" "$@"
    exit 100
}

# To run kubeconfig run
# `KUBECONFIG=/path/to/file openreplay -s`
[[ -f /etc/rancher/k3s/k3s.yaml ]] && k3s_path="/etc/rancher/k3s/k3s.yaml"
[[ -f "${HOME}/.kube/config" ]] && local_kube_config_path="${HOME}/.kube/config"
k3s_path=${k3s_path:-$local_kube_config_path}
export KUBECONFIG=${KUBECONFIG:-$k3s_path}
[[ -z $KUBECONFIG ]] && log warn "No kubeconfig file found. Exiting"
log info "Using KUBECONFIG $KUBECONFIG"
tmp_dir=$(mktemp -d)

function install_packages() {

    [[ -e "$OR_DIR/eget" ]] || {
        cd "$tmp_dir" || log err "Not able to cd to tmp dir $tmp_dir"
        curl --version &>/dev/null || log err "curl not found. Please install"
        curl -SsL https://zyedidia.github.io/eget.sh | sh - >/dev/null
        sudo mv eget $OR_DIR
        err_cd -
    }

    for package in "${tools[@]}"; do
        log info Installing "$(awk -F/ '{print $2}' <<<$package)"
        sudo /var/lib/openreplay/eget -q --upgrade-only --to "${OR_DIR}" "$package"
    done
    log info Installing k9s
    sudo /var/lib/openreplay/eget -q --upgrade-only --to "$OR_DIR" derailed/k9s --asset=tar.gz --asset=^sbom
    log info Installing yq
    sudo /var/lib/openreplay/eget -q --upgrade-only --to "$OR_DIR" mikefarah/yq --asset=^tar.gz
    log info Installing helm
    sudo /var/lib/openreplay/eget -q --upgrade-only --to "$OR_DIR" https://get.helm.sh/helm-v3.10.2-linux-${pkg_type}.tar.gz -f helm
    log info Installing kubectl
    sudo /var/lib/openreplay/eget -q --upgrade-only --to "$OR_DIR" https://dl.k8s.io/release/v1.25.0/bin/linux/${pkg_type}/kubectl
    [[ $pkg_type != "arm64" ]] && {
        log info Installing Busybox
        sudo /var/lib/openreplay/eget -q --upgrade-only --to "$OR_DIR" https://busybox.net/downloads/binaries/1.35.0-x86_64-linux-musl/busybox
    }
    date | sudo tee $OR_DIR/packages.lock &>/dev/null
}

function help() {

    echo -e ${BWHITE}
    cat <<"EOF"
  ___                   ____            _
 / _ \ _ __   ___ _ __ |  _ \ ___ _ __ | | __ _ _   _
| | | | '_ \ / _ \ '_ \| |_) / _ \ '_ \| |/ _` | | | |
| |_| | |_) |  __/ | | |  _ <  __/ |_) | | (_| | |_| |
 \___/| .__/ \___|_| |_|_| \_\___| .__/|_|\__,_|\__, |
      |_|                        |_|            |___/
EOF
    echo -e ${NC}

    log info "
  Usage: openreplay [ -h | --help ]
                    [ -s | --status ]
                    [ -i | --install DOMAIN_NAME ]
                    [ -u | --upgrade (fetch lastest patches for installed release. ${BWHITE}RELEASE_UPGRADE=1 openreplay -u${GREEN} to upgrade release.)]
                    [ -U | --deprecated-upgrade /path/to/old_vars.yaml]
                    [ -r | --restart ]
                    [ -R | --Reload ]
                    [ -c | --cleanup N(in days) (optional --force to skip the confirmation)]
                    [ -e | --edit ]
                    [ -p | --install-packages ]
                    [ -l | --logs SERVICE ] 
         Services: alerts assets assist chalice 
                   db ender frontend heuristics 
                   http integrations nginx-controller 
                   peers sink sourcemapreader storage 
    "
    return
}

function status() {
    log info OpenReplay Version
    # awk '(NR<2)' < "$OR_DIR/vars.yaml"
    awk '/fromVersion/{print $2}' <"${OR_DIR}/vars.yaml"
    yq .config.version /var/lib/openreplay/vars.yaml || true
    log info Disk
    df -h /var
    log info Memory
    free -mh
    log info CPU
    uname -a
    # Print only the fist line.
    awk '(NR<2)' </etc/os-release
    echo "CPU Count: $(nproc)"
    log info Kubernetes
    kubecolor version --short
    log info Openreplay Component
    kubecolor get po -n "${APP_NS}"
    kubecolor get po -n "${DB_NS}"
    return
}

# Function to upgrade helm openreplay app.
function or_helm_upgrade() {
    set -o pipefail
    log_file="${tmp_dir}/helm.log"
    state=$1
    chart_names=(
        toolings
        openreplay
    )
    [[ $UPGRADE_OR_ONLY -eq 1 ]] && chart_names=(openreplay)
    # Cleaning up toolings
    [[ $CLEANUP_TOOLING -eq 1 ]] && {
        helm uninstall toolings -n "$APP_NS"
    }
    if [[ $state == "reload" ]]; then
        chart_names=(openreplay)
        HELM_OPTIONS="${HELM_OPTIONS} --set skipMigration=true"
    fi
    for chart in "${chart_names[@]}"; do
        [[ -z $OR_VERSION ]] || HELM_OPTIONS="${HELM_OPTIONS} --set dbMigrationUpstreamBranch=${OR_VERSION}"
        log info helm upgrade --install "$chart" ./"$chart" -n "$APP_NS" --wait -f ./vars.yaml --atomic --debug $HELM_OPTIONS 2>&1 | tee -a "${log_file}"
        if ! helm upgrade --install "$chart" ./"$chart" -n "$APP_NS" --wait -f ./vars.yaml --atomic --debug $HELM_OPTIONS 2>&1 | tee -a "${log_file}"; then
            log err "
        Installation failed, run ${BWHITE}cat ${log_file}${RED} for more info

        If logs aren't verbose, run ${BWHITE}openreplay --status${RED}

        If pods are in failed state, run ${BWHITE}openreplay --logs <pod-name>${RED}
        "
        fi
    done
    set +o pipefail
    return
}

function upgrade_old() {
    old_vars_path="$1"
    [[ -f $old_vars_path ]] || log err "No configuration file ${BWHITE}$old_vars_path${RED}.
          If you're updating from version older than ${BWHITE}v1.10.0${RED}, for example ${BWHITE}v1.9.0${RED}:
            ${BWHITE}RELEASE_UPGRADE=1 openreplay --deprecated-upgrade ~/openreplay_v1.9.0/scripts/helmcharts/vars.yaml${RED}.
          If you're having a custom installation,
            ${BWHITE}RELEASE_UPGRADE=1 openreplay --deprecated-upgrade /path/to/vars.yaml${RED}.
          "
    or_version=$(busybox awk '/fromVersion/{print $2}' <"${old_vars_path}")
    sudo cp "${old_vars_path}" ${OR_DIR}/vars.yaml.backup."${or_version//\"/}"_"$(date +%Y%m%d-%H%M%S)" || log err "Not able to copy old vars.yaml"
    sudo cp "${old_vars_path}" ${OR_DIR}/vars.yaml || log err "Not able to copy old vars.yaml"
    upgrade
}

function clone_repo() {
    err_cd "$tmp_dir"
    log info "Working directory $tmp_dir"
    git_options="-b ${OR_VERSION:-main}"
    log info "git clone ${OR_REPO} --depth 1 $git_options"
    eval git clone "${OR_REPO}" --depth 1 $git_options
    return
}

function install() {
    domain_name=$1
    # Check existing installation
    [[ -f ${OR_DIR}/vars.yaml ]] && {
        or_version=$(busybox awk '/fromVersion/{print $2}' <"${OR_DIR}/vars.yaml")
        log err "Openreplay installation ${BWHITE}${or_version}${RED} found. If you want to upgrade, run ${BWHITE}openreplay -u${RED}"
    }
    # Installing OR
    install_packages
    log title "Installing OpenReplay"
    clone_repo
    err_cd "$tmp_dir/openreplay/scripts/helmcharts"
    # process override file while installing
    # OVERRIDE_FILE=/tmp/override.yaml openreplay -i
    [[ -z $OVERRIDE_FILE ]] || {
        [[ -f $OVERRIDE_FILE ]] || {
            echo "OVERRIDE_FILE does not exist, exiting"
            exit 100
        }
        override_tmpdir=$(pwd)
        err_cd /tmp
        cp "$OVERRIDE_FILE" new_vars.yaml
        cp "$override_tmpdir/vars.yaml" old_vars.yaml
        yq '(load("new_vars.yaml") | .. | select(tag != "!!map" and tag != "!!seq")) as $i ireduce(.; setpath($i | path; $i))' old_vars.yaml >vars.yaml
        mv vars.yaml "$override_tmpdir"
        err_cd "$override_tmpdir"
    }
    SKIP_K8S_TOOLS=1 DOMAIN_NAME=$domain_name bash init.sh
    return
}

function cleanup() {
    # Confirmation for deletion. Do you want to delete Postgres/Minio(session) data before $date ?
    delete_from_number_days=$1
    delete_from_date=$(date +%Y-%m-%d -d "$delete_from_number_days day ago")

    # Check if --force flag is present
    if [[ $2 == --force ]]; then
        log info "Deleting data without confirmation..."
    else
        log debug "Do you want to delete the data captured on and before ${BWHITE}$delete_from_date${YELLOW}?"
        read -p "Are you sure[y/n]? " -n 1 -r
        echo # (optional) move to a new line
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            log err "Cancelling data deletion"
            return 1 # Exit with an error code to indicate cancellation
        fi
    fi

    # Run pg cleanup
    pguser=$(awk '/postgresqlUser/{print $2}' <"${OR_DIR}/vars.yaml" | xargs)
    pgpassword=$(awk '/postgresqlPassword/{print $2}' <"${OR_DIR}/vars.yaml" | xargs)
    pghost=$(awk '/postgresqlHost/{print $2}' <"${OR_DIR}/vars.yaml" | xargs)
    pgport=$(awk '/postgresqlPort/{print $2}' <"${OR_DIR}/vars.yaml" | xargs)
    pgdatabase=$(awk '/postgresqlDatabase/{print $2}' <"${OR_DIR}/vars.yaml" | xargs)
    cleanup_query="DELETE FROM public.sessions WHERE start_ts < extract(epoch from '${delete_from_date}'::date) * 1000;"
    [[ $EE ]] && cleanup_query="DELETE FROM public.sessions WHERE start_ts < extract(epoch from '${delete_from_date}'::date) * 1000 AND session_id NOT IN (SELECT session_id FROM user_favorite_sessions);"
    kubectl delete po -n "${APP_NS}" pg-cleanup &>/dev/null || true
    kubectl run pg-cleanup -n "${APP_NS}" \
        --restart=Never \
        --env PGHOST="$pghost" \
        --env PGUSER="$pguser" \
        --env PGDATABASE="$pgdatabase" \
        --env PGPASSWORD="$pgpassword" \
        --env PGPORT="$pgport" \
        --image bitnami/postgresql -- psql -c "$cleanup_query"
    # Run minio cleanup
    MINIO_ACCESS_KEY=$(awk '/accessKey/{print $NF}' <"${OR_DIR}/vars.yaml" | tail -n1 | xargs)
    MINIO_SECRET_KEY=$(awk '/secretKey/{print $NF}' <"${OR_DIR}/vars.yaml" | tail -n1 | xargs)
    MINIO_HOST=$(awk '/endpoint/{print $NF}' <"${OR_DIR}/vars.yaml" | tail -n1 | xargs)
    kubectl delete po -n "${APP_NS}" minio-cleanup &>/dev/null || true
    kubectl run minio-cleanup -n "${APP_NS}" \
        --restart=Never \
        --env MINIO_HOST="$pghost" \
        --image bitnami/minio:2020.10.9-debian-10-r6 -- /bin/sh -c "
      mc alias set minio $MINIO_HOST $MINIO_ACCESS_KEY $MINIO_SECRET_KEY &&
      mc rm --recursive --dangerous --force --older-than ${delete_from_number_days}d minio/mobs
      "
    log info "Postgres data cleanup process initiated. Postgres will automatically vacuum deleted rows when the database is idle. This may take up a few days to free the disk space."
    log info "Minio (where recordings are stored) cleanup process initiated."
    log info "Run ${BWHITE}openreplay -s${GREEN} to check the status of the cleanup process and available disk space."
    return
}

function upgrade() {
    # TODO:
    # 1. store vars.yaml in central place.
    # 3. In upgrade you'll have to clone the repo
    # 3. How to update package. Because openreplay -u will be done from old update script
    # 4. Update from Version
    exists git || log err "Git not found. Please install"
    [[ -f ${OR_DIR}/vars.yaml ]] || log err "No configuration file ${BWHITE}${OR_DIR}/vars.yaml${RED}.
        If you're updating from version older than ${BWHITE}v1.10.0${RED}, for example ${BWHITE}v1.9.0${RED}:
          ${BWHITE}RELEASE_UPGRADE=1 openreplay --deprecated-upgrade ~/openreplay_v1.9.0/scripts/helmcharts/vars.yaml${RED}.
        If you're having a custom installation,
          ${BWHITE}RELEASE_UPGRADE=1 openreplay --deprecated-upgrade /path/to/vars.yaml${RED}.
        "
    or_version=$(busybox awk '/fromVersion/{print $2}' <"${OR_DIR}/vars.yaml") || {
        log err "${BWHITE}${OR_DIR}/vars.yaml${RED} not found. 
        Please do ${BWHITE}openreplay --deprecated-upgrade /path/to/vars.yaml${RED}
      "
    }

    # Unless its upgrade release, always checkout same tag.
    [[ $RELEASE_UPGRADE -eq 1 ]] || OR_VERSION=${OR_VERSION:-$or_version}

    time_now=$(date +%m-%d-%Y-%I%M%S)
    # Creating backup dir of current installation
    [[ -d "$OR_DIR/openreplay" ]] && sudo mv "$OR_DIR/openreplay" "$OR_DIR/openreplay_${or_version//\"/}_${time_now}"

    clone_repo
    err_cd openreplay/scripts/helmcharts
    install_packages
    [[ -d /openreplay ]] && sudo chown -R 1001:1001 /openreplay

    # Merge prefrerences
    cp $OR_DIR/vars.yaml old_vars.yaml
    or_new_version=$(awk '/fromVersion/{print $2}' <"vars.yaml")
    yq '(load("old_vars.yaml") | .. | select(tag != "!!map" and tag != "!!seq")) as $i ireduce(.; setpath($i | path; $i))' vars.yaml >new_vars.yaml
    mv new_vars.yaml vars.yaml
    or_helm_upgrade

    # Update the version
    busybox sed -i "s/fromVersion.*/fromVersion: ${or_new_version}/" vars.yaml
    patch_version
    sudo mv ./openreplay-cli /bin/openreplay
    sudo chmod +x /bin/openreplay
    sudo mv ./vars.yaml "$OR_DIR"
    sudo cp -rf ../../../openreplay $OR_DIR/
    log info "Configuration file is saved in /var/lib/openreplay/vars.yaml"
    log info "Run ${BWHITE}openreplay -h${GREEN} to see the cli information to manage OpenReplay."

    err_cd -
    return
}

function patch_version() {
    # Patching config version for console
    version=$(/var/lib/openreplay/yq '.fromVersion' /var/lib/openreplay/vars.yaml)-$(sudo git rev-parse --short HEAD)
    # Create a temporary YAML file with the generated version
    cat <<EOF >/tmp/version_snippet.yaml
config:
  version: $version
EOF
    sudo /var/lib/openreplay/yq eval-all 'select(fileIndex == 0) * select(fileIndex == 1)' vars.yaml /tmp/version_snippet.yaml -i
    rm /tmp/version_snippet.yaml
}

function reload() {
    err_cd $OR_DIR/openreplay/scripts/helmcharts
    sudo cp -f $OR_DIR/vars.yaml .
    or_helm_upgrade reload
    return
}

function clean_tmp_dir() {
    [[ -z $SKIP_DELETE_TMP_DIR ]] && rm -rf "${tmp_dir}"
}

[[ -f $OR_DIR/packages.lock ]] || {
    log title Installing packages "${NC}"
    install_packages
}

PARSED_ARGUMENTS=$(busybox getopt -a -n openreplay -o Rrevpi:uhsl:U:c: --long reload,edit,restart,verbose,install-packages,install:,upgrade,help,status,logs,deprecated-upgrade:,cleanup:,force -- "$@")
VALID_ARGUMENTS=$?
if [[ "$VALID_ARGUMENTS" != "0" ]]; then
    help
    exit 100
fi

eval set -- "$PARSED_ARGUMENTS"
while :; do
    case "$1" in
    -v | --verbose)
        VERBOSE=1
        echo $VERBOSE
        clean_tmp_dir
        shift
        ;;
    -h | --help)
        help
        clean_tmp_dir
        exit 0
        ;;
    -i | --install)
        log title "Installing OpenReplay"
        install "$2"
        clean_tmp_dir
        exit 0
        ;;
    -p | --install-packages)
        log title "Updating/Installing dependency packages"
        install_packages
        clean_tmp_dir
        exit 0
        ;;
    -u | --upgrade)
        if [[ $RELEASE_UPGRADE -eq 1 ]]; then
            log title "Upgrading OpenReplay to Latest Release"
            CLEANUP_TOOLING=1
        else
            log title "Applying Latest OpenReplay Patches"
            UPGRADE_OR_ONLY=${UPGRADE_OR_ONLY:-1}
        fi
        upgrade
        clean_tmp_dir
        exit 0
        ;;
    -U | --deprecated-upgrade)
        log title "[Deprected] Upgrading OpenReplay"
        upgrade_old "$2"
        clean_tmp_dir
        exit 0
        ;;
    -c | --cleanup)
        log title "Cleaning up data older than $2 days"
        cleanup "$2" "$3"
        clean_tmp_dir
        exit 0
        ;;
    -r | --restart)
        log title "Restarting OpenReplay Components"
        kubecolor rollout restart deployment -n "${APP_NS}"
        kubecolor rollout status deployment -n "${APP_NS}"
        clean_tmp_dir
        exit 0
        ;;
    -R | --reload)
        log title "Reloading OpenReplay Components"
        reload
        clean_tmp_dir
        exit 0
        ;;
    -e | --edit)
        log title "Editing OpenReplay"
        [[ -f ${OR_DIR}/vars.yaml ]] || {
            log err "
              Couldn't open ${BWHITE}${OR_DIR}/vars.yaml${RED}. Seems like a custom installation.
              Edit the proper ${BWHITE}vars.yaml${RED} and run ${BWHITE}openreplay -R${RED}
              Or ${BWHITE}helm upgrade openreplay -n app openreplay/scripts/helmcharts/openreplay -f openreplay/scripts/helmcharts/vars.yaml --debug --atomic"
            exit 100
        }
        busybox md5sum /var/lib/openreplay/vars.yaml >"${tmp_dir}/var.yaml.md5"
        sudo vim -n ${OR_DIR}/vars.yaml
        /var/lib/openreplay/yq 'true' /var/lib/openreplay/vars.yaml &>/dev/null || {
            log debug "seems like the edit is not correct. Rerun ${BWHITE}openreplay -e${YELLOW} and fix the issue in config file."
            clean_tmp_dir
            exit 100
        }
        if busybox md5sum -c "${tmp_dir}/var.yaml.md5" &>/dev/null; then
            log info "No change detected in ${BWHITE}${OR_DIR}/vars.yaml${GREEN}. Not reloading"
        else
            reload
        fi
        clean_tmp_dir
        exit 0
        ;;
    -s | --status)
        log title "Checking OpenReplay Components Status"
        status
        clean_tmp_dir
        exit 0
        ;;
    -l | --logs)
        # Skipping double quotes because we want globbing. For example
        # ./openreplay -l "chalice --tail 10"
        stern -A --container-state=running,terminated $2
        clean_tmp_dir
        exit 0
        ;;
    # -- means the end of the arguments; drop this, and break out of the while loop
    --)
        shift
        break
        ;;
    # If invalid options were passed, then getopt should have reported an error,
    # which we checked as VALID_ARGUMENTS when getopt was called...
    *)
        echo "Unexpected option: $1 - this should not happen."
        help
        clean_tmp_dir
        ;;
    esac
done

[ $# -eq 0 ] && help
clean_tmp_dir
