#!/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

ENV_FILE="${ENV_FILE:-.env}"
DS_AUTH_URL="https://www.ds.us-west-2.aws.observability.tidbcloud.com/client/auth"
GS_API_URL="http://www.gs.us-west-2.aws.observability.tidbcloud.com/api/v1/clusters/"

o11y_ds_metrics_base() {
    local region="${1:-us-west-2}"
    local vendor="${2:-aws}"
    echo "https://www.ds.${region}.${vendor}.observability.tidbcloud.com/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 param=""
    for param in "$@"; do
        [ -n "$param" ] && args+=(--data-urlencode "$param")
    done

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

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

    http_code=$(curl -s -o /dev/null -w '%{http_code}' \
        "$(o11y_ds_metrics_base "$region" "$vendor")/${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_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=""

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

    for candidate in "$primary" "$fallback"; do
        [ -z "$candidate" ] && continue
        if o11y_probe_metrics_target "$candidate" "$region" "$vendor"; 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
        prod-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 [ -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)}"

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

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

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

    local app_uuid
    app_uuid=$(o11y_pick_metrics_uuid "$app_json" "$region" "$vendor")
    if [ -z "$app_uuid" ]; then
        echo "ERROR: Failed to resolve application metrics UUID for cluster ${cluster_id}" >&2
        echo "$app_resp" >&2
        return 1
    fi

    local tenant_id
    tenant_id=$(echo "$app_json" | jq -r '.tenantID // 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 infra_json='{}'
    local keyspace_id=""
    local serverless_hosts='{"instances":[],"nodes":{}}'

    if [ -z "$namespace" ]; then
        if [[ "$biz_type" == *"devtier"* ]] || [ "$deploy_type" = "starter" ] || [ "$deploy_type" = "essential" ]; then
            namespace="tidb-serverless"
        else
            namespace="tidb${cluster_id}"
        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=""
    if [ "$(echo "$infra_json" | jq -r 'if . == null then "" else .id // empty end')" != "" ]; then
        infra_uuid=$(o11y_pick_metrics_uuid "$infra_json" "$region" "$vendor")
    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 '
                        if (.data.result | length) > 0
                        then (.data.result[0].value[1] | tonumber | floor)
                        else 0
                        end')
                if [ "$match_count" -gt 0 ]; then
                    infra_json=$(echo "$infra_candidates" | jq ".[$idx]")
                    infra_uuid="$candidate_uuid"
                    break
                fi
            done
        fi
    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

    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 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,
            infra_uuid: $infra_uuid,
            infra_application_id: $infra_application_id,
            keyspace_id: $keyspace_id,
            tikv_instances: ($serverless_hosts.instances // []),
            tikv_nodes: ($serverless_hosts.nodes // {})
        }'
}

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

    curl -s "$(o11y_ds_metrics_base "$region" "$vendor")/${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]}"
    local query="${2:?Usage: o11y_metrics_query <o11y_id> <promql> [start_ts] [end_ts] [step] [region] [vendor]}"
    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}"

    curl -s "$(o11y_ds_metrics_base "$region" "$vendor")/${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() {
    local region="${1:-us-west-2}"
    local vendor="${2:-aws}"
    echo "https://www.ds.${region}.${vendor}.observability.tidbcloud.com/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
    dedicated)             base="tidbcloud" ;;
    premium|resource_pool) base="tidbcloud-nextgen" ;;
    starter|essential)     base="devtier" ;;
    byoc)                  base="tidbcloud" ;;
    *)
        echo "ERROR: Unknown deployType '$deploy_type'. Must be dedicated/premium/starter/essential/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_resolve_serverless_keyspace o11y_resolve_serverless_tikv_hosts
    export -f o11y_gs_query o11y_find_cluster_record
fi