#!/usr/bin/env bash

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ROOT_DIR="$(cd "${SCRIPT_DIR}/../.." && pwd)"
ENTERPRISE_SH="${SCRIPT_DIR}/enterprise.sh"

DEFAULT_REMOTE_DIR="/opt/security-center/enterprise"

ENV_FILE="${SCRIPT_DIR}/enterprise.prod.env"
AWS_PROFILE_NAME=""
AWS_REGION_NAME=""
INSTANCE_ID=""
REMOTE_DIR="${DEFAULT_REMOTE_DIR}"
S3_BUCKET=""
S3_PREFIX=""
DB_BACKUP_S3_BUCKET=""
DB_BACKUP_S3_PREFIX=""
DB_BACKUP_S3_REGION=""
APPLY="false"
CONTINUE_DEPLOY="false"
SKIP_BACKUP="false"

usage() {
  cat <<EOF
Safely fix known production schema_migrations checksum drift for CAASM identity migrations.

This script:
  1. optionally runs the enterprise DB backup command
  2. checks the remote ledger checksums for known CAASM identity migrations
  3. verifies the live schema already satisfies the migration intent
  4. updates only known drifted ledger rows when --apply is set
  5. optionally resumes build-deploy with --skip-upload --skip-build

Default mode is read-only. Pass --apply to mutate the ledger.

Usage:
  bash scripts/enterprise/fix-202604110001-checksum-drift.sh \\
    --env-file ./scripts/enterprise/enterprise.prod.env \\
    --profile <aws-profile> \\
    --region <aws-region> \\
    --instance-id <i-xxxxxxxxxxxxxxxxx> \\
    --apply \\
    --continue-deploy

Options:
  --env-file <path>       Enterprise env file. Default: scripts/enterprise/enterprise.prod.env
  --profile <name>        AWS CLI profile
  --region <name>         AWS region
  --instance-id <id>      Target EC2 instance ID for SSM
  --remote-dir <path>     Remote enterprise root. Default: ${DEFAULT_REMOTE_DIR}
  --s3-bucket <name>      Artifact staging bucket for continue-deploy, if not in env
  --s3-prefix <prefix>    Artifact staging prefix for continue-deploy
  --backup-s3-bucket <n>  DB backup bucket. If omitted during --apply, an account/region bucket is generated.
  --backup-s3-prefix <p>  DB backup prefix. Default: security-center-enterprise
  --backup-s3-region <r>  DB backup bucket region. Default: --region
  --apply                 Actually update schema_migrations; without this, check only
  --continue-deploy       After successful fix, run build-deploy --skip-upload --skip-build
  --skip-backup           Skip enterprise DB backup. Not recommended for production.
  --help                  Show this help
EOF
}

die() {
  echo "ERROR: $*" >&2
  exit 1
}

info() {
  echo "INFO: $*"
}

