#!/usr/bin/env bash
# O11Y OLAP toolkit — auth, cluster lookup, Diagnosis Query (Athena) helpers.
#
# IMPORTANT: use bash (macOS /bin/bash 3.2+ or brew bash). Do not `source` from zsh.
#
#   bash -c 'source skills/o11y-olap/scripts/o11y-olap.sh && olap_auth && olap_tables us-west-2-staging'
#
# Env:
#   GITHUB_PAT / ENV_FILE — credentials (.env, searched in cwd + common paths)
#   O11Y_TOKEN — JWT for OLAP (auth response .token) — sent as o11ytoken
#   O11Y_AUTH0_TOKEN — GS cluster list (auth response .auth0_token)

# Refuse accidental zsh source (prints confusing function dumps / breaks export -f)
if [ -n "${ZSH_VERSION:-}" ]; then
  echo "ERROR: o11y-olap.sh must be sourced with bash, not zsh." >&2
  echo "  bash -c 'source skills/o11y-olap/scripts/o11y-olap.sh && olap_auth'" >&2
  return 1 2>/dev/null || exit 1
fi

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

OLAP_AWS_REGIONS=(
  us-west-2 us-east-1 us-east-2 us-west-1
  ap-northeast-1 ap-northeast-2 ap-northeast-3
  ap-southeast-1 ap-southeast-2 ap-southeast-3 ap-southeast-5 ap-southeast-7
  ap-east-1 ap-south-1
  eu-central-1 eu-west-1
  sa-east-1 af-south-1
)
OLAP_BIZ_TYPES=(
  tidbcloud
  tidbcloud-nextgen
  prod-tidbcloud-nextgen
  staging-tidbcloud-nextgen
  dev-tidbcloud-nextgen
  devtier
  prod-devtier
  staging-devtier
  dev-devtier
)

# ────────────────────────────────────────────────────────────
# Env file discovery
# ────────────────────────────────────────────────────────────
_olap_resolve_env_file() {
  if [ -n "${ENV_FILE:-}" ] && [ -f "$ENV_FILE" ]; then
    echo "$ENV_FILE"
    return 0
  fi
  local c
  for c in \
    .env \
    skills/o11y-olap/.env \
    "$HOME/code/test/observability-client/.env" \
    "$HOME/code/observability/.env"
  do
    if [ -f "$c" ]; then
      echo "$c"
      return 0
    fi
  done
  echo "${ENV_FILE:-.env}"
}

# ────────────────────────────────────────────────────────────
# Host vs x-region
# ────────────────────────────────────────────────────────────
# DS hostname uses the *logical* region (bucket / DNS), e.g. us-west-2-staging.
# Header x-region must be the *physical* cloud region from diagnosis-query
# deployment (helm `region:`), e.g. us-west-2 — not the logical suffix.
olap_base_url() {
  local logical="${1:?Usage: olap_base_url <logical_region> [vendor]}"
  local vendor="${2:-aws}"
  echo "https://www.ds.${logical}.${vendor}.observability.tidbcloud.com/internal/o11y/diagnosis/api/v1"
}

olap_x_region() {
  local logical="${1:?Usage: olap_x_region <logical_region>}"
  case "$logical" in
    *-staging) echo "${logical%-staging}" ;;
    *-premium) echo "${logical%-premium}" ;;
    *-serverless) echo "${logical%-serverless}" ;;
    *-ng) echo "${logical%-ng}" ;;
    *-na-test) echo "${logical%-na-test}" ;;
    *-byoc-*) echo "${logical%%-byoc-*}" ;;
    *) echo "$logical" ;;
  esac
}

# Normalize scope/list date values to yyyy-MM-dd when needed (S3 may return yyyyMMdd).
olap_normalize_date() {
  local d="${1:?}"
  if echo "$d" | grep -Eq '^[0-9]{8}$'; then
    echo "$d" | sed -E 's/([0-9]{4})([0-9]{2})([0-9]{2})/\1-\2-\3/'
  else
    echo "$d"
  fi
}

