#!/bin/bash
# O11Y Data Access Toolkit
# Provides unified authentication, S3 exploration, and API queries.
#
# Usage: source o11y-data.sh
#
# Functions:
#   o11y_auth              - Authenticate and export credentials
#   o11y_cluster_info      - Query cluster info from GS API
#   o11y_resolve_metrics_targets - Resolve app + infra O11Y UUIDs for a cluster
#   o11y_metrics_query     - Run Prometheus query against a cluster
#   o11y_metrics_instant_query - Run instant Prometheus query against a cluster
#   o11y_s3_list           - List S3 objects under a prefix
#   o11y_s3_download       - Download a file from S3
#   o11y_log_stats         - Download and display log-stats SQLite
#   o11y_list_logs         - List available log files for a cluster
#   o11y_download_logs     - Download merged logs
#   o11y_s3_bucket         - Get bucket name for a region

if [ -n "${BASH_SOURCE[0]:-}" ]; then
    _O11Y_DATA_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
else
    _O11Y_DATA_ROOT="$(cd "$(dirname "$0")" && pwd)"
fi

_O11Y_INSPECT_SCRIPT="${_O11Y_DATA_ROOT}/o11y-inspect.py"

ENV_FILE="${ENV_FILE:-.env}"

# DS_AUTH_URL is used only in direct mode: the proxy performs the /client/auth
# exchange itself and does not proxy it (a proxied POST /o11y/ds/client/auth is
# rejected), so it stays the real host regardless of mode.
DS_AUTH_URL="https://www.ds.us-west-2.aws.observability.tidbcloud.com/client/auth"
if [[ -n "${DS_PROXY_URL:-}" ]]; then
    GS_API_URL="${DS_PROXY_URL}/o11y/gs/api/v1/clusters/"
else
    GS_API_URL="http://www.gs.us-west-2.aws.observability.tidbcloud.com/api/v1/clusters/"
fi

# o11y_ds_base <region> <vendor> -- mode-aware base URL for the o11y datasource
# (metrics, logs). Proxy mode routes through datasource-proxy, which injects the
# auth0 token and strips any client Authorization header, so queries keep working
# with an empty O11Y_AUTH0_TOKEN. Direct mode targets the real regional host.
o11y_ds_base() {
    local region="${1:-us-west-2}" vendor="${2:-aws}"
    if [[ -n "${DS_PROXY_URL:-}" ]]; then
        echo "${DS_PROXY_URL}/o11y/ds/${region}/${vendor}"
    else
        echo "https://www.ds.${region}.${vendor}.observability.tidbcloud.com"
    fi
}

# o11y_gs_base -- mode-aware base URL for the o11y global server (cluster lookup,
# scaling). gs is always us-west-2/aws.
o11y_gs_base() {
    if [[ -n "${DS_PROXY_URL:-}" ]]; then
        echo "${DS_PROXY_URL}/o11y/gs"
    else
        echo "http://www.gs.us-west-2.aws.observability.tidbcloud.com"
    fi
}

o11y_ds_metrics_base() {
    echo "$(o11y_ds_base "${1:-us-west-2}" "${2:-aws}")/internal/metrics"
}

o11y_auth_request() {
    local provider="${1:?Usage: o11y_auth_request <provider>}"
    local response=""
    local body=""
    local http_code=""
    local attempt=""
    local delay=1

    case "$provider" in
    aws|gcp|alicloud) ;;
    azure)
        echo "ERROR: Azure provider is not yet supported for credential retrieval." >&2
        return 1
        ;;
    *)
        echo "ERROR: Unknown provider '${provider}'. Must be aws, gcp, or alicloud." >&2
        return 1
        ;;
    esac

    for attempt in 1 2 3; do
        response=$(curl -s -w '\n%{http_code}' -X POST "$DS_AUTH_URL" \
            -H 'Content-Type: application/json' \
            -d "{\"github_token\": \"${GITHUB_PAT}\", \"provider\": \"${provider}\"}")
        http_code="${response##*$'\n'}"
        body="${response%$'\n'*}"

        if [ "$http_code" = "200" ]; then
            printf '%s\n' "$body"
            return 0
        fi

        case "$http_code" in
        000|403|429|5??)
            if [ "$attempt" -lt 3 ]; then
                echo "WARN: /client/auth attempt ${attempt} failed with HTTP ${http_code}; retrying..." >&2
                sleep "$delay"
                delay=$((delay * 2))
                continue
            fi
            ;;
        esac

        echo "ERROR: /client/auth failed with HTTP ${http_code}" >&2
        printf '%s\n' "$body" >&2
        return 1
    done

    echo "ERROR: /client/auth failed after 3 attempts" >&2
    printf '%s\n' "$body" >&2
    return 1
}

o11y_gs_query() {
    local vendor="${1:?Usage: o11y_gs_query <vendor> <region> [key=value ...]}"
    local region="${2:?Usage: o11y_gs_query <vendor> <region> [key=value ...]}"
    shift 2

    local args=(
        -G
        --data-urlencode "vendor=${vendor}"
        --data-urlencode "region=${region}"
    )
    local headers=()
    local param=""
    local query_biz_type=""
    for param in "$@"; do
        if [[ "$param" == bizType=* ]]; then
            query_biz_type="${param#bizType=}"
        fi
        [ -n "$param" ] && args+=(--data-urlencode "$param")
    done

    if [ -n "${O11Y_GS_ORG_ID:-}" ]; then
        headers+=(-H "x-org-id: ${O11Y_GS_ORG_ID}")
    fi

    local header_biz_type="${O11Y_GS_BIZ_TYPE:-$query_biz_type}"
    if [ -n "$header_biz_type" ]; then
        headers+=(-H "x-biz-type: ${header_biz_type}")
    fi

    local header_region_id="${O11Y_GS_REGION_ID:-${vendor}-${region}}"
    if [ -n "$header_region_id" ]; then
        headers+=(-H "x-region-id: ${header_region_id}")
    fi

    curl -s "$GS_API_URL" \
        -H "Authorization: Bearer ${O11Y_AUTH0_TOKEN}" \
        "${headers[@]}" \
        "${args[@]}"
}