trim() {
  local value="${1:-}"
  value="${value#"${value%%[![:space:]]*}"}"
  value="${value%"${value##*[![:space:]]}"}"
  printf '%s' "${value}"
}

get_env_value() {
  local env_file="$1"
  local key="$2"
  [[ -f "${env_file}" ]] || return 0
  python3 - "${env_file}" "${key}" <<'PY'
from pathlib import Path
import sys

path = Path(sys.argv[1])
key = sys.argv[2]
for raw in path.read_text().splitlines():
    line = raw.strip()
    if not line or line.startswith("#") or "=" not in line:
        continue
    lhs, rhs = line.split("=", 1)
    if lhs.strip() == key:
        print(rhs.strip())
        break
PY
}

normalize_s3_bucket_name() {
  local raw="${1:-}"
  local normalized
  normalized="$(printf '%s' "${raw}" | tr '[:upper:]' '[:lower:]' | sed -E 's/[^a-z0-9.-]+/-/g; s/^[.-]+//; s/[.-]+$//; s/-+/-/g; s/\.+/./g')"
  normalized="${normalized:0:63}"
  normalized="$(printf '%s' "${normalized}" | sed -E 's/^[.-]+//; s/[.-]+$//')"
  if [[ -z "${normalized}" ]]; then
    normalized="security-center-db-backups"
  fi
  printf '%s' "${normalized}"
}

aws_cli() {
  local -a cmd=(aws)
  [[ -n "${AWS_PROFILE_NAME}" ]] && cmd+=(--profile "${AWS_PROFILE_NAME}")
  [[ -n "${AWS_REGION_NAME}" ]] && cmd+=(--region "${AWS_REGION_NAME}")
  cmd+=("$@")
  "${cmd[@]}"
}

require_inputs() {
  [[ -n "${INSTANCE_ID}" ]] || die "--instance-id is required"
  [[ -f "${ENV_FILE}" ]] || die "env file not found: ${ENV_FILE}"
  command -v aws >/dev/null 2>&1 || die "aws CLI is required"
  command -v python3 >/dev/null 2>&1 || die "python3 is required to encode the SSM command"
}

ensure_backup_defaults() {
  DB_BACKUP_S3_BUCKET="${DB_BACKUP_S3_BUCKET:-$(trim "$(get_env_value "${ENV_FILE}" "SC_DB_BACKUP_S3_BUCKET")")}"
  DB_BACKUP_S3_PREFIX="${DB_BACKUP_S3_PREFIX:-$(trim "$(get_env_value "${ENV_FILE}" "SC_DB_BACKUP_S3_PREFIX")")}"
  DB_BACKUP_S3_REGION="${DB_BACKUP_S3_REGION:-$(trim "$(get_env_value "${ENV_FILE}" "SC_DB_BACKUP_S3_REGION")")}"

  DB_BACKUP_S3_REGION="${DB_BACKUP_S3_REGION:-${AWS_REGION_NAME}}"
  DB_BACKUP_S3_PREFIX="${DB_BACKUP_S3_PREFIX:-security-center-enterprise}"

  if [[ -z "${DB_BACKUP_S3_BUCKET}" ]]; then
    local account_id region_for_name
    account_id="$(aws_cli sts get-caller-identity --query Account --output text)"
    region_for_name="${DB_BACKUP_S3_REGION:-${AWS_REGION_NAME:-unknown-region}}"
    DB_BACKUP_S3_BUCKET="$(normalize_s3_bucket_name "security-center-db-backups-${account_id}-${region_for_name}")"
    info "auto-selected DB backup bucket: ${DB_BACKUP_S3_BUCKET}"
  fi
}

run_enterprise_backup() {
  if [[ "${SKIP_BACKUP}" == "true" ]]; then
    info "skipping enterprise DB backup because --skip-backup was provided"
    return 0
  fi

  ensure_backup_defaults
  info "running enterprise DB backup before ledger mutation"
  local -a cmd=("${ENTERPRISE_SH}" db-backup --env-file "${ENV_FILE}" --instance-id "${INSTANCE_ID}" --resolve-ssm --non-interactive)
  [[ -n "${AWS_PROFILE_NAME}" ]] && cmd+=(--profile "${AWS_PROFILE_NAME}")
  [[ -n "${AWS_REGION_NAME}" ]] && cmd+=(--region "${AWS_REGION_NAME}")
  [[ -n "${DB_BACKUP_S3_BUCKET}" ]] && cmd+=(--backup-s3-bucket "${DB_BACKUP_S3_BUCKET}")
  [[ -n "${DB_BACKUP_S3_PREFIX}" ]] && cmd+=(--backup-s3-prefix "${DB_BACKUP_S3_PREFIX}")
  [[ -n "${DB_BACKUP_S3_REGION}" ]] && cmd+=(--backup-s3-region "${DB_BACKUP_S3_REGION}")
  "${cmd[@]}"
}

remote_fix_script() {
  cat <<REMOTE
set -euo pipefail

APPLY="${APPLY}"
REMOTE_DIR="${REMOTE_DIR}"

MIGRATION_ASSET="202604110001_add_asset_identity_unique_key.sql"
ASSET_OLD_CHECKSUM="8e52d82791f7ba2cb26ddaa7cf849ec933b51a8db71edcae719c5defd59396cc"
ASSET_NEW_CHECKSUM="5f9f50a0f0869593851d37d7c27d9ab6355752ea4469b11fe20c35ca026eb561"
MIGRATION_DNS="202604110002_add_dns_record_identity_key_columns.sql"
DNS_OLD_CHECKSUM="296bbda5e57320b98cda07e1057036da8522dc80cc889376d5200b5a6fcceae4"
DNS_NEW_CHECKSUM="4796a57418a426d60bb7deccb0be69bae952d65fcaca00ec88afc8bea629296e"
MIGRATIONS_DIR="./security-center-backend/migrations"

cd "\${REMOTE_DIR}/rendered"
test -f .env.runtime || { echo "missing .env.runtime in \${REMOTE_DIR}/rendered" >&2; exit 20; }
test -f docker-compose.enterprise.yml || { echo "missing docker-compose.enterprise.yml in \${REMOTE_DIR}/rendered" >&2; exit 21; }
test -d "\${MIGRATIONS_DIR}" || { echo "missing migrations directory: \${MIGRATIONS_DIR}" >&2; exit 22; }

mysql_query() {
  local sql="\$1"
  docker compose --env-file .env.runtime -f docker-compose.enterprise.yml exec -T mysql \\
    sh -lc 'mysql -u"\${MYSQL_USER}" -p"\${MYSQL_PASSWORD}" "\${MYSQL_DATABASE}" --batch --skip-column-names -e "\$1"' sh "\${sql}"
}

file_checksum() {
  local filename="\$1"
  sha256sum "\${MIGRATIONS_DIR}/\${filename}" | awk '{print \$1}'
}

ledger_checksum() {
  local filename="\$1"
  mysql_query "SELECT COALESCE((SELECT checksum FROM schema_migrations WHERE filename='\${filename}'), '');" | tr -d '\r'
}

allowed_old_checksum() {
  local filename="\$1"
  case "\${filename}" in
    "\${MIGRATION_ASSET}") printf '%s' "\${ASSET_OLD_CHECKSUM}" ;;
    "\${MIGRATION_DNS}") printf '%s' "\${DNS_OLD_CHECKSUM}" ;;
    *) return 1 ;;
  esac
}

