#!/bin/bash
set -e
# This script applies rules for different environments and stacks
# Usage: ./apply-rules.sh <env> <stack> [region]
# env: dev or prod
# stack: eks,gke,azure,alicloud

ENV=$1
STACK=$2
REGION_FILTER=${3:-}

WORKSPACE=$(unset CDPATH && cd $(dirname "${BASH_SOURCE[0]}") && pwd)

# regions metadata
# Get all region objects as JSON strings
AWS_REGIONS=($(jq -c '.aws[]' "$WORKSPACE/regions_metadata.json"))
GCP_REGIONS=($(jq -c '.gcp[]' "$WORKSPACE/regions_metadata.json"))
AZURE_REGIONS=($(jq -c '.azure[]' "$WORKSPACE/regions_metadata.json"))
ALICLOUD_REGIONS=($(jq -c '.alicloud[]' "$WORKSPACE/regions_metadata.json"))

# Paths
DATA_PATH_DEV='newrules/dev/tidb-cloud/data-plane/'
CONTROL_PATH_DEV='newrules/dev/tidb-cloud/control-plane/'
DATA_PATH_PROD='newrules/prod/tidb-cloud/data-plane/'
CONTROL_PATH_PROD='newrules/prod/tidb-cloud/control-plane/'
LOGGING_PATH_PROD='newrules/prod/tidb-cloud/logging'

# rule IDs
DATA_PLANE_DEV_TIER_RULES_ID='tidb-cloud-data-plane-dev-tier-rule'
DATA_PLANE_DEDICATED_TIER_RULES_ID='95072d4e-5676-4084-bdf1-09aca85881f2'
DATA_PLANE_DATA_MIGRATION_RULES_ID='data-migration-data-plane-rule'
DATA_PLANE_SERVERLESS_RULES_ID='tidb-cloud-data-plane-serverless-rule'
DATA_PLANE_NEXTGEN_RULES_ID='tidb-cloud-data-plane-nextgen-rule'
CONTROL_PLANE_DEV_TIER_RULES_ID='tidb-cloud-control-plane-dev-tier-rule'
CONTROL_PLANE_DEDICATED_TIER_RULES_ID='d5025c09-d861-417c-9cfe-797bd97ce0c3'
CONTROL_PLANE_DEDICATED_SEED_RULES_ID='tidb-cloud-control-plane-dedicated-seed-rule'
CONTROL_PLANE_BASE_RULES_ID='tidb-cloud-control-plane-base-rule'
CONTROL_PLANE_CLINIC_RULES_ID='clinic-control-plane-clinic-rule'
CONTROL_PLANE_NEXTGEN_K8S_INFRA_RULES_ID='tidb-cloud-control-plane-nextgen-k8s-infra-rule'
CONTROL_PLANE_MEM9_RULES_ID='mem9-control-plane-mem9-rule'
CONTROL_PLANE_DEVOPS_EKS_RULES_ID='devops-eks-control-plane-rule'

GLOBAL_DOMAIN="www.gs.us-west-2.aws.observability.tidbcloud.com"
# GLOBAL_DOMAIN="www.gs.us-west-2.aws.observability-dev.pingcap.cloud"

# Dev/staging (`ENV=dev`): for logical region us-west-2-premium only, nextgen data-plane template
# omits Prometheus recording rules (`record:`); alerts are unchanged. Prod and all other regions are unaffected.
strip_recording_rules_from_rules_json() {
  local rules=$1
  echo "$rules" | jq '.groups = (
    (.groups // []) | map(
      (.rules // []) as $r |
      .rules = ($r | map(select(has("alert")))) |
      select((.rules | length) > 0)
    )
  )'
}