o11y_restore_gs_context() {
    local had_org_id="${1:-0}"
    local old_org_id="${2-}"
    local had_region_id="${3:-0}"
    local old_region_id="${4-}"
    local had_biz_type="${5:-0}"
    local old_biz_type="${6-}"

    if [ "$had_org_id" = "1" ]; then
        export O11Y_GS_ORG_ID="$old_org_id"
    else
        unset O11Y_GS_ORG_ID
    fi

    if [ "$had_region_id" = "1" ]; then
        export O11Y_GS_REGION_ID="$old_region_id"
    else
        unset O11Y_GS_REGION_ID
    fi

    if [ "$had_biz_type" = "1" ]; then
        export O11Y_GS_BIZ_TYPE="$old_biz_type"
    else
        unset O11Y_GS_BIZ_TYPE
    fi
}

o11y_try_clinic_cluster_context() {
    local cluster_id="${1:?Usage: o11y_try_clinic_cluster_context <cluster_id>}"
    local python_bin="${PYTHON_BIN:-python3}"

    if ! command -v "$python_bin" >/dev/null 2>&1 || [ ! -f "$_O11Y_INSPECT_SCRIPT" ]; then
        echo '{}'
        return 0
    fi

    "$python_bin" "$_O11Y_INSPECT_SCRIPT" "$cluster_id" --clinic-context 2>/dev/null || echo '{}'
}

o11y_try_python_metrics_targets() {
    local cluster_id="${1:?Usage: o11y_try_python_metrics_targets <cluster_id> [vendor] [region] [time_ts]}"
    local vendor="${2:-aws}"
    local region="${3:-}"
    local time_ts="${4:-$(date +%s)}"
    local python_bin="${PYTHON_BIN:-python3}"

    if ! command -v "$python_bin" >/dev/null 2>&1 || [ ! -f "$_O11Y_INSPECT_SCRIPT" ]; then
        echo '{}'
        return 0
    fi

    local args=("$python_bin" "$_O11Y_INSPECT_SCRIPT" "$cluster_id" --provider "$vendor" --resolve-targets --time-ts "$time_ts")
    if [ -n "$region" ]; then
        args+=(--region "$region")
    fi

    "${args[@]}" 2>/dev/null || echo '{}'
}

o11y_probe_metrics_target() {
    local o11y_id="${1:?Usage: o11y_probe_metrics_target <o11y_id> [region] [vendor] [metrics_base]}"
    local region="${2:-us-west-2}"
    local vendor="${3:-aws}"
    local metrics_base="${4:-$(o11y_ds_metrics_base "$region" "$vendor")}"
    local http_code=""

    http_code=$(curl -s -o /dev/null -w '%{http_code}' \
        "${metrics_base}/${o11y_id}/api/v1/query" \
        -G \
        --data-urlencode "query=vector(1)" \
        --data-urlencode "time=$(date +%s)" \
        -H "Authorization: Bearer ${O11Y_AUTH0_TOKEN}")

    [ "$http_code" = "200" ]
}