allowed_new_checksum() {
  local filename="\$1"
  case "\${filename}" in
    "\${MIGRATION_ASSET}") printf '%s' "\${ASSET_NEW_CHECKSUM}" ;;
    "\${MIGRATION_DNS}") printf '%s' "\${DNS_NEW_CHECKSUM}" ;;
    *) return 1 ;;
  esac
}

print_known_state() {
  local filename ledger file expected_old expected_new
  for filename in "\${MIGRATION_ASSET}" "\${MIGRATION_DNS}"; do
    ledger="\$(ledger_checksum "\${filename}")"
    file="\$(file_checksum "\${filename}")"
    expected_old="\$(allowed_old_checksum "\${filename}")"
    expected_new="\$(allowed_new_checksum "\${filename}")"
    echo "migration: \${filename}"
    echo "  ledger checksum: \${ledger:-<missing>}"
    echo "  file checksum: \${file}"
    echo "  allowed old: \${expected_old}"
    echo "  allowed new: \${expected_new}"
  done
}

verify_asset_schema() {
  local index_shape
  index_shape="\$(mysql_query "SELECT COALESCE((SELECT GROUP_CONCAT(CONCAT(SEQ_IN_INDEX, ':', COLUMN_NAME, ':', NON_UNIQUE) ORDER BY SEQ_IN_INDEX SEPARATOR ',') FROM information_schema.STATISTICS WHERE TABLE_SCHEMA=DATABASE() AND TABLE_NAME='assets' AND INDEX_NAME='idx_assets_cloud_account_native'), '');" | tr -d '\r')"
  echo "asset index shape: \${index_shape}"
  if [[ "\${index_shape}" != "1:cloud_account_id:0,2:native_id:0" ]]; then
    echo "unexpected idx_assets_cloud_account_native shape; refusing to update ledger" >&2
    exit 30
  fi
}

