#!/usr/bin/env bash
# metrics-query: Execute a metrics query against Axiom MetricsDB
#
# Usage: metrics-query [-p name=value]... [-w pixels] [--pixel-per-point n] \
#        <deployment> <mpl> <startTime> <endTime>
#
# Times: RFC3339 (e.g. 2025-01-01T00:00:00Z) or relative (e.g. now-1h, now-1d).
#
# Adaptive resolution ($__interval):
#   Reference $__interval anywhere a Duration is expected (e.g.
#   `align to $__interval using avg`, `bucket to $__interval ...`) and the
#   server resolves it to a "nice" step computed from the query time range and
#   the target chart width. No `param $__interval` declaration is needed -- the
#   metrics service registers it automatically. Tune the density with:
#     -w / --chart-width <pixels>   target chart width; the server aims for
#                                   ~chart-width/pixel-per-point buckets
#                                   (default ~500 buckets when -w is omitted).
#     --pixel-per-point <n>         pixels per data point (server default 10).
#   Both are forwarded under the request body's queryOptions object.
#
# Parameter values (-p / --param name=value, repeatable):
#   For each MPL parameter declared in the query (e.g. `param $svc: string;`),
#   pass the variable name without the leading `$` and an MPL literal as the
#   value. The script applies the `param__` prefix to the key and forwards the
#   value verbatim under the request body's `params` object — i.e. `-p svc=...`
#   becomes `params.param__svc`. Callers are responsible for formatting the
#   value as a valid MPL literal (per the metrics spec) and for any quoting
#   the literal requires; the server parses it according to the variable's
#   declared type. Parameters declared optional in the MPL (see metrics-spec)
#   may be omitted; required parameters must be supplied or the server returns
#   HTTP 400.
#
# Examples:
#   metrics-query prod '`otel-metrics`:`http.server.duration` | align to 5m using avg | group by `endpoint` using sum' \
#     '2025-06-01T00:00:00Z' '2025-06-02T00:00:00Z'
#   metrics-query prod '`otel-metrics`:`http.server.duration` | align to 5m using avg' now-1h now
#   # With parameters $svc: string and $window: Duration:
#   metrics-query -p svc='"frontend"' -p window='5m' prod \
#     'param $svc: string; param $window: Duration; `otel-metrics`:`http.server.duration` | where `service.name` == $svc | align to $window using avg' \
#     now-1h now

set -euo pipefail

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