# Emit unix begin/end for a UTC calendar day as JSON fragment for time_range.
# Prefer begin_time/end_time keys — older diagnosis-query builds ignore begin/end.
olap_day_time_range_json() {
  local day="${1:?Usage: olap_day_time_range_json <yyyy-MM-dd|yyyyMMdd>}"
  day=$(olap_normalize_date "$day")
  local begin end
  if begin=$(date -u -j -f "%Y-%m-%d %H:%M:%S" "${day} 00:00:00" "+%s" 2>/dev/null); then
    :
  elif begin=$(date -u -d "${day} 00:00:00" "+%s" 2>/dev/null); then
    :
  else
    echo "ERROR: cannot parse date $day" >&2
    return 1
  fi
  end=$((begin + 86400))
  jq -nc --argjson b "$begin" --argjson e "$end" '{begin_time:$b, end_time:$e}'
}

# Shared curl headers. Always send o11ytoken (some gateways ignore Bearer alone).
_olap_curl() {
  local method="${1:?}"
  local url="${2:?}"
  local vendor="${3:?}"
  local x_region="${4:?}"
  local cluster_id="${5:-all}"
  local org_id="${6:-}"
  local deploy_type="${7:-}"
  local body="${8:-}"

  if [ -z "${O11Y_TOKEN:-}" ]; then
    echo "ERROR: O11Y_TOKEN unset — run olap_auth first" >&2
    return 1
  fi

  local args=(
    -sS -X "$method" "$url"
    -H "o11ytoken: ${O11Y_TOKEN}"
    -H "Authorization: Bearer ${O11Y_TOKEN}"
    -H "x-provider: ${vendor}"
    -H "x-region: ${x_region}"
    -H "x-cluster-id: ${cluster_id}"
  )
  [ -n "$org_id" ] && args+=(-H "x-org-id: ${org_id}")
  [ -n "$deploy_type" ] && args+=(-H "x-deploy-type: ${deploy_type}")
  if [ -n "$body" ]; then
    args+=(-H "Content-Type: application/json" -d "$body")
  fi
  curl "${args[@]}"
}

# ────────────────────────────────────────────────────────────
# Auth
# ────────────────────────────────────────────────────────────
olap_auth() {
  ENV_FILE="$(_olap_resolve_env_file)"
  export ENV_FILE
  if [ -f "$ENV_FILE" ]; then
    # shellcheck disable=SC1090
    set -a; source "$ENV_FILE"; set +a
  fi

  if [ -z "${GITHUB_PAT:-}" ]; then
    echo "ERROR: GITHUB_PAT not set. Provide your GitHub Personal Access Token." >&2
    echo -n "Enter GitHub PAT: " >&2
    read -r GITHUB_PAT
    if [ -z "$GITHUB_PAT" ]; then
      echo "ERROR: No token provided" >&2
      return 1
    fi
    [ ! -f "$ENV_FILE" ] && touch "$ENV_FILE"
    if ! grep -q '^GITHUB_PAT=' "$ENV_FILE" 2>/dev/null; then
      echo "GITHUB_PAT=${GITHUB_PAT}" >>"$ENV_FILE"
      echo "Saved GITHUB_PAT to $ENV_FILE" >&2
    fi
  fi

  local now
  now=$(date +%s)
  if [ -n "${O11Y_TOKEN:-}" ] && [ -n "${O11Y_TOKEN_EXPIRES:-}" ]; then
    local buf=$((O11Y_TOKEN_EXPIRES - 300))
    if [ "$now" -lt "$buf" ]; then
      echo "Using cached O11Y_TOKEN (expires in $((O11Y_TOKEN_EXPIRES - now))s) [env=$ENV_FILE]" >&2
      export O11Y_TOKEN O11Y_AUTH0_TOKEN="${O11Y_AUTH0_TOKEN:-}"
      if [ -n "${O11Y_AWS_ACCESS_KEY_ID:-}" ]; then
        export AWS_ACCESS_KEY_ID="$O11Y_AWS_ACCESS_KEY_ID"
        export AWS_SECRET_ACCESS_KEY="$O11Y_AWS_SECRET_ACCESS_KEY"
        export AWS_SESSION_TOKEN="$O11Y_AWS_SESSION_TOKEN"
      fi
      return 0
    fi
  fi

  local resp
  resp=$(curl -sS -X POST "$DS_AUTH_URL" \
    -H 'Content-Type: application/json' \
    -d "{\"github_token\": \"${GITHUB_PAT}\"}")

  local token auth0 ak sk st
  token=$(echo "$resp" | jq -r '.token // empty')
  auth0=$(echo "$resp" | jq -r '.auth0_token // empty')
  ak=$(echo "$resp" | jq -r '.access_key_id // empty')
  sk=$(echo "$resp" | jq -r '.secret_access_key // empty')
  st=$(echo "$resp" | jq -r '.session_token // empty')

  if [ -z "$token" ] || [ "$token" = "null" ]; then
    echo "ERROR: Authentication failed (missing .token)" >&2
    echo "$resp" | jq -c '.' 2>/dev/null || echo "$resp" >&2
    return 1
  fi

  local exp=$((now + 43200))
  local tmp="${ENV_FILE}.tmp"
  if [ -f "$ENV_FILE" ]; then
    grep -vE '^(O11Y_TOKEN|O11Y_AUTH0_TOKEN|O11Y_TOKEN_EXPIRES|O11Y_AWS_ACCESS_KEY_ID|O11Y_AWS_SECRET_ACCESS_KEY|O11Y_AWS_SESSION_TOKEN)=' \
      "$ENV_FILE" >"$tmp" || true
  else
    : >"$tmp"
  fi

  cat >>"$tmp" <<EOF
O11Y_TOKEN=${token}
O11Y_AUTH0_TOKEN=${auth0}
O11Y_TOKEN_EXPIRES=${exp}
O11Y_AWS_ACCESS_KEY_ID=${ak}
O11Y_AWS_SECRET_ACCESS_KEY=${sk}
O11Y_AWS_SESSION_TOKEN=${st}
EOF
  mv "$tmp" "$ENV_FILE"

  export O11Y_TOKEN="$token"
  export O11Y_AUTH0_TOKEN="$auth0"
  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}"

  echo "Authenticated successfully (O11Y_TOKEN + AWS STS, ~12h) [env=$ENV_FILE]" >&2
  if [ -z "$auth0" ]; then
    echo "WARN: auth0_token empty — cluster list via GS may fail" >&2
  fi
}