verify_dns_schema() {
  local column_count index_shape mismatch_count duplicate_count
  column_count="\$(mysql_query "SELECT COUNT(*) FROM information_schema.COLUMNS WHERE TABLE_SCHEMA=DATABASE() AND TABLE_NAME='dns_records' AND COLUMN_NAME IN ('name_key','type_key','set_identifier_key','region_key','failover_key');" | tr -d '\r')"
  echo "dns key column count: \${column_count}"
  if [[ "\${column_count}" != "5" ]]; then
    echo "dns_records identity key columns are incomplete; refusing to update ledger" >&2
    exit 40
  fi

  index_shape="\$(mysql_query "SELECT COALESCE((SELECT GROUP_CONCAT(CONCAT(SEQ_IN_INDEX, ':', COLUMN_NAME, ':', NON_UNIQUE) ORDER BY SEQ_IN_INDEX SEPARATOR ',') FROM information_schema.STATISTICS WHERE TABLE_SCHEMA=DATABASE() AND TABLE_NAME='dns_records' AND INDEX_NAME='idx_dns_records_identity'), '');" | tr -d '\r')"
  echo "dns index shape: \${index_shape}"
  if [[ "\${index_shape}" != "1:domain_asset_id:0,2:name_key:0,3:type_key:0,4:set_identifier_key:0,5:region_key:0,6:failover_key:0" ]]; then
    echo "unexpected idx_dns_records_identity shape; refusing to update ledger" >&2
    exit 41
  fi

  mismatch_count="\$(mysql_query "SELECT COUNT(*) FROM dns_records WHERE COALESCE(name_key, '') <> LOWER(TRIM(TRAILING '.' FROM TRIM(COALESCE(name, '')))) OR COALESCE(type_key, '') <> UPPER(TRIM(COALESCE(type, ''))) OR COALESCE(set_identifier_key, '') <> CASE WHEN set_identifier IS NULL OR TRIM(set_identifier) = '' OR LOWER(TRIM(set_identifier)) = '<nil>' THEN '' ELSE TRIM(set_identifier) END OR COALESCE(region_key, '') <> CASE WHEN region IS NULL OR TRIM(region) = '' OR LOWER(TRIM(region)) = '<nil>' THEN '' ELSE LOWER(TRIM(region)) END OR COALESCE(failover_key, '') <> CASE WHEN failover IS NULL OR TRIM(failover) = '' OR LOWER(TRIM(failover)) = '<nil>' THEN '' ELSE UPPER(TRIM(failover)) END;" | tr -d '\r')"
  echo "dns key mismatch count: \${mismatch_count}"
  if [[ "\${mismatch_count}" != "0" ]]; then
    echo "dns_records identity keys do not match current normalization; refusing to update ledger" >&2
    exit 42
  fi

  duplicate_count="\$(mysql_query "SELECT COUNT(*) FROM (SELECT domain_asset_id, COALESCE(name_key, '') AS name_key, COALESCE(type_key, '') AS type_key, COALESCE(set_identifier_key, '') AS set_identifier_key, COALESCE(region_key, '') AS region_key, COALESCE(failover_key, '') AS failover_key FROM dns_records GROUP BY domain_asset_id, name_key, type_key, set_identifier_key, region_key, failover_key HAVING COUNT(*) > 1) duplicate_dns_identities;" | tr -d '\r')"
  echo "dns duplicate identity group count: \${duplicate_count}"
  if [[ "\${duplicate_count}" != "0" ]]; then
    echo "duplicate dns identity groups exist; refusing to update ledger" >&2
    exit 43
  fi
}