o11y_extract_metrics_uuid() {
    local cluster_json="${1:?Usage: o11y_extract_metrics_uuid <cluster_json>}"

    echo "$cluster_json" | jq -r '
        (
            try ((.internalReadUri // "") | capture("/metrics/(?<id>[^/]+)/api/v1/").id)
            catch empty
        ) // (.id // empty)
    '
}

o11y_extract_metrics_base() {
    local cluster_json="${1:?Usage: o11y_extract_metrics_base <cluster_json>}"

    echo "$cluster_json" | jq -r '
        (
            try (
                (.internalReadUri // "")
                | capture("^(?<base>https?://.*/metrics)/[^/]+/api/v1/?$").base
            )
            catch empty
        ) // empty
    '
}

o11y_pick_metrics_uuid() {
    local cluster_json="${1:?Usage: o11y_pick_metrics_uuid <cluster_json> [region] [vendor]}"
    local region="${2:-us-west-2}"
    local vendor="${3:-aws}"
    local primary=""
    local fallback=""
    local candidate=""
    local metrics_base=""

    primary=$(o11y_extract_metrics_uuid "$cluster_json")
    fallback=$(echo "$cluster_json" | jq -r '.id // empty')
    metrics_base=$(o11y_extract_metrics_base "$cluster_json")

    for candidate in "$primary" "$fallback"; do
        [ -z "$candidate" ] && continue
        if o11y_probe_metrics_target "$candidate" "$region" "$vendor" "$metrics_base"; then
            printf '%s\n' "$candidate"
            return 0
        fi
    done

    [ -n "$primary" ] && { printf '%s\n' "$primary"; return 0; }
    [ -n "$fallback" ] && { printf '%s\n' "$fallback"; return 0; }
    return 1
}

o11y_find_cluster_record() {
    local vendor="${1:?Usage: o11y_find_cluster_record <vendor> <region> <cluster_id>}"
    local region="${2:?Usage: o11y_find_cluster_record <vendor> <region> <cluster_id>}"
    local cluster_id="${3:?Usage: o11y_find_cluster_record <vendor> <region> <cluster_id>}"
    local biz_types=(
        tidbcloud
        staging-tidbcloud
        dev-tidbcloud
        prod-tidbcloud-nextgen
        staging-tidbcloud-nextgen
        dev-tidbcloud-nextgen
        tidbcloud-nextgen
        prod-devtier
        staging-devtier
        dev-devtier
        devtier
    )
    local biz_type=""
    local resp=""
    local hit=""

    for biz_type in "${biz_types[@]}"; do
        resp=$(o11y_gs_query "$vendor" "$region" \
            "bizType=${biz_type}" \
            "applicationID=${cluster_id}")
        hit=$(echo "$resp" | jq '[
            .data.clusters[]?
            | select(.applicationID != "k8s-infra")
            | select((.projectID // "") != "seed" and (.projectID // "") != "base")
        ] | .[0]')
        if [ "$(echo "$hit" | jq -r '.id // empty')" != "" ]; then
            printf '%s\n' "$hit"
            return 0
        fi
    done

    return 1
}

o11y_resolve_serverless_keyspace() {
    local app_uuid="${1:?Usage: o11y_resolve_serverless_keyspace <app_uuid> <cluster_id> [region] [vendor]}"
    local cluster_id="${2:?Usage: o11y_resolve_serverless_keyspace <app_uuid> <cluster_id> [region] [vendor]}"
    local region="${3:-us-west-2}"
    local vendor="${4:-aws}"
    local query=""
    local keyspace_id=""

    for query in \
        "sum by (keyspace_id) (0 * gateway_connection_flow_bytes_total{cluster_id=\"${cluster_id}\"})" \
        "sum by (keyspace_id) (0 * tidb_server_connections{cluster_id=\"${cluster_id}\"})"; do
        keyspace_id=$(o11y_metrics_instant_query "$app_uuid" "$query" "" "$region" "$vendor" | jq -r '.data.result[0].metric.keyspace_id // empty')
        if [ -n "$keyspace_id" ]; then
            printf '%s\n' "$keyspace_id"
            return 0
        fi
    done

    return 1
}

o11y_resolve_serverless_tikv_hosts() {
    local app_uuid="${1:?Usage: o11y_resolve_serverless_tikv_hosts <app_uuid> <keyspace_id> [time_ts] [region] [vendor]}"
    local keyspace_id="${2:?Usage: o11y_resolve_serverless_tikv_hosts <app_uuid> <keyspace_id> [time_ts] [region] [vendor]}"
    local ts="${3:-$(date +%s)}"
    local region="${4:-us-west-2}"
    local vendor="${5:-aws}"
    local instances_json='[]'
    local nodes_json='{}'
    local instance=""
    local node=""

    instances_json=$(o11y_metrics_instant_query "$app_uuid" \
        "sum(tikv_store_size_bytes{keyspace_id=\"${keyspace_id}\", type=\"used\"}) by (instance)" \
        "$ts" "$region" "$vendor" | jq '[
            .data.result[]?
            | select((.value[1] | tonumber) > 0)
            | .metric.instance
        ] | unique')

    while IFS= read -r instance; do
        [ -z "$instance" ] && continue
        node=$(o11y_metrics_instant_query "$app_uuid" \
            "kube_pod_info{namespace=\"tidb-serverless\", pod=\"${instance}\"}" \
            "$ts" "$region" "$vendor" | jq -r '.data.result[0].metric.node // empty')
        if [ -n "$node" ]; then
            nodes_json=$(echo "$nodes_json" | jq --arg instance "$instance" --arg node "$node" '. + {($instance): $node}')
        fi
    done < <(echo "$instances_json" | jq -r '.[]')

    jq -n \
        --argjson instances "$instances_json" \
        --argjson nodes "$nodes_json" \
        '{
            instances: $instances,
            nodes: $nodes
        }'
}

# ────────────────────────────────────────────────────────────
# Authentication
# ────────────────────────────────────────────────────────────
o11y_auth() {
    local provider="${1:-}"

    if [[ -n "${DS_PROXY_URL:-}" ]]; then
        # The proxy injects the auth0 token for metrics/logs/cluster queries, so
        # the caller holds no query token. Fetch only the direct-flow creds the
        # proxy does not proxy: O11Y_TOKEN (export) and AWS STS (S3), served at
        # /o11y/credentials when O11Y_EXPOSE_CREDENTIALS is enabled on the proxy.
        export O11Y_AUTH0_TOKEN="" O11Y_PROVIDER="${provider:-aws}"
        local _resp _code
        _resp=$(curl -s -w '\n%{http_code}' "${DS_PROXY_URL}/o11y/credentials")
        _code="${_resp##*$'\n'}"
        _resp="${_resp%$'\n'*}"
        if [ "$_code" = "200" ]; then
            export O11Y_TOKEN=$(echo "$_resp" | jq -r '.o11y_token')
            export AWS_ACCESS_KEY_ID=$(echo "$_resp" | jq -r '.access_key_id // empty')
            export AWS_SECRET_ACCESS_KEY=$(echo "$_resp" | jq -r '.secret_access_key // empty')
            export AWS_SESSION_TOKEN=$(echo "$_resp" | jq -r '.session_token // empty')
            echo "O11Y proxy mode: queries via proxy; export/S3 creds fetched from /o11y/credentials" >&2
        else
            export O11Y_TOKEN=""
            echo "O11Y proxy mode: queries via proxy. /o11y/credentials -> HTTP ${_code}; export/S3 unavailable (enable O11Y_EXPOSE_CREDENTIALS on the proxy)" >&2
        fi
        return 0
    fi

    if [ -f "$ENV_FILE" ]; then
        source "$ENV_FILE"
    fi

    if [ -z "$GITHUB_PAT" ]; then
        if [ -f "$HOME/.observability/.github" ]; then
            GITHUB_PAT=$(tr -d '[:space:]' < "$HOME/.observability/.github")
        fi
    fi

    if [ -z "$GITHUB_PAT" ]; then
        echo "ERROR: GITHUB_PAT not set. Please provide your GitHub Personal Access Token." >&2
        echo "Set it in .env, ~/.observability/.github, or environment." >&2
        return 1
    fi

    if [ -z "$provider" ]; then
        echo "ERROR: Provider is required. Usage: o11y_auth <aws|gcp|alicloud>" >&2
        return 1
    fi

    case "$provider" in
    aws|gcp|alicloud) ;;
    azure)
        echo "ERROR: Azure provider is not yet supported for credential retrieval." >&2
        return 1
        ;;
    *)
        echo "ERROR: Unknown provider '${provider}'. Must be aws, gcp, or alicloud." >&2
        return 1
        ;;
    esac

    local now=$(date +%s)
    local token_expires="${O11Y_TOKEN_EXPIRES:-$O11Y_AUTH0_TOKEN_EXPIRES}"
    local cached_provider="${O11Y_PROVIDER:-}"
    if [ -n "$O11Y_AUTH0_TOKEN" ] && [ -n "$token_expires" ] && [ "$cached_provider" = "$provider" ]; then
        local buf=$((token_expires - 300))
        if [ "$now" -lt "$buf" ]; then
            echo "Using cached ${provider} credentials (expires in $((token_expires - now))s)" >&2
            export O11Y_TOKEN O11Y_AUTH0_TOKEN O11Y_PROVIDER
            case "$provider" in
            aws)
                export AWS_ACCESS_KEY_ID="$O11Y_AWS_ACCESS_KEY_ID" \
                       AWS_SECRET_ACCESS_KEY="$O11Y_AWS_SECRET_ACCESS_KEY" \
                       AWS_SESSION_TOKEN="$O11Y_AWS_SESSION_TOKEN"
                ;;
            gcp)
                export O11Y_GCP_ACCESS_TOKEN="$O11Y_GCP_ACCESS_TOKEN"
                ;;
            alicloud)
                export O11Y_ALI_ACCESS_KEY_ID="$O11Y_ALI_ACCESS_KEY_ID" \
                       O11Y_ALI_SECRET_ACCESS_KEY="$O11Y_ALI_SECRET_ACCESS_KEY" \
                       O11Y_ALI_SECURITY_TOKEN="$O11Y_ALI_SECURITY_TOKEN"
                ;;
            esac
            return 0
        fi
    fi

    local resp
    resp=$(o11y_auth_request "$provider") || return 1

    local token=$(echo "$resp" | jq -r '.token')
    local auth0=$(echo "$resp" | jq -r '.auth0_token')

    if [ -z "$token" ] || [ "$token" = "null" ] || [ -z "$auth0" ] || [ "$auth0" = "null" ]; then
        echo "ERROR: Authentication failed" >&2
        echo "$resp" >&2
        return 1
    fi

    local exp=$((now + 43200))

    grep -v "^O11Y_TOKEN=" "$ENV_FILE" 2>/dev/null |
        grep -v "^O11Y_AUTH0_TOKEN=" |
        grep -v "^O11Y_AUTH0_TOKEN_EXPIRES=" |
        grep -v "^O11Y_TOKEN_EXPIRES=" |
        grep -v "^O11Y_PROVIDER=" |
        grep -v "^O11Y_AWS_ACCESS_KEY_ID=" |
        grep -v "^O11Y_AWS_SECRET_ACCESS_KEY=" |
        grep -v "^O11Y_AWS_SESSION_TOKEN=" |
        grep -v "^O11Y_GCP_ACCESS_TOKEN=" |
        grep -v "^O11Y_ALI_ACCESS_KEY_ID=" |
        grep -v "^O11Y_ALI_SECRET_ACCESS_KEY=" |
        grep -v "^O11Y_ALI_SECURITY_TOKEN=" \
        >"${ENV_FILE}.tmp" 2>/dev/null

    cat >>"${ENV_FILE}.tmp" <<EOF