# Returns rules JSON for tidb-cloud-data-plane-nextgen-rule; strips recording rules only when
# env is dev/staging and logical_region is us-west-2-premium.
prepare_nextgen_dataplane_rules_for_region() {
  local logical_region=$1
  local rules=$2
  if [[ "$ENV" == "dev" && "$logical_region" == "us-west-2-premium" ]]; then
    echo "Recording rules omitted from ${DATA_PLANE_NEXTGEN_RULES_ID}${SUFFIX} for logical_region=${logical_region}" >&2
    strip_recording_rules_from_rules_json "$rules"
  else
    echo "$rules"
  fi
}

# Function to get rules from path
get_rules() {
  local path=$1
  local rules=""
  local files=$(find "$path" -name "*.yaml")

  for file in $files; do
    content=$(yq eval -o=json "$file" | jq -c '.groups')
    if [ -n "$content" ]; then
      rules="${rules}${content},"
    fi
  done

  echo "${rules%,}" | jq -s '{groups: add} '
}

# Function to write rules
write_rule() {
  local vendor=$1
  local region_obj=$2
  local rule_id=$3
  local rules=$4

  # Extract region info from the region object
  local region=$(echo "$region_obj" | jq -r '.region')
  local logical_region=$(echo "$region_obj" | jq -r '.logical_region')
  local tenant_id=$(echo "$region_obj" | jq -r '.tenantID')
  local premium=$(echo "$region_obj" | jq -r '.premium')
  local bizType=$(echo "$region_obj" | jq -r '.bizType')

  echo "Writing rules: ${rule_id} to ${vendor}, region: ${region}, logical_region: ${logical_region}, premium: ${premium}, bizType: ${bizType}"

  # Create JSON payload using region
  data=$(
    cat <<EOF
{
  "vendor": "${vendor}",
  "region": "${region}",
  "templateID": "${rule_id}",
  "monitorPromRules": ${rules}
}
EOF
  )
  if [[ $rule_id == "logging" ]]; then
    data=$(
      cat <<EOF
{
  "vendor": "${vendor}",
  "region": "${region}",
  "rulesConfig": ${rules}
}
EOF
    )
  fi

  # Write data to temporary file to avoid "Argument list too long" error
  tmp_file=$(mktemp)
  echo "${data}" >"${tmp_file}"

  # Build curl headers
  headers=(
    --header "Authorization: Bearer ${TOKEN}"
    --header "Content-Type: application/json"
  )
  headers+=(--header "X-Region-Id: ${vendor}-${region}")
  headers+=(--header "X-Logical-Region: ${logical_region}")

  # Normalize: jq returns "null" as string when field is missing
  if [ "$bizType" == "null" ]; then
    bizType=""
  fi

  if [ "$tenant_id" == "null" ]; then
    tenant_id=""
  fi

  if [ -z "$bizType" ]; then
    # Add tenant-specific headers if tenantID is not empty and bizType is empty
    if [ -n "$tenant_id" ]; then
      bizType="byoc"
    # Add biztype header for premium regions, if not set
    elif [ "$premium" == "true" ] || [[ "$logical_region" == *"premium"* ]] || [[ "$logical_region" == *"-ng" ]]; then
      bizType="tidbcloud-nextgen"
    fi
  fi

  if [ -n "$bizType" ]; then
    # Use `fake` as tenant ID header if not set
    if [ -z "$tenant_id" ]; then
      tenant_id="fake"
    fi
    headers+=(--header "X-Org-Id: ${tenant_id}")
    headers+=(--header "X-Biz-Type: ${bizType}")
    echo "Adding tenant headers for tenant ID: ${tenant_id}, biztype: ${bizType}"
  fi

  url="http://${GLOBAL_DOMAIN}/api/v1/alertrule-template"
  if [[ $rule_id == "logging" ]]; then
    url="http://${GLOBAL_DOMAIN}/api/v1/alertrule-template/logging"
  fi

  # Send HTTP request using file input instead of direct data
  # Retry logic for HTTP request
  max_retries=3
  retry_delay=5
  attempt=0
  while true; do
    attempt=$((attempt + 1))

    local response http_code curl_exit=0

    # Execute curl and capture response with HTTP code
    response=$(curl --silent --show-error --location --connect-timeout 10 -m 60 --request POST "$url" \
      "${headers[@]}" \
      --data @"${tmp_file}" \
      --write-out "\n%{http_code}") || curl_exit=$?

    # Extract HTTP code from response (last line)
    http_code=$(echo "$response" | tail -n1)
    response=$(echo "$response" | sed '$d')

    # Clean up temp file
    rm -f "${tmp_file}"

    # Check curl exit code
    if [ $curl_exit -ne 0 ]; then
      if [ $attempt -lt $max_retries ]; then
        sleep $retry_delay
        # Recreate temp file for retry
        tmp_file=$(mktemp) || { echo "  [ERROR] Failed to create temp file for retry"; return 1; }
        echo "${data}" >"${tmp_file}"
        continue
      fi
      echo "  [ERROR] curl failed for rule_id=${rule_id}, vendor=${vendor}, region=${region}, logical_region=${logical_region}, exit_code=${curl_exit}"
      return 1
    fi

    # Check HTTP status code
    if [ "$http_code" -lt 200 ] || [ "$http_code" -ge 300 ]; then
      if [ $attempt -lt $max_retries ]; then
        sleep $retry_delay
        # Recreate temp file for retry
        tmp_file=$(mktemp) || { echo "  [ERROR] Failed to create temp file for retry"; return 1; }
        echo "${data}" >"${tmp_file}"
        continue
      fi
      echo "  [ERROR] HTTP request failed for rule_id=${rule_id}, vendor=${vendor}, region=${region}, logical_region=${logical_region}, http_code=${http_code}, response=${response}"
      return 1
    fi

    # Check API response code
    local code
    code=$(echo "$response" | jq -r '.code // empty' 2>/dev/null || echo "parse_error")

    if [ "$code" == "0" ] || [ "$code" == "null" ] || [ -z "$code" ]; then
      return 0
    fi

    if [ $attempt -lt $max_retries ]; then
      sleep $retry_delay
      # Recreate temp file for retry
      tmp_file=$(mktemp) || { echo "  [ERROR] Failed to create temp file for retry"; return 1; }
      echo "${data}" >"${tmp_file}"
      continue
    fi
    echo "  [ERROR] API returned non-zero code for rule_id=${rule_id}, vendor=${vendor}, region=${region}, logical_region=${logical_region}, code=${code}, response=${response}"
    return 1
  done

  # Clean up temporary file
  rm "${tmp_file}"

  echo "Write rule: ${rule_id} to ${vendor}, region: ${region}, logical_region: ${logical_region}, done"
}