scan_drift() {
  local mode="\${1:-before}"
  local drift_count=0
  local filename ledger file old new
  echo "schema_migrations drift scan (\${mode})"

  while IFS=\$'\t' read -r filename ledger; do
    [[ -n "\${filename}" ]] || continue
    if [[ ! -f "\${MIGRATIONS_DIR}/\${filename}" ]]; then
      echo "unexpected drift: \${filename} ledger=\${ledger} file=<missing>"
      drift_count=\$((drift_count + 1))
      continue
    fi

    file="\$(file_checksum "\${filename}")"
    if [[ "\${ledger}" == "\${file}" ]]; then
      continue
    fi

    if old="\$(allowed_old_checksum "\${filename}" 2>/dev/null)" && new="\$(allowed_new_checksum "\${filename}" 2>/dev/null)" && [[ "\${ledger}" == "\${old}" && "\${file}" == "\${new}" ]]; then
      if [[ "\${mode}" == "before" ]]; then
        echo "allowed drift: \${filename} ledger=\${ledger} file=\${file}"
        continue
      fi
    fi

    echo "unexpected drift: \${filename} ledger=\${ledger} file=\${file}"
    drift_count=\$((drift_count + 1))
  done < <(mysql_query "SELECT filename, checksum FROM schema_migrations ORDER BY filename;" | tr -d '\r')

  if [[ "\${drift_count}" != "0" ]]; then
    echo "unexpected schema_migrations drift detected: \${drift_count}" >&2
    exit 50
  fi
  echo "drift scan \${mode}: ok"
}

update_known_drift() {
  local filename="\$1"
  local old_checksum="\$2"
  local new_checksum="\$3"
  local ledger updated_rows verified_checksum

  ledger="\$(ledger_checksum "\${filename}")"
  if [[ "\${ledger}" == "\${new_checksum}" ]]; then
    echo "\${filename}: ledger already has new checksum"
    return 0
  fi
  if [[ "\${ledger}" != "\${old_checksum}" ]]; then
    echo "\${filename}: ledger checksum is neither expected old nor new value; refusing to update" >&2
    exit 60
  fi
  if [[ "\${APPLY}" != "true" ]]; then
    echo "\${filename}: dry-run only, would update checksum to \${new_checksum}"
    return 0
  fi

  updated_rows="\$(mysql_query "UPDATE schema_migrations SET checksum='\${new_checksum}' WHERE filename='\${filename}' AND checksum='\${old_checksum}'; SELECT ROW_COUNT();" | tail -n 1 | tr -d '\r')"
  echo "\${filename}: updated rows: \${updated_rows}"
  if [[ "\${updated_rows}" != "1" ]]; then
    echo "\${filename}: expected exactly one ledger row update, got \${updated_rows}" >&2
    exit 61
  fi

  verified_checksum="\$(ledger_checksum "\${filename}")"
  echo "\${filename}: verified checksum: \${verified_checksum}"
  if [[ "\${verified_checksum}" != "\${new_checksum}" ]]; then
    echo "\${filename}: checksum verification failed after update" >&2
    exit 62
  fi
}

print_known_state
verify_asset_schema
verify_dns_schema
scan_drift before

update_known_drift "\${MIGRATION_ASSET}" "\${ASSET_OLD_CHECKSUM}" "\${ASSET_NEW_CHECKSUM}"
update_known_drift "\${MIGRATION_DNS}" "\${DNS_OLD_CHECKSUM}" "\${DNS_NEW_CHECKSUM}"

if [[ "\${APPLY}" == "true" ]]; then
  scan_drift after
  echo "known CAASM checksum drift fixed safely"
else
  echo "dry-run only: pass --apply to update allowed drift rows"
fi

REMOTE
}