PARAMS=()
POSITIONAL=()
CHART_WIDTH=""
PIXEL_PER_POINT=""
while [[ $# -gt 0 ]]; do
    case "$1" in
        -p|--param)
            if [[ $# -lt 2 ]]; then
                echo "Error: $1 requires an argument of the form name=value" >&2
                exit 1
            fi
            PARAMS+=("$2")
            shift 2
            ;;
        --param=*)
            PARAMS+=("${1#--param=}")
            shift
            ;;
        -w|--chart-width)
            if [[ $# -lt 2 ]]; then
                echo "Error: $1 requires a pixel-width argument" >&2
                exit 1
            fi
            CHART_WIDTH="$2"
            shift 2
            ;;
        --chart-width=*)
            CHART_WIDTH="${1#--chart-width=}"
            shift
            ;;
        --pixel-per-point)
            if [[ $# -lt 2 ]]; then
                echo "Error: $1 requires an integer argument" >&2
                exit 1
            fi
            PIXEL_PER_POINT="$2"
            shift 2
            ;;
        --pixel-per-point=*)
            PIXEL_PER_POINT="${1#--pixel-per-point=}"
            shift
            ;;
        --)
            shift
            while [[ $# -gt 0 ]]; do POSITIONAL+=("$1"); shift; done
            ;;
        *)
            POSITIONAL+=("$1")
            shift
            ;;
    esac
done

DEPLOYMENT="${POSITIONAL[0]:-}"
MPL="${POSITIONAL[1]:-}"
START_TIME="${POSITIONAL[2]:-}"
END_TIME="${POSITIONAL[3]:-}"

if [[ -z "$DEPLOYMENT" || -z "$MPL" || -z "$START_TIME" || -z "$END_TIME" ]]; then
    echo "Usage: metrics-query [-p name=value]... [-w pixels] [--pixel-per-point n] <deployment> <mpl> <startTime> <endTime>" >&2
    echo "" >&2
    echo "Times: RFC3339 (e.g. 2025-01-01T00:00:00Z) or relative (e.g. now-1h, now-1d)." >&2
    echo "" >&2
    echo "-p / --param name=value (repeatable): supply an MPL parameter value." >&2
    echo "  name  - variable name without the leading \$ (e.g. 'svc' for \$svc)." >&2
    echo "  value - MPL literal, forwarded verbatim under params.param__<name>." >&2
    echo "" >&2
    echo "-w / --chart-width <pixels>   target chart width; lets the server resolve" >&2
    echo "                              \$__interval to a nice step (queryOptions)." >&2
    echo "--pixel-per-point <n>         pixels per data point (server default 10)." >&2
    exit 1
fi

# Validate and split params into parallel name/value arrays.
PARAM_NAMES=()
PARAM_VALUES=()
if [[ ${#PARAMS[@]} -gt 0 ]]; then
    for entry in "${PARAMS[@]}"; do
        if [[ "$entry" != *"="* ]]; then
            echo "Error: -p expects name=value, got: $entry" >&2
            exit 1
        fi
        name="${entry%%=*}"
        value="${entry#*=}"
        if [[ -z "$name" ]]; then
            echo "Error: -p name is empty in: $entry" >&2
            exit 1
        fi
        if [[ "${name:0:1}" == "\$" ]]; then
            echo "Error: -p name must not include the leading \$ (got: $name); pass the bare variable name" >&2
            exit 1
        fi
        PARAM_NAMES+=("$name")
        PARAM_VALUES+=("$value")
    done
fi

# Validate the optional chart-sizing options. They must be positive integers;
# they are forwarded under queryOptions so the server can resolve $__interval.
if [[ -n "$CHART_WIDTH" && ! "$CHART_WIDTH" =~ ^[1-9][0-9]*$ ]]; then
    echo "Error: --chart-width must be a positive integer (got: $CHART_WIDTH)" >&2
    exit 1
fi
if [[ -n "$PIXEL_PER_POINT" && ! "$PIXEL_PER_POINT" =~ ^[1-9][0-9]*$ ]]; then
    echo "Error: --pixel-per-point must be a positive integer (got: $PIXEL_PER_POINT)" >&2
    exit 1
fi

# Extract dataset name from MPL: `dataset`:`metric` ... or dataset:`metric` ...
# Strip leading `param <name>: <type>;` declarations first so their `:` doesn't
# get mistaken for the dataset:metric separator.
PIPELINE="$MPL"
while [[ "$PIPELINE" =~ ^[[:space:]]*[Pp]aram[[:space:]] ]]; do
    if [[ "$PIPELINE" != *";"* ]]; then break; fi
    PIPELINE="${PIPELINE#*;}"
done
DATASET=$(echo "$PIPELINE" | sed 's/`//g' | cut -d: -f1 | tr -d '[:space:]')

QUERY_EDGE_DEPLOYMENT=""
if [[ -n "$DATASET" ]]; then
    RESOLVED_URL=$("$SCRIPT_DIR/resolve-url" "$DEPLOYMENT" "$DATASET" 2>/dev/null || true)
    if [[ -n "$RESOLVED_URL" ]]; then
        export AXIOM_URL_OVERRIDE="$RESOLVED_URL"
        # Derive queryEdgeDeployment from edge URL: https://eu-central-1.aws.edge.axiom.co → cloud.eu-central-1.aws
        if [[ "$RESOLVED_URL" == *".edge.axiom.co" ]]; then
            QUERY_EDGE_DEPLOYMENT="cloud.$(echo "$RESOLVED_URL" | sed 's|https://||;s|\.edge\.axiom\.co||')"
        fi
    fi
fi

JQ_ARGS=(
    --arg apl "$MPL"
    --arg startT "$START_TIME"
    --arg endT "$END_TIME"
)
JQ_EXPR='{apl: $apl, startTime: $startT, endTime: $endT}'

if [[ -n "$QUERY_EDGE_DEPLOYMENT" ]]; then
    JQ_ARGS+=(--arg edgeDeployment "$QUERY_EDGE_DEPLOYMENT")
    JQ_EXPR="$JQ_EXPR + {queryEdgeDeployment: \$edgeDeployment}"
fi

if [[ ${#PARAM_NAMES[@]} -gt 0 ]]; then
    PARAMS_EXPR=""
    for i in "${!PARAM_NAMES[@]}"; do
        JQ_ARGS+=(--arg "pn$i" "${PARAM_NAMES[$i]}" --arg "pv$i" "${PARAM_VALUES[$i]}")
        if [[ -n "$PARAMS_EXPR" ]]; then PARAMS_EXPR+=" + "; fi
        PARAMS_EXPR+="{(\"param__\" + \$pn$i): \$pv$i}"
    done
    JQ_EXPR="$JQ_EXPR + {params: ($PARAMS_EXPR)}"
fi

# Forward chart-sizing hints under queryOptions. The edge translates these into
# the x-axiom-chart-width / x-axiom-pixel-per-point headers, which the metrics
# service uses to resolve $__interval. Values are JSON numbers (--argjson).
if [[ -n "$CHART_WIDTH" || -n "$PIXEL_PER_POINT" ]]; then
    QO_EXPR=""
    if [[ -n "$CHART_WIDTH" ]]; then
        JQ_ARGS+=(--argjson chartWidth "$CHART_WIDTH")
        QO_EXPR="{\"chart-width\": \$chartWidth}"
    fi
    if [[ -n "$PIXEL_PER_POINT" ]]; then
        JQ_ARGS+=(--argjson pixelPerPoint "$PIXEL_PER_POINT")
        if [[ -n "$QO_EXPR" ]]; then QO_EXPR+=" + "; fi
        QO_EXPR+="{\"pixel-per-point\": \$pixelPerPoint}"
    fi
    JQ_EXPR="$JQ_EXPR + {queryOptions: ($QO_EXPR)}"
fi

BODY=$(jq -n "${JQ_ARGS[@]}" "$JQ_EXPR")

AXIOM_ACCEPT="application/json+metrics.v2" "$SCRIPT_DIR/axiom-api" "$DEPLOYMENT" POST "/v1/query/_mpl" "$BODY"