# ────────────────────────────────────────────────────────────
# Cluster list / find
# ────────────────────────────────────────────────────────────
olap_list_clusters() {
  local vendor="${1:?Usage: olap_list_clusters <vendor> <region> [bizType] [applicationID]}"
  local region="${2:?}"
  local biz="${3:-tidbcloud}"
  local app_id="${4:-}"

  if [ -z "${O11Y_AUTH0_TOKEN:-}" ]; then
    echo "ERROR: O11Y_AUTH0_TOKEN unset — run olap_auth first" >&2
    return 1
  fi

  local url="${GS_API_BASE}?bizType=${biz}&vendor=${vendor}&region=${region}"
  [ -n "$app_id" ] && url+="&applicationID=${app_id}"

  curl -sS "$url" -H "Authorization: Bearer ${O11Y_AUTH0_TOKEN}" | jq '.data.clusters // .data // .'
}

olap_find_cluster() {
  local app_id="${1:?Usage: olap_find_cluster <applicationID> [vendor] [hint_region]}"
  local vendor="${2:-aws}"
  local hint="${3:-}"

  if [ -z "${O11Y_AUTH0_TOKEN:-}" ]; then
    echo "ERROR: O11Y_AUTH0_TOKEN unset — run olap_auth first" >&2
    return 1
  fi

  local regions=()
  if [ -n "$hint" ]; then
    regions+=("$hint")
    for r in "${OLAP_AWS_REGIONS[@]}"; do
      [ "$r" != "$hint" ] && regions+=("$r")
    done
  else
    regions=("${OLAP_AWS_REGIONS[@]}")
  fi

  local biz region url resp count
  for biz in "${OLAP_BIZ_TYPES[@]}"; do
    for region in "${regions[@]}"; do
      url="${GS_API_BASE}?bizType=${biz}&vendor=${vendor}&region=${region}&applicationID=${app_id}"
      resp=$(curl -sS "$url" -H "Authorization: Bearer ${O11Y_AUTH0_TOKEN}" 2>/dev/null || true)
      count=$(echo "$resp" | jq -r '(.data.clusters // []) | length' 2>/dev/null || echo 0)
      if [ "$count" != "0" ] && [ "$count" != "null" ]; then
        echo "$resp" | jq --arg biz "$biz" --arg region "$region" --arg vendor "$vendor" '
          .data.clusters[0] + {bizType: $biz, region: $region, vendor: $vendor,
            deploy_type: (
              if ($biz | test("devtier")) then "shared"
              elif ($biz | test("nextgen")) then "premium"
              else "dedicated" end
            ),
            logical_region: (
              if ($biz | test("devtier")) then ($region + "-serverless")
              elif ($biz | test("nextgen")) then ($region + "-premium")
              else $region end
            )
          }'
        return 0
      fi
    done
  done

  echo "ERROR: cluster applicationID=${app_id} not found (tried ${#OLAP_BIZ_TYPES[@]} bizTypes × ${#regions[@]} regions)" >&2
  return 1
}