run_remote_fix() {
  local params_file command_id status output err
  params_file="$(mktemp "${TMPDIR:-/tmp}/checksum-fix-params.XXXXXX")"
  trap 'rm -f "${params_file}"' RETURN

  remote_fix_script | python3 -c 'import json, sys; print(json.dumps({"commands": [sys.stdin.read()]}))' > "${params_file}"

  info "sending checksum fix command to ${INSTANCE_ID}"
  command_id="$(aws_cli ssm send-command \
    --instance-ids "${INSTANCE_ID}" \
    --document-name AWS-RunShellScript \
    --comment "security-center CAASM checksum drift fix" \
    --parameters "file://${params_file}" \
    --query 'Command.CommandId' \
    --output text)"

  info "SSM command id: ${command_id}"
  while true; do
    status="$(aws_cli ssm get-command-invocation \
      --command-id "${command_id}" \
      --instance-id "${INSTANCE_ID}" \
      --query 'Status' \
      --output text 2>/dev/null || true)"
    case "${status}" in
      Success|Cancelled|TimedOut|Failed|Cancelling)
        break
        ;;
      *)
        info "SSM command ${command_id} status=${status:-Pending}"
        sleep 5
        ;;
    esac
  done

  output="$(aws_cli ssm get-command-invocation --command-id "${command_id}" --instance-id "${INSTANCE_ID}" --query 'StandardOutputContent' --output text || true)"
  err="$(aws_cli ssm get-command-invocation --command-id "${command_id}" --instance-id "${INSTANCE_ID}" --query 'StandardErrorContent' --output text || true)"
  [[ -n "${output}" && "${output}" != "None" ]] && printf '%s\n' "${output}"
  [[ -n "${err}" && "${err}" != "None" ]] && printf '%s\n' "${err}" >&2

  [[ "${status}" == "Success" ]] || die "remote checksum fix failed with SSM status=${status}"
}

continue_deploy() {
  if [[ "${CONTINUE_DEPLOY}" != "true" ]]; then
    info "checksum step complete. Re-run build-deploy with --skip-upload --skip-build when ready."
    return 0
  fi

  if [[ "${APPLY}" != "true" ]]; then
    die "--continue-deploy requires --apply"
  fi

  info "resuming build-deploy with --skip-upload --skip-build"
  local -a cmd=("${ENTERPRISE_SH}" build-deploy --env-file "${ENV_FILE}" --instance-id "${INSTANCE_ID}" --source-dir "${ROOT_DIR}" --resolve-ssm --non-interactive --skip-upload --skip-build)
  [[ -n "${AWS_PROFILE_NAME}" ]] && cmd+=(--profile "${AWS_PROFILE_NAME}")
  [[ -n "${AWS_REGION_NAME}" ]] && cmd+=(--region "${AWS_REGION_NAME}")
  [[ -n "${S3_BUCKET}" ]] && cmd+=(--s3-bucket "${S3_BUCKET}")
  [[ -n "${S3_PREFIX}" ]] && cmd+=(--s3-prefix "${S3_PREFIX}")
  "${cmd[@]}"
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --env-file)
      ENV_FILE="$2"
      shift 2
      ;;
    --profile)
      AWS_PROFILE_NAME="$2"
      shift 2
      ;;
    --region)
      AWS_REGION_NAME="$2"
      shift 2
      ;;
    --instance-id)
      INSTANCE_ID="$2"
      shift 2
      ;;
    --remote-dir)
      REMOTE_DIR="$2"
      shift 2
      ;;
    --s3-bucket)
      S3_BUCKET="$2"
      shift 2
      ;;
    --s3-prefix)
      S3_PREFIX="$2"
      shift 2
      ;;
    --backup-s3-bucket)
      DB_BACKUP_S3_BUCKET="$2"
      shift 2
      ;;
    --backup-s3-prefix)
      DB_BACKUP_S3_PREFIX="$2"
      shift 2
      ;;
    --backup-s3-region)
      DB_BACKUP_S3_REGION="$2"
      shift 2
      ;;
    --apply)
      APPLY="true"
      shift
      ;;
    --continue-deploy)
      CONTINUE_DEPLOY="true"
      shift
      ;;
    --skip-backup)
      SKIP_BACKUP="true"
      shift
      ;;
    --help|-h)
      usage
      exit 0
      ;;
    *)
      die "unknown argument: $1"
      ;;
  esac
done

cd "${ROOT_DIR}"
require_inputs
if [[ "${APPLY}" == "true" ]]; then
  run_enterprise_backup
else
  info "read-only mode: pass --apply to update the ledger"
fi
run_remote_fix
continue_deploy