O11Y_TOKEN=${token}
O11Y_AUTH0_TOKEN=${auth0}
O11Y_AUTH0_TOKEN_EXPIRES=${exp}
O11Y_TOKEN_EXPIRES=${exp}
O11Y_PROVIDER=${provider}
EOF

    case "$provider" in
    aws)
        local ak=$(echo "$resp" | jq -r '.access_key_id')
        local sk=$(echo "$resp" | jq -r '.secret_access_key')
        local st=$(echo "$resp" | jq -r '.session_token')
        cat >>"${ENV_FILE}.tmp" <<EOF
O11Y_AWS_ACCESS_KEY_ID=${ak}
O11Y_AWS_SECRET_ACCESS_KEY=${sk}
O11Y_AWS_SESSION_TOKEN=${st}
EOF
        export AWS_ACCESS_KEY_ID="$ak"
        export AWS_SECRET_ACCESS_KEY="$sk"
        export AWS_SESSION_TOKEN="$st"
        export AWS_DEFAULT_REGION="${AWS_DEFAULT_REGION:-us-west-2}"
        ;;
    gcp)
        local gcp_token=$(echo "$resp" | jq -r '.access_token')
        cat >>"${ENV_FILE}.tmp" <<EOF
O11Y_GCP_ACCESS_TOKEN=${gcp_token}
EOF
        export O11Y_GCP_ACCESS_TOKEN="$gcp_token"
        ;;
    alicloud)
        local ali_ak=$(echo "$resp" | jq -r '.access_key_id')
        local ali_sk=$(echo "$resp" | jq -r '.secret_access_key')
        local ali_st=$(echo "$resp" | jq -r '.security_token')
        cat >>"${ENV_FILE}.tmp" <<EOF
