#!/usr/bin/bash

if (( ! ZLC_DEVMODE )) && [[ $(id -un) != zlccli ]]; then
    exec sudo -u zlccli -n "$0" "$@"
    # unreachable
fi

(( XDEBUG )) && set -x

set -f
export LC_ALL=C
umask 0027
cd
PROGNAME=${0##*/}
PID=$$
OIFS=$IFS
ZLC_ENV=${ZLC_ENV:-/etc/logcenter/logcenter.conf}

function exit_usage() {
    printf 'Usage: %s COMMAND [ARG...]\n' "$PROGNAME"
    printf 'Helper for logcenter operations\n\n'
    printf 'Available commands:\n\n'
    sed -nre 's/^#: ?(.*)/\1/p' "${BASH_SOURCE[0]}"
    exit 0
}

exec 3>&2
exec 2> >(while read -r; do
    echo "$REPLY" >&2
    logger -t "$PROGNAME[$PID]" -p notice -- "$REPLY"
done)
LOGGER_PID=$!
trap 'exec 2>&3 3>&-; wait $LOGGER_PID' EXIT

function _log() {
    local now=$(date +%Y-%m-%dT%H:%M:%S.%3N%:z)
    local STDOUT_LABEL=${STDOUT_LABEL:-INFO}
    echo "$now $STDOUT_LABEL H=${HOSTNAME%%.*}\
${REDIRECTED_FROM:+ R=${REDIRECTED_FROM%%.*}} $PROGNAME: $*" >&2
}
function info() { STDOUT_LABEL=INFO _log "$@"; }
function fatal() { STDOUT_LABEL=FATAL _log "$@"; exit 2; }

function load_env_or_fatal() {
    # proper $PATH may be needed for sigfiled
    [[ -f /etc/profile.d/logcenter-path.sh ]] && source /etc/profile.d/logcenter-path.sh
    source "$ZLC_ENV" || fatal "Failed to load env"
    return 0
}

# $1: target host
function reexec_by_ssh() {
    local host=$1; shift
    [[ $host == 127.* || $host == localhost ]] && return 0
    [[ -n $REDIRECTED_FROM ]] && return 0
    # assume this script is the user shell
    local args=( -R "$HOSTNAME" "${ARGS[@]}" )
    info "reexec_by_ssh $host -- ${args[*]@Q}"
    exec ssh -i "$ZLC_CLI_SSH_IDENTITY" -o PreferredAuthentications=publickey \
        -o StrictHostKeyChecking=yes -o PreferredAuthentications=publickey \
        -o ConnectTimeout=5 "$host" -- "${args[*]@Q}"
}

#: help
#:     Display this help
#:
function do_help() {
    exit_usage 0
}

# $1: after select, default '*'
# $2: after where, default '1'
function archives_db_query() {
    REPLY=$'pragma query_only = ON;\n'
    local dbfile dbid dbids=() i
    while read -r dbfile; do
        dbid=${dbfile##*/stats.}; dbid=${dbid%.db}; dbids+=( "$dbid" )
        REPLY+="attach '$dbfile' as $dbid;"$'\n'
    done < <(find "$ZLC_ARCHIVES_DIR/@stats/" -mindepth 1 -maxdepth 1 \
                \( -name 'stats.db' -o -name 'stats.*.db' \))
    [[ -z $dbids ]] && fatal 'Database not found!'
    REPLY+="select ${1:-*} from ("$'\n'
    for (( i = 0; i < ${#dbids[@]}; i++ )); do
        (( i > 0 )) && REPLY+=' union '
        REPLY+="select '${dbids[i]}' as db, * from ${dbids[i]}.counters"$'\n'
    done
    REPLY+=$')\n'" where ${2:-1}"
}

#: list-archives [SQL-WHERE]
#:     List archive files with a query to the stats counters
#:     database. Return TSV ouput with header line.
#:
function do_list_archives() {
    [[ -n $ZLC_ARCHIVES_ZLCCLI_HOST ]] && reexec_by_ssh "$ZLC_ARCHIVES_ZLCCLI_HOST"
    [[ -z $ZLC_ARCHIVES_DIR ]] && fatal 'Bad environment! ZLC_ARCHIVES_DIR'
    archives_db_query '*' "${1:-1} order by fdate desc, fhour desc, src asc"
    sqlite3 -tabs -header :memory: "$REPLY"
}

#: list-archives-aggregated SELECT-CLAUSE WHERE-CLAUSE GROUP-BY
#:     List archive files with custom SELECT and optional GROUP BY for aggregation
#:     database. Return TSV ouput with header line.
#:
function do_list_archives_aggregated() {
    [[ -n $ZLC_ARCHIVES_ZLCCLI_HOST ]] && reexec_by_ssh "$ZLC_ARCHIVES_ZLCCLI_HOST"
    [[ -z $ZLC_ARCHIVES_DIR ]] && fatal 'Bad environment! ZLC_ARCHIVES_DIR'
    local select_clause="${1:-*}"
    local where_clause="${2:-1}"
    local group_by="${3:-}"

    REPLY=$'pragma query_only = ON;\n'
    local dbfile dbid dbids=() i
    while read -r dbfile; do
        dbid=${dbfile##*/stats.}; dbid=${dbid%.db}; dbids+=( "$dbid" )
        REPLY+="attach '$dbfile' as $dbid;"$'\n'
    done < <(find "$ZLC_ARCHIVES_DIR/@stats/" -mindepth 1 -maxdepth 1 \
                \( -name 'stats.db' -o -name 'stats.*.db' \))
    [[ -z $dbids ]] && fatal 'Database not found!'

    REPLY+="select $select_clause from ("$'\n'
    for (( i = 0; i < ${#dbids[@]}; i++ )); do
        (( i > 0 )) && REPLY+=' union all '
        REPLY+="select '${dbids[i]}' as db, * from ${dbids[i]}.counters"$'\n'
    done
    REPLY+=$')\n'" where $where_clause"
    [[ -n $group_by ]] && REPLY+=$'\n'"$group_by"
    REPLY+=$'\n'"order by fdate desc, fhour desc, src asc"

    sqlite3 -tabs -header :memory: "$REPLY"
}

#: list-indices [SQL-WHERE]
#:     List elastic indices with a query to the cache indices counters
#:     database. Return TSV ouput with header line.
#:
function do_list_indices() {
    [[ -n $ZLC_ELASTIC_ZLCCLI_HOST ]] && reexec_by_ssh "$ZLC_ELASTIC_ZLCCLI_HOST"

    local query
    local IFS=" "
    # Get indices data

    query=(
      "SELECT"
      "host as hostname,"
      "doc_count as docs,"
      "size_bytes as size,"
      "strftime('%Y-%m-%d %H:%M:%S', timestamp) as date"
      "FROM host_hour_data"
      "${1:+WHERE $1}"
      "ORDER BY date DESC, hostname asc;"
    )
    sqlite3 -tabs -header ${ZLC_INDICES_DB} "${query[*]}"
}

#: list-indices-aggregated SELECT-FIELDS WHERE-CLAUSE GROUP-BY
#:     List elastic indices with custom SELECT and optional GROUP BY for aggregation
#:     database. Return TSV ouput with header line.
#:
function do_list_indices_aggregated() {
    [[ -n $ZLC_ELASTIC_ZLCCLI_HOST ]] && reexec_by_ssh "$ZLC_ELASTIC_ZLCCLI_HOST"

    local select_fields="${1:-host as hostname, doc_count as docs, size_bytes as size, strftime('%Y-%m-%d %H:%M:%S', timestamp) as date}"
    local where_clause="${2:-1}"
    local group_by="${3:-}"
    local query
    local IFS=" "
    
    query=(
      "SELECT"
      "$select_fields"
      "FROM host_hour_data"
      "WHERE $where_clause"
      "${group_by:+$group_by}"
      "ORDER BY date DESC, hostname asc;"
    )
    
    # Use ZLC_INDICES_DB if set, otherwise default path
    sqlite3 -tabs -header "${ZLC_INDICES_DB}" "${query[*]}"
}

#: get-archive [--proto] SRC,YYYY-MM-DD,H...
#:     Stream an archive file, with optional proto header
#:
function do_get_archive() {
    [[ -n $ZLC_ARCHIVES_ZLCCLI_HOST ]] && reexec_by_ssh "$ZLC_ARCHIVES_ZLCCLI_HOST"
    [[ -z $ZLC_ARCHIVES_DIR ]] && fatal 'Bad environment! ZLC_ARCHIVES_DIR'
    local proto=; [[ $1 == --proto ]] && { proto=1; shift; }
    local IFS=,; set -- $1; IFS=$OIFS
    local src=$1 fdate=$2 fhour=$3 q="'"
    local file=$(archives_db_query 'file' "src='${src//$q/$q$q}' and
        fdate='${fdate//$q/$q$q}' and fhour='${fhour//$q/$q$q}'
        order by db asc limit 1"
        sqlite3 :memory: "$REPLY")
    [[ -z $file ]] && fatal 'No match!'

    # FIXME: Only the first matching file is sent. There may be several
    # files in case of multiple archivers; same to support multiple
    # SRC,YYYY-MM-DD,H arguments. In such cases, a tarball with all matching
    # files should be to sent.
    local suffix fsfile
    for suffix in .xz ''; do
        fsfile="$ZLC_ARCHIVES_DIR/${file}${suffix}"
        [[ -f $fsfile ]] || continue
        [[ -n $proto ]] && printf '\x16\x17\x18%s\x16%s\n' \
            "${src}_${fsfile##*/}" $(stat -c %s "$fsfile")
        exec cat "$fsfile"
    done
    fatal 'Not found!'
}

#: get-archives-usage
#:     Output archives used bytes and max bytes on stdout.
#:     Values are sperated by a tab.
function do_get_archives_usage() {
    [[ -n $ZLC_ARCHIVES_ZLCCLI_HOST ]] && reexec_by_ssh "$ZLC_ARCHIVES_ZLCCLI_HOST"
    [[ -z $ZLC_ARCHIVES_DIR ]] && fatal 'Bad environment! ZLC_ARCHIVES_DIR'

    local usage=( $(df -P "$ZLC_ARCHIVES_DIR" -B1 |tail -n +2) )
    [[ -z ${usage[2]} ]] && fatal 'Failed to get archives disk usage!'
    local max=${ZLC_ARCHIVES_MAX_SIZE:-${usage[1]}}
    [[ -z $max ]] && fatal 'Failed to get archives disk max size!'
    printf "%s\t%s\n" "${usage[2]}" "$max"
}

#: get-indices-usage
#:     Output Elasticsearch used bytes and max bytes on stdout.
#:     Values are sperated by a tab.
function do_get_elastic_usage() {
    [[ -n $ZLC_ELASTIC_ZLCCLI_HOST ]] && reexec_by_ssh "$ZLC_ELASTIC_ZLCCLI_HOST"

    local usage=( $(es-curl _nodes/stats/fs |jq -r '.nodes |to_entries |
        map( .value |select(.roles|join(",")|test("data")) ) |
        reduce .[] as $i ({ used:0, total:0 };
            .used += ($i.fs.total.total_in_bytes - $i.fs.total.available_in_bytes) |
            .total += $i.fs.total.total_in_bytes
        ) |
        "\(.used // 0) \(.total // 0)"') )
    [[ -z ${usage[0]} ]] && fatal 'Failed to get elastic disk usage!'
    local max=${ZLC_ELASTIC_MAX_SIZE:-${usage[1]}}
    [[ -z $max ]] && fatal 'Failed to get elastic disk max size!'
    printf "%s\t%s\n" "${usage[0]}" "$max"
}

#: list-aliases
#:     Dump aliases
#:
function do_list_aliases() {
    [[ -n $ZLC_ARCHIVES_ZLCCLI_HOST ]] && reexec_by_ssh "$ZLC_ARCHIVES_ZLCCLI_HOST"
    [[ -z $ZLC_ALIASES_DB ]] && fatal 'Bad environment! ZLC_ALIASES_DB'

    if [[ ! -f "$ZLC_ALIASES_DB" ]]; then
        sqlite3 "$ZLC_ALIASES_DB" <<"__EOF__"
            CREATE TABLE IF NOT EXISTS aliases (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                ip TEXT NOT NULL UNIQUE,
                hostname TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            );
            CREATE INDEX IF NOT EXISTS idx_aliases_ip ON aliases(ip);
__EOF__
    fi

    # Query the database and format as JSON
    # Assuming the table structure has 'ip' and 'hostname' columns
    sqlite3 -tabs -header "$ZLC_ALIASES_DB" "SELECT ip, hostname FROM aliases;"
}

#: set-alias
#:     Set or delete an alias for an IP address
#:     Usage: set-alias IP [ALIAS]
#:     If ALIAS is empty or not provided, the alias for IP will be deleted
#:
function do_set_alias() {
    [[ -n $ZLC_ARCHIVES_ZLCCLI_HOST ]] && reexec_by_ssh "$ZLC_ARCHIVES_ZLCCLI_HOST"
    [[ -z $ZLC_ALIASES_DB ]] && fatal 'Bad environment! ZLC_ALIASES_DB'

    if [[ ! -f "$ZLC_ALIASES_DB" ]]; then
        # initialize database
        do_list_aliases > /dev/null
    fi

    local ipaddr=$1; shift
    local alias=${1:-}; shift
    local db=$ZLC_ALIASES_DB

    # IP address is always valid at this point
    # Check if this IP already exists in the database
    existing_hostname=$(sqlite3 "$db" "SELECT hostname FROM aliases WHERE ip='$ipaddr' LIMIT 1")

    # If alias is empty or not provided, delete the entry
    if [[ -z "$alias" ]]; then
        if [[ -n "$existing_hostname" ]]; then
            result=$(sqlite3 "$db" "DELETE FROM aliases WHERE ip='$ipaddr'")
            info "Deleted alias: $ipaddr ($existing_hostname)"
        else
            info "No alias found for IP: $ipaddr"
        fi
        return
    fi

    if [[ -n "$existing_hostname" ]]; then
        # IP exists, just update the hostname
        result=$(sqlite3 "$db" "UPDATE aliases SET hostname='$alias' WHERE ip='$ipaddr'")
        info "Updated alias: $ipaddr ($existing_hostname) → $alias"
    else
        # IP doesn't exist, we need to insert
        # But first check if the new hostname is already used for another IP
        existing_ip=$(sqlite3 "$db" "SELECT ip FROM aliases WHERE hostname='$alias' LIMIT 1")

        if [[ -n "$existing_ip" ]]; then
            # This hostname is already assigned to another IP, update that record
            result=$(sqlite3 "$db" "UPDATE aliases SET ip='$ipaddr' WHERE hostname='$alias'")
            info "Reassigned hostname $alias from $existing_ip to $ipaddr"
        else
            # New hostname isn't used, safe to insert
            result=$(sqlite3 "$db" "INSERT INTO aliases (ip, hostname) VALUES ('$ipaddr', '$alias')")
            info "Added new alias: $ipaddr → $alias"
        fi
    fi
}

load_env_or_fatal

if [[ $1 == -c ]]; then
    # user shell
    declare -ga "ARGS=( $2 )"
else
    # classic exec
    ARGS=( "$@" )
fi
set --

for (( i=0; i<${#ARGS[@]}; i++ )); do
    case "${ARGS[i]}" in
        -R|--redirected-from) REDIRECTED_FROM=${ARGS[i++ +1]} ;;
        -h|--help) exit_usage 0 ;;
        --) ((i++)); break ;;
        *) break ;;
    esac
done
ARGS=( "${ARGS[i]}" "${ARGS[@]:i+1}" )
info "ARGS: ${ARGS[*]@Q}"

fn="do_${ARGS[0]//[^[:alnum:]]/_}"
declare -f -F "$fn" >/dev/null || fatal 'Bad command!'
"$fn" "${ARGS[@]:1}"