olap_biz_to_deploy() {
  local biz="${1:?}"
  case "$biz" in
    *devtier*) echo shared ;;
    *nextgen*) echo premium ;;
    *) echo dedicated ;;
  esac
}

# ────────────────────────────────────────────────────────────
# OLAP API (logical_region = DNS host segment)
# ────────────────────────────────────────────────────────────
olap_tables() {
  local logical="${1:?Usage: olap_tables <logical_region> [vendor] [cluster_id] [org_id] [deploy_type]}"
  local vendor="${2:-aws}"
  local cluster_id="${3:-all}"
  local org_id="${4:-}"
  local deploy_type="${5:-}"
  local base xreg
  base=$(olap_base_url "$logical" "$vendor")
  xreg=$(olap_x_region "$logical")
  _olap_curl GET "${base}/olap/tables" "$vendor" "$xreg" "$cluster_id" "$org_id" "$deploy_type" | jq '.'
}

# Body MUST include dataset + field (+ ancestor scope for non-root levels).
# Example: '{"dataset":"log-stats","field":"date","scope":{}}'
olap_scope_list() {
  local logical="${1:?Usage: olap_scope_list <logical_region> '<json with dataset+field>' [vendor] [cluster_id] [org_id] [deploy_type]}"
  local body="${2:?}"
  local vendor="${3:-aws}"
  local cluster_id="${4:-all}"
  local org_id="${5:-}"
  local deploy_type="${6:-}"
  local base xreg
  base=$(olap_base_url "$logical" "$vendor")
  xreg=$(olap_x_region "$logical")
  _olap_curl POST "${base}/olap/scope/list" "$vendor" "$xreg" "$cluster_id" "$org_id" "$deploy_type" "$body" | jq '.'
}

olap_submit() {
  local logical="${1:?Usage: olap_submit <logical_region> '<json>' [vendor] [cluster_id] [org_id] [deploy_type]}"
  local body="${2:?}"
  local vendor="${3:-aws}"
  local cluster_id="${4:-all}"
  local org_id="${5:-}"
  local deploy_type="${6:-}"
  local base xreg
  base=$(olap_base_url "$logical" "$vendor")
  xreg=$(olap_x_region "$logical")
  _olap_curl POST "${base}/olap" "$vendor" "$xreg" "$cluster_id" "$org_id" "$deploy_type" "$body"
}

olap_status() {
  local logical="${1:?}"
  local qid="${2:?}"
  local vendor="${3:-aws}"
  local cluster_id="${4:-all}"
  local org_id="${5:-}"
  local deploy_type="${6:-}"
  local base xreg
  base=$(olap_base_url "$logical" "$vendor")
  xreg=$(olap_x_region "$logical")
  _olap_curl GET "${base}/olap/${qid}" "$vendor" "$xreg" "$cluster_id" "$org_id" "$deploy_type"
}

olap_results() {
  local logical="${1:?}"
  local qid="${2:?}"
  local vendor="${3:-aws}"
  local cluster_id="${4:-all}"
  local org_id="${5:-}"
  local deploy_type="${6:-}"
  local base xreg
  base=$(olap_base_url "$logical" "$vendor")
  xreg=$(olap_x_region "$logical")
  _olap_curl GET "${base}/olap/${qid}/results" "$vendor" "$xreg" "$cluster_id" "$org_id" "$deploy_type"
}