# Apply rules based on environment and stack
case "$ENV" in
  dev)
    DATA_PATH="${DATA_PATH_DEV}"
    CONTROL_PATH="${CONTROL_PATH_DEV}"
    SUFFIX="-dev"
    ;;
  prod)
    DATA_PATH="${DATA_PATH_PROD}"
    CONTROL_PATH="${CONTROL_PATH_PROD}"
    LOGGING_PATH="${LOGGING_PATH_PROD}"
    SUFFIX=""
    ;;
  *)
    echo "Unsupported environment: $ENV. Only dev and prod are supported"
    exit 1
    ;;
esac

# Apply rules for the specified stack
if [ "$STACK" == "eks" ]; then
  VENDOR="aws"
  REGIONS=("${AWS_REGIONS[@]}")
# azure and alicloud
elif [ "$STACK" == "azure" ]; then
  VENDOR="azure"
  REGIONS=("${AZURE_REGIONS[@]}")
elif [ "$STACK" == "alicloud" ]; then
  VENDOR="alicloud"
  REGIONS=("${ALICLOUD_REGIONS[@]}")
else
  VENDOR="gcp"
  REGIONS=("${GCP_REGIONS[@]}")
fi

if [ -n "$REGION_FILTER" ]; then
  FILTERED_REGIONS=()

  for region_obj in "${REGIONS[@]}"; do
    region=$(echo "$region_obj" | jq -r '.region')
    logical_region=$(echo "$region_obj" | jq -r '.logical_region')

    if [ "$region" == "$REGION_FILTER" ] || [ "$logical_region" == "$REGION_FILTER" ]; then
      FILTERED_REGIONS+=("$region_obj")
    fi
  done

  if [ ${#FILTERED_REGIONS[@]} -eq 0 ]; then
    echo "No region matched filter '$REGION_FILTER' in scripts/regions_metadata.json for stack '$STACK'"
    exit 1
  fi

  REGIONS=("${FILTERED_REGIONS[@]}")
fi

# Get rules for each component
DATA_PLANE_DEV_TIER_RULES=$(get_rules "${DATA_PATH}dev-tier")
DATA_PLANE_DEDICATED_TIER_RULES=$(get_rules "${DATA_PATH}dedicated-tier")
DATA_PLANE_DATA_MIGRATION_RULES=$(get_rules "${DATA_PATH}data-migration")
DATA_PLANE_SERVERLESS_RULES=$(get_rules "${DATA_PATH}serverless")
DATA_PLANE_NEXTGEN_RULES=$(get_rules "${DATA_PATH}nextgen")
# control plane
CONTROL_PLANE_DEV_TIER_RULES=$(get_rules "${CONTROL_PATH}serverless")
CONTROL_PLANE_DEDICATED_TIER_RULES=$(get_rules "${CONTROL_PATH}dedicated-tier")
CONTROL_PLANE_DEDICATED_SEED_RULES=$(get_rules "${CONTROL_PATH}dedicated-seed")
CONTROL_PLANE_BASE_RULES=$(get_rules "${CONTROL_PATH}base")
CONTROL_PLANE_CLINIC_RULES=$(get_rules "${CONTROL_PATH}clinic")
CONTROL_PLANE_NEXTGEN_RULES=$(get_rules "${CONTROL_PATH}nextgen-k8s-infra")
CONTROL_PLANE_MEM9_RULES=$(get_rules "${CONTROL_PATH}mem9")
CONTROL_PLANE_DEVOPS_EKS_RULES=$(get_rules "${CONTROL_PATH}devops-eks")
# logging
LOGGING_RULES=$(get_rules "${LOGGING_PATH}")

echo "get all rules done"

# Initialize arrays to track success and failure
SUCCESSFUL_REGIONS=()
FAILED_REGIONS=()
FAILED_BYOC_REGIONS=()

# # Apply rules to all regions
for region_obj in "${REGIONS[@]}"; do
  region=$(echo "$region_obj" | jq -r '.region')
  logical_region=$(echo "$region_obj" | jq -r '.logical_region')
  region_failed=false

  echo "Processing region: $region (logical: $logical_region)"

  # Track failures for this region
  # start apply dataplane
  write_rule "$VENDOR" "$region_obj" "${DATA_PLANE_DEV_TIER_RULES_ID}${SUFFIX}" "$DATA_PLANE_DEV_TIER_RULES" || region_failed=true
  write_rule "$VENDOR" "$region_obj" "${DATA_PLANE_DEDICATED_TIER_RULES_ID}${SUFFIX}" "$DATA_PLANE_DEDICATED_TIER_RULES" || region_failed=true
  write_rule "$VENDOR" "$region_obj" "${DATA_PLANE_DATA_MIGRATION_RULES_ID}${SUFFIX}" "$DATA_PLANE_DATA_MIGRATION_RULES" || region_failed=true
  write_rule "$VENDOR" "$region_obj" "${DATA_PLANE_SERVERLESS_RULES_ID}${SUFFIX}" "$DATA_PLANE_SERVERLESS_RULES" || region_failed=true
  nextgen_payload=$(prepare_nextgen_dataplane_rules_for_region "$logical_region" "$DATA_PLANE_NEXTGEN_RULES")
  write_rule "$VENDOR" "$region_obj" "${DATA_PLANE_NEXTGEN_RULES_ID}${SUFFIX}" "$nextgen_payload" || region_failed=true

  # start apply controlplane
  write_rule "$VENDOR" "$region_obj" "${CONTROL_PLANE_DEV_TIER_RULES_ID}${SUFFIX}" "$CONTROL_PLANE_DEV_TIER_RULES" || region_failed=true
  write_rule "$VENDOR" "$region_obj" "${CONTROL_PLANE_DEDICATED_TIER_RULES_ID}${SUFFIX}" "$CONTROL_PLANE_DEDICATED_TIER_RULES" || region_failed=true
  write_rule "$VENDOR" "$region_obj" "${CONTROL_PLANE_DEDICATED_SEED_RULES_ID}${SUFFIX}" "$CONTROL_PLANE_DEDICATED_SEED_RULES" || region_failed=true
  write_rule "$VENDOR" "$region_obj" "${CONTROL_PLANE_BASE_RULES_ID}${SUFFIX}" "$CONTROL_PLANE_BASE_RULES" || region_failed=true
  write_rule "$VENDOR" "$region_obj" "${CONTROL_PLANE_CLINIC_RULES_ID}${SUFFIX}" "$CONTROL_PLANE_CLINIC_RULES" || region_failed=true
  write_rule "$VENDOR" "$region_obj" "${CONTROL_PLANE_NEXTGEN_K8S_INFRA_RULES_ID}${SUFFIX}" "$CONTROL_PLANE_NEXTGEN_RULES" || region_failed=true
  write_rule "$VENDOR" "$region_obj" "${CONTROL_PLANE_MEM9_RULES_ID}${SUFFIX}" "$CONTROL_PLANE_MEM9_RULES" || region_failed=true
  write_rule "$VENDOR" "$region_obj" "${CONTROL_PLANE_DEVOPS_EKS_RULES_ID}${SUFFIX}" "$CONTROL_PLANE_DEVOPS_EKS_RULES" || region_failed=true

  if [[ $ENV == "prod" ]]; then
    # start apply logging
    write_rule "$VENDOR" "$region_obj" "logging" "$LOGGING_RULES" || region_failed=true
  fi

  # Add to appropriate array
  if [ "$region_failed" = true ]; then
    FAILED_REGIONS+=("$logical_region")
    # Check if this is a BYOC region (has tenantID)
    tenant_id=$(echo "$region_obj" | jq -r '.tenantID')
    if [ "$tenant_id" != "null" ] && [ -n "$tenant_id" ]; then
      FAILED_BYOC_REGIONS+=("$logical_region")
    fi
  else
    SUCCESSFUL_REGIONS+=("$logical_region")
  fi
done

# Print summary
echo ""
echo "========================================="
echo "DEPLOYMENT SUMMARY"
echo "========================================="
echo "Environment: $ENV"
echo "Stack: $STACK"
echo "Vendor: $VENDOR"
echo "Region filter: ${REGION_FILTER:-all}"
echo ""

if [ ${#SUCCESSFUL_REGIONS[@]} -gt 0 ]; then
  echo "✅ SUCCESSFUL REGIONS (${#SUCCESSFUL_REGIONS[@]}):"
  for region in "${SUCCESSFUL_REGIONS[@]}"; do
    echo "  - $region"
  done
  echo ""
fi

if [ ${#FAILED_REGIONS[@]} -gt 0 ]; then
  echo "❌ FAILED REGIONS (${#FAILED_REGIONS[@]}):"
  for region in "${FAILED_REGIONS[@]}"; do
    echo "  - $region"
  done
  echo ""

  # Check if all failed regions are BYOC regions
  if [ ${#FAILED_REGIONS[@]} -eq ${#FAILED_BYOC_REGIONS[@]} ]; then
    echo "ℹ️  All failed regions are BYOC regions - treating as success"
    echo "Exit code: 0 (only BYOC regions failed)"
    exit 0
  else
    echo "Exit code: 1 (some non-BYOC regions failed)"
    exit 1
  fi
else
  echo "🎉 All regions deployed successfully!"
  echo "Exit code: 0"
fi