O11Y_ALI_ACCESS_KEY_ID=${ali_ak}
O11Y_ALI_SECRET_ACCESS_KEY=${ali_sk}
O11Y_ALI_SECURITY_TOKEN=${ali_st}
EOF
        export O11Y_ALI_ACCESS_KEY_ID="$ali_ak"
        export O11Y_ALI_SECRET_ACCESS_KEY="$ali_sk"
        export O11Y_ALI_SECURITY_TOKEN="$ali_st"
        ;;
    esac

    mv "${ENV_FILE}.tmp" "$ENV_FILE"

    export O11Y_TOKEN="$token"
    export O11Y_AUTH0_TOKEN="$auth0"
    export O11Y_PROVIDER="$provider"

    echo "Authenticated successfully with ${provider} (valid for 12h)" >&2
}

# ────────────────────────────────────────────────────────────
# Cluster Info
# ────────────────────────────────────────────────────────────
o11y_cluster_info() {
    local vendor="${1:?Usage: o11y_cluster_info <vendor> <region> [tenantID] [clusterID] [bizType]}"
    local region="${2:?Usage: o11y_cluster_info <vendor> <region> [tenantID] [clusterID] [bizType]}"
    local tenant_id="$3"
    local cluster_id="$4"
    local biz_type="$5"

    if [ -n "$biz_type" ]; then
        o11y_gs_query "$vendor" "$region" "bizType=${biz_type}" \
            "${tenant_id:+tenantID=${tenant_id}}" \
            "${cluster_id:+applicationID=${cluster_id}}" | jq '.data.clusters'
    elif [ -n "$cluster_id" ]; then
        o11y_find_cluster_record "$vendor" "$region" "$cluster_id" | jq -s '.'
    else
        o11y_gs_query "$vendor" "$region" "bizType=tidbcloud" \
            "${tenant_id:+tenantID=${tenant_id}}" | jq '.data.clusters'
    fi
}