olap_download() {
  local logical="${1:?}"
  local qid="${2:?}"
  local vendor="${3:-aws}"
  local cluster_id="${4:-all}"
  local org_id="${5:-}"
  local deploy_type="${6:-}"
  local base xreg
  base=$(olap_base_url "$logical" "$vendor")
  xreg=$(olap_x_region "$logical")
  _olap_curl GET "${base}/olap/${qid}/download" "$vendor" "$xreg" "$cluster_id" "$org_id" "$deploy_type"
}

olap_wait() {
  local logical="${1:?}"
  local qid="${2:?}"
  local vendor="${3:-aws}"
  local cluster_id="${4:-all}"
  local org_id="${5:-}"
  local deploy_type="${6:-}"
  local max_wait="${7:-120}"
  local elapsed=0
  local st resp

  while [ "$elapsed" -lt "$max_wait" ]; do
    resp=$(olap_status "$logical" "$qid" "$vendor" "$cluster_id" "$org_id" "$deploy_type")
    st=$(echo "$resp" | jq -r '.data.status // empty')
    case "$st" in
      SUCCEEDED|FAILED|CANCELLED)
        echo "$resp"
        return 0
        ;;
    esac
    sleep 2
    elapsed=$((elapsed + 2))
  done
  echo "$resp"
  echo "WARN: still running after ${max_wait}s (query_id=${qid}, status=${st})" >&2
  return 2
}

olap_query() {
  local logical="${1:?Usage: olap_query <logical_region> '<json>' [vendor] [cluster_id] [org_id] [deploy_type] [max_wait_s]}"
  local body="${2:?}"
  local vendor="${3:-aws}"
  local cluster_id="${4:-all}"
  local org_id="${5:-}"
  local deploy_type="${6:-}"
  local max_wait="${7:-120}"

  local submit qid wait_resp st
  submit=$(olap_submit "$logical" "$body" "$vendor" "$cluster_id" "$org_id" "$deploy_type")
  qid=$(echo "$submit" | jq -r '.data.query_id // empty')
  if [ -z "$qid" ]; then
    echo "ERROR: submit failed" >&2
    echo "$submit" | jq '.' 2>/dev/null || echo "$submit" >&2
    return 1
  fi
  echo "query_id=${qid} (host=$(olap_base_url "$logical" "$vendor") x-region=$(olap_x_region "$logical"))" >&2

  wait_resp=$(olap_wait "$logical" "$qid" "$vendor" "$cluster_id" "$org_id" "$deploy_type" "$max_wait")
  st=$(echo "$wait_resp" | jq -r '.data.status // empty')
  echo "$wait_resp" | jq '{status: .data.status, scanned_bytes: .data.scanned_bytes, error: .data.error_message}' >&2

  if [ "$st" != "SUCCEEDED" ]; then
    echo "$wait_resp" | jq '.'
    return 1
  fi

  olap_results "$logical" "$qid" "$vendor" "$cluster_id" "$org_id" "$deploy_type" | jq '.'
}

# First value from scope/list (or empty).
olap_scope_first() {
  local logical="${1:?}"
  local body="${2:?}"
  local vendor="${3:-aws}"
  olap_scope_list "$logical" "$body" "$vendor" | jq -r '.data.values[0] // empty'
}

if [[ "${BASH_SOURCE[0]}" != "${0}" ]]; then
  :
else
  cat <<'EOF'
Use bash:

  bash -c 'source skills/o11y-olap/scripts/o11y-olap.sh && olap_auth'

Key helpers:
  olap_auth
  olap_x_region <logical_region>     # physical x-region (strips -staging/-premium/…)
  olap_base_url <logical_region>
  olap_normalize_date / olap_day_time_range_json
  olap_scope_list <logical> '{"dataset":"log-stats","field":"date","scope":{}}'
  olap_query <logical> '<json>' [vendor] [cluster_id] [org_id] [deploy_type]
  bash skills/o11y-olap/scripts/smoke-test.sh us-west-2-staging

See skills/o11y-olap/SKILL.md
EOF
fi