o11y_resolve_metrics_targets() {
    local vendor="${1:?Usage: o11y_resolve_metrics_targets <vendor> <region> <clusterID> [namespace] [time_ts]}"
    local region="${2:?Usage: o11y_resolve_metrics_targets <vendor> <region> <clusterID> [namespace] [time_ts]}"
    local cluster_id="${3:?Usage: o11y_resolve_metrics_targets <vendor> <region> <clusterID> [namespace] [time_ts]}"
    local namespace="${4:-}"
    local time_ts="${5:-$(date +%s)}"
    local app_lookup_id="${O11Y_CLUSTER_APP_ID:-$cluster_id}"
    local had_org_id=0
    local old_org_id="${O11Y_GS_ORG_ID:-}"
    local had_region_id=0
    local old_region_id="${O11Y_GS_REGION_ID:-}"
    local had_biz_type=0
    local old_biz_type="${O11Y_GS_BIZ_TYPE:-}"
    local nextgen_org_id=""
    local clinic_context='{}'
    local clinic_gs_biz_type=""

    [ "${O11Y_GS_ORG_ID+x}" = "x" ] && had_org_id=1
    [ "${O11Y_GS_REGION_ID+x}" = "x" ] && had_region_id=1
    [ "${O11Y_GS_BIZ_TYPE+x}" = "x" ] && had_biz_type=1

    if [[ "$namespace" =~ ^[0-9]+$ ]]; then
        time_ts="$namespace"
        namespace=""
    fi

    local app_resp
    app_resp=$(o11y_find_cluster_record "$vendor" "$region" "$app_lookup_id" 2>/dev/null || true)

    local app_json
    app_json=$(printf '%s\n' "${app_resp:-{}}" | jq -c '.' 2>/dev/null || echo '{}')
    clinic_context=$(o11y_try_clinic_cluster_context "$cluster_id")

    local clinic_deploy_type=""
    clinic_deploy_type=$(echo "$clinic_context" | jq -r '.deploy_type // empty' 2>/dev/null || true)
    if [[ "$clinic_deploy_type" == "premium" || "$clinic_deploy_type" == "resource_pool" ]]; then
        local python_targets='{}'
        python_targets=$(o11y_try_python_metrics_targets "$cluster_id" "$vendor" "$region" "$time_ts")
        if [ "$(echo "$python_targets" | jq -r '.app_uuid // empty' 2>/dev/null || true)" != "" ]; then
            o11y_restore_gs_context "$had_org_id" "$old_org_id" "$had_region_id" "$old_region_id" "$had_biz_type" "$old_biz_type"
            printf '%s\n' "$python_targets"
            return 0
        fi
    fi

    if [ "$(echo "$app_json" | jq -r '.id // empty')" = "" ]; then
        if [ "$(echo "$clinic_context" | jq -r '.deploy_type // empty')" = "premium" ] || [ "$(echo "$clinic_context" | jq -r '.deploy_type // empty')" = "resource_pool" ]; then
            vendor=$(echo "$clinic_context" | jq -r --arg fallback "$vendor" '(.vendor // "") as $value | if $value == "" then $fallback else $value end')
            region=$(echo "$clinic_context" | jq -r --arg fallback "$region" '(.region // "") as $value | if $value == "" then $fallback else $value end')
            nextgen_org_id=$(echo "$clinic_context" | jq -r '.tenant_id // empty')
            if [ -z "$nextgen_org_id" ]; then
                nextgen_org_id=$(echo "$clinic_context" | jq -r '.org_id // empty')
            fi
            if [ -n "$nextgen_org_id" ]; then
                export O11Y_GS_ORG_ID="$nextgen_org_id"
            else
                unset O11Y_GS_ORG_ID
            fi
            export O11Y_GS_REGION_ID="${vendor}-${region}"
            clinic_gs_biz_type=$(echo "$clinic_context" | jq -r '.gs_biz_type // empty')
            if [ -n "$clinic_gs_biz_type" ]; then
                export O11Y_GS_BIZ_TYPE="$clinic_gs_biz_type"
            else
                unset O11Y_GS_BIZ_TYPE
            fi
            app_resp=$(o11y_find_cluster_record "$vendor" "$region" "$app_lookup_id" 2>/dev/null || true)
            app_json=$(printf '%s\n' "${app_resp:-{}}" | jq -c '.' 2>/dev/null || echo '{}')
        fi
    fi

    local app_uuid
    app_uuid=$(o11y_pick_metrics_uuid "$app_json" "$region" "$vendor")
    if [ -z "$app_uuid" ]; then
        o11y_restore_gs_context "$had_org_id" "$old_org_id" "$had_region_id" "$old_region_id" "$had_biz_type" "$old_biz_type"
        echo "ERROR: Failed to resolve application metrics UUID for cluster ${cluster_id}" >&2
        echo "$app_resp" >&2
        return 1
    fi
    local app_metrics_base
    app_metrics_base=$(o11y_extract_metrics_base "$app_json")
    [ -z "$app_metrics_base" ] && app_metrics_base=$(o11y_ds_metrics_base "$region" "$vendor")

    local tenant_id
    tenant_id=$(echo "$app_json" | jq -r '.tenantID // empty')
    local org_id
    org_id=$(echo "$app_json" | jq -r '.orgID // empty')
    local project_id
    project_id=$(echo "$app_json" | jq -r '.projectID // empty')
    local biz_type
    biz_type=$(echo "$app_json" | jq -r '.bizType // empty')
    local deploy_type
    deploy_type=$(echo "$app_json" | jq -r '.deployType // empty')
    local parent_id
    parent_id=$(echo "$app_json" | jq -r '.parentID // empty')
    local data_plane_name
    data_plane_name=$(echo "$app_json" | jq -r '.dataPlaneName // empty')
    local infra_json='{}'
    local keyspace_id=""
    local serverless_hosts='{"instances":[],"nodes":{}}'

    [ -z "$tenant_id" ] && tenant_id=$(echo "$clinic_context" | jq -r '.tenant_id // empty')
    [ -z "$org_id" ] && org_id=$(echo "$clinic_context" | jq -r '.org_id // empty')
    [ -z "$deploy_type" ] && deploy_type=$(echo "$clinic_context" | jq -r '.deploy_type // empty')
    [ -z "$parent_id" ] && parent_id=$(echo "$clinic_context" | jq -r '.parent_id // empty')
    [ -z "$data_plane_name" ] && data_plane_name=$(echo "$clinic_context" | jq -r '.data_plane_name // empty')
    if [ -z "$biz_type" ] && [ -n "$(echo "$clinic_context" | jq -r '.gs_biz_type // empty')" ]; then
        biz_type=$(echo "$clinic_context" | jq -r '.gs_biz_type // empty')
    fi

    if [ -z "$namespace" ]; then
        if [ -n "${O11Y_CLUSTER_NAMESPACE:-}" ]; then
            namespace="${O11Y_CLUSTER_NAMESPACE}"
        elif [[ "$biz_type" == *"devtier"* ]] || [ "$deploy_type" = "starter" ] || [ "$deploy_type" = "essential" ]; then
            namespace="tidb-serverless"
        elif [[ "$biz_type" == *"nextgen"* ]] || [ "$deploy_type" = "premium" ] || [ "$deploy_type" = "resource_pool" ]; then
            if [ -n "$parent_id" ]; then
                namespace="tidb${parent_id}"
            else
                namespace="tidb${cluster_id}"
            fi
        else
            namespace="tidb${cluster_id}"
        fi
    fi

    if [[ "$biz_type" == *"nextgen"* ]] || [ "$deploy_type" = "premium" ] || [ "$deploy_type" = "resource_pool" ]; then
        nextgen_org_id="${tenant_id:-$org_id}"
        if [ -n "$nextgen_org_id" ]; then
            export O11Y_GS_ORG_ID="$nextgen_org_id"
        else
            unset O11Y_GS_ORG_ID
        fi
        export O11Y_GS_REGION_ID="${vendor}-${region}"
        if [ -n "$biz_type" ]; then
            export O11Y_GS_BIZ_TYPE="$biz_type"
        else
            unset O11Y_GS_BIZ_TYPE
        fi
    fi

    if [ "$biz_type" = "tidbcloud" ] && [ -n "$project_id" ]; then
        infra_json=$(o11y_gs_query "$vendor" "$region" \
            "bizType=${biz_type}" \
            "applicationID=k8s-infra" \
            "projectID=${project_id}" | jq --arg project_id "$project_id" '[
                .data.clusters[]?
                | select(.applicationID == "k8s-infra" and .projectID == $project_id)
            ] | .[0]')
    fi

    local infra_uuid
    infra_uuid=""
    local infra_metrics_base=""
    if [ "$(echo "$infra_json" | jq -r 'if . == null then "" else .id // empty end')" != "" ]; then
        infra_uuid=$(o11y_pick_metrics_uuid "$infra_json" "$region" "$vendor")
        infra_metrics_base=$(o11y_extract_metrics_base "$infra_json")
    fi
    local dataplane_lookup_id="${O11Y_CLUSTER_DATAPLANE_NAME:-$data_plane_name}"
    if [ -z "$infra_uuid" ] && [ -n "$dataplane_lookup_id" ]; then
        infra_json=$(o11y_find_cluster_record "$vendor" "$region" "$dataplane_lookup_id" 2>/dev/null || true)
        if [ "$(echo "$infra_json" | jq -r '.id // empty')" != "" ]; then
            infra_uuid=$(o11y_pick_metrics_uuid "$infra_json" "$region" "$vendor")
            infra_metrics_base=$(o11y_extract_metrics_base "$infra_json")
        fi
    fi
    if [ -z "$infra_uuid" ] && [ -n "$tenant_id" ]; then
        local infra_candidates
        infra_candidates=$(o11y_gs_query "$vendor" "$region" \
            "${biz_type:+bizType=${biz_type}}" \
            "tenantID=${tenant_id}" | jq '[
                .data.clusters[]?
                | select((.applicationID // "") | startswith("dp-"))
            ]')

        local candidate_count
        candidate_count=$(echo "$infra_candidates" | jq 'length')
        if [ "$candidate_count" -eq 1 ]; then
            infra_json=$(echo "$infra_candidates" | jq '.[0]')
            infra_uuid=$(o11y_pick_metrics_uuid "$infra_json" "$region" "$vendor")
        elif [ "$candidate_count" -gt 1 ]; then
            local idx=""
            local candidate_uuid=""
            local match_count=""
            for idx in $(seq 0 $((candidate_count - 1))); do
                candidate_uuid=$(o11y_pick_metrics_uuid "$(echo "$infra_candidates" | jq ".[$idx]")" "$region" "$vendor")
                [ -z "$candidate_uuid" ] && continue
                match_count=$(o11y_metrics_instant_query "$candidate_uuid" \
                    "count(kube_pod_info{namespace=\"${namespace}\"})" \
                    "" "$region" "$vendor" | jq -r '
                        try (
                            if (.data.result | length) > 0
                            then (.data.result[0].value[1] | tonumber | floor)
                            else 0
                            end
                        ) catch 0')
                case "$match_count" in
                    ''|*[!0-9]*)
                        match_count=0
                        ;;
                esac
                if [ "$match_count" -gt 0 ]; then
                    infra_json=$(echo "$infra_candidates" | jq ".[$idx]")
                    infra_uuid="$candidate_uuid"
                    infra_metrics_base=$(o11y_extract_metrics_base "$infra_json")
                    break
                fi
            done
        fi
    fi
    if [ -n "$infra_uuid" ] && [ -z "$infra_metrics_base" ]; then
        infra_metrics_base=$(o11y_ds_metrics_base "$region" "$vendor")
    fi

    if [[ "$biz_type" == *"devtier"* ]] || [ "$deploy_type" = "starter" ] || [ "$deploy_type" = "essential" ]; then
        keyspace_id=$(o11y_resolve_serverless_keyspace "$app_uuid" "$cluster_id" "$region" "$vendor" 2>/dev/null || true)
        if [ -n "$keyspace_id" ]; then
            serverless_hosts=$(o11y_resolve_serverless_tikv_hosts "$app_uuid" "$keyspace_id" "$time_ts" "$region" "$vendor")
        fi
    fi

    local resolved_json
    resolved_json=$(jq -n \
        --arg cluster_id "$cluster_id" \
        --arg namespace "$namespace" \
        --arg vendor "$vendor" \
        --arg region "$region" \
        --arg app_uuid "$app_uuid" \
        --arg app_application_id "$(echo "$app_json" | jq -r '.applicationID // empty')" \
        --arg tenant_id "$tenant_id" \
        --arg project_id "$project_id" \
        --arg biz_type "$biz_type" \
        --arg deploy_type "$deploy_type" \
        --arg infra_uuid "$infra_uuid" \
        --arg infra_application_id "$(echo "$infra_json" | jq -r '.applicationID // empty')" \
        --arg app_metrics_base "$app_metrics_base" \
        --arg infra_metrics_base "$infra_metrics_base" \
        --arg keyspace_id "$keyspace_id" \
        --argjson serverless_hosts "$serverless_hosts" \
        '{
            cluster_id: $cluster_id,
            namespace: $namespace,
            vendor: $vendor,
            region: $region,
            biz_type: $biz_type,
            deploy_type: $deploy_type,
            tenant_id: $tenant_id,
            project_id: $project_id,
            app_uuid: $app_uuid,
            app_application_id: $app_application_id,
            app_metrics_base: $app_metrics_base,
            infra_uuid: $infra_uuid,
            infra_application_id: $infra_application_id,
            infra_metrics_base: $infra_metrics_base,
            keyspace_id: $keyspace_id,
            tikv_instances: ($serverless_hosts.instances // []),
            tikv_nodes: ($serverless_hosts.nodes // {})
        }')

    o11y_restore_gs_context "$had_org_id" "$old_org_id" "$had_region_id" "$old_region_id" "$had_biz_type" "$old_biz_type"

    printf '%s\n' "$resolved_json"
}

# ────────────────────────────────────────────────────────────
# Metrics Query
# ────────────────────────────────────────────────────────────
o11y_metrics_instant_query() {
    local o11y_id="${1:?Usage: o11y_metrics_instant_query <o11y_id> <promql> [time_ts] [region] [vendor] [metrics_base]}"
    local query="${2:?Usage: o11y_metrics_instant_query <o11y_id> <promql> [time_ts] [region] [vendor] [metrics_base]}"
    local ts="${3:-$(date +%s)}"
    local region="${4:-us-west-2}"
    local vendor="${5:-aws}"
    local metrics_base="${6:-$(o11y_ds_metrics_base "$region" "$vendor")}"

    curl -s "${metrics_base}/${o11y_id}/api/v1/query" \
        -G \
        --data-urlencode "query=${query}" \
        --data-urlencode "time=${ts}" \
        -H "Authorization: Bearer ${O11Y_AUTH0_TOKEN}" | jq '.'
}

o11y_metrics_query() {
    local o11y_id="${1:?Usage: o11y_metrics_query <o11y_id> <promql> [start_ts] [end_ts] [step] [region] [vendor] [metrics_base]}"
    local query="${2:?Usage: o11y_metrics_query <o11y_id> <promql> [start_ts] [end_ts] [step] [region] [vendor] [metrics_base]}"
    local now=$(date +%s)
    local start="${3:-$((now - 3600))}"
    local end="${4:-$now}"
    local step="${5:-60}"
    local region="${6:-us-west-2}"
    local vendor="${7:-aws}"
    local metrics_base="${8:-$(o11y_ds_metrics_base "$region" "$vendor")}"

    curl -s "${metrics_base}/${o11y_id}/api/v1/query_range?query=$(python3 -c "import urllib.parse; print(urllib.parse.quote('$query'))")&start=${start}&end=${end}&step=${step}" \
        -H "Authorization: Bearer ${O11Y_AUTH0_TOKEN}" | jq '.'
}

# ────────────────────────────────────────────────────────────
# Loki Helpers
# ────────────────────────────────────────────────────────────
o11y_loki_base() {
    echo "$(o11y_ds_base "${1:-us-west-2}" "${2:-aws}")/loki/self-monitoring/loki/api/v1"
}

o11y_loki_query() {
    local region="${1:?Usage: o11y_loki_query <region> <vendor> <logql> [start_ts] [end_ts] [limit]}"
    local vendor="${2:?Usage: o11y_loki_query <region> <vendor> <logql> [start_ts] [end_ts] [limit]}"
    local query="${3:?Usage: o11y_loki_query <region> <vendor> <logql> [start_ts] [end_ts] [limit]}"
    local now=$(date +%s)
    local start="${4:-$((now - 3600))}"
    local end="${5:-$now}"
    local limit="${6:-100}"

    curl -s --max-time 30 --globoff \
        "$(o11y_loki_base "$region" "$vendor")/query_range" \
        -G \
        --data-urlencode "query=${query}" \
        --data-urlencode "start=${start}.0" \
        --data-urlencode "end=${end}.0" \
        --data-urlencode "limit=${limit}" \
        -H "Authorization: Bearer ${O11Y_AUTH0_TOKEN}" | jq '.'
}

o11y_loki_labels() {
    local region="${1:-us-west-2}"
    local vendor="${2:-aws}"
    curl -s --max-time 10 \
        "$(o11y_loki_base "$region" "$vendor")/labels" \
        -H "Authorization: Bearer ${O11Y_AUTH0_TOKEN}" | jq '.'
}

o11y_loki_label_values() {
    local label="${1:?Usage: o11y_loki_label_values <label> [region] [vendor]}"
    local region="${2:-us-west-2}"
    local vendor="${3:-aws}"
    curl -s --max-time 10 \
        "$(o11y_loki_base "$region" "$vendor")/label/${label}/values" \
        -H "Authorization: Bearer ${O11Y_AUTH0_TOKEN}" | jq '.'
}

# ────────────────────────────────────────────────────────────
# S3 Helpers
# ────────────────────────────────────────────────────────────
o11y_s3_bucket() {
    local region="${1:-us-west-2}"
    local suffix="${2:-}"
    if [ -n "$suffix" ]; then
        echo "o11y-prod-shared-${region}-${suffix}"
    else
        echo "o11y-prod-shared-${region}"
    fi
}

o11y_s3_list() {
    local prefix="${1:?Usage: o11y_s3_list <prefix> [region] [bucket_suffix]}"
    local region="${2:-us-west-2}"
    local suffix="${3:-}"
    aws s3 ls "s3://$(o11y_s3_bucket "$region" "$suffix")/${prefix}"
}

o11y_s3_download() {
    local key="${1:?Usage: o11y_s3_download <s3_key> <local_path> [region] [bucket_suffix]}"
    local dest="${2:?Usage: o11y_s3_download <s3_key> <local_path> [region] [bucket_suffix]}"
    local region="${3:-us-west-2}"
    local suffix="${4:-}"
    mkdir -p "$(dirname "$dest")"
    aws s3 cp "s3://$(o11y_s3_bucket "$region" "$suffix")/${key}" "$dest"
}

# ────────────────────────────────────────────────────────────
# BizType Resolver
# ────────────────────────────────────────────────────────────
o11y_resolve_biz_type() {
    local deploy_type="$1"
    local env="${2:-}"
    local base=""

    case "$deploy_type" in
    essential)             base="devtier" ;;
    dedicated)             base="tidbcloud" ;;
    starter)               base="devtier" ;;
    premium|resource_pool) base="tidbcloud-nextgen" ;;
    byoc)                  base="tidbcloud" ;;
    *)
        echo "ERROR: Unknown deployType '$deploy_type'. Must be essential/dedicated/starter/premium/byoc." >&2
        return 1
        ;;
    esac

    if [ -z "$env" ]; then
        echo "$base"
    elif [ "$env" = "prod" ] && [ "$base" = "tidbcloud" ]; then
        echo "tidbcloud"
    else
        echo "${env}-${base}"
    fi
}

if [ -n "${BASH_VERSION:-}" ]; then
    export -f o11y_auth o11y_auth_request o11y_cluster_info o11y_resolve_metrics_targets
    export -f o11y_metrics_instant_query o11y_metrics_query
    export -f o11y_loki_base o11y_loki_query o11y_loki_labels o11y_loki_label_values
    export -f o11y_s3_bucket o11y_s3_list o11y_s3_download o11y_resolve_biz_type
    export -f o11y_ds_metrics_base o11y_probe_metrics_target o11y_extract_metrics_uuid o11y_pick_metrics_uuid
    export -f o11y_extract_metrics_base o11y_restore_gs_context o11y_try_clinic_cluster_context
    export -f o11y_try_python_metrics_targets
    export -f o11y_resolve_serverless_keyspace o11y_resolve_serverless_tikv_hosts
    export -f o11y_gs_query o11y_find_cluster_record
fi