#!/usr/bin/env bash
set -Eeuo pipefail

readonly ACCOUNT_ID=804837308162
readonly AWS_PROFILE=dev-offline-afterworld
readonly AWS_REGION=us-west-2

usage() {
  echo "usage: $0 prepare <source-development|parallel|private|artifacts|security|gcp-regional> <evidence-dir> | cleanup <evidence-dir>" >&2
  exit 2
}

future_ttl() {
  if date -u -v+10H +%Y-%m-%dT%H:%M:%SZ >/dev/null 2>&1; then
    date -u -v+10H +%Y-%m-%dT%H:%M:%SZ
  else
    date -u -d '+10 hours' +%Y-%m-%dT%H:%M:%SZ
  fi
}

case_defaults() {
  TILT_MODE=up
  case "$1" in
    source-development) CASE_CODE=r; SERVICES=config-center; REPOSITORIES=config-center ;;
    parallel) CASE_CODE=p; SERVICES=config-center,config-center-regional,cluster-ng-server; REPOSITORIES=config-center,cluster-service-ng ;;
    private)  CASE_CODE=g; SERVICES=tempo-proxy,tempo-wrapper; REPOSITORIES=tidb-management-service ;;
    artifacts) CASE_CODE=a; SERVICES=; REPOSITORIES= ;;
    security) CASE_CODE=s; SERVICES=; REPOSITORIES=; TILT_MODE=generate-only ;;
    gcp-regional) CASE_CODE=u; SERVICES=; REPOSITORIES= ;;
    *) usage ;;
  esac
}

clone_repository() {
  local name=$1 destination=$2 repository
  case "$name" in
    config-center) repository=tidbcloud/config-center ;;
    cluster-service-ng) repository=tidbcloud/cluster-service-ng ;;
    tidb-management-service) repository=tidbcloud/tidb-management-service ;;
    *) echo "unsupported source repository: $name" >&2; return 1 ;;
  esac
  gh repo clone "$repository" "$destination" -- --depth=1
}

cp9_aws() {
  local cp9_root=$1 owner=$2 session=$3 credentials access_key secret_key session_token
  shift 3
  credentials=$(bash "$cp9_root/internal/aws/cp9-credential-process.sh" \
    "$AWS_PROFILE" "arn:aws:iam::$ACCOUNT_ID:role/cp9-developer" \
    "$session" 3600 "$owner")
  access_key=$(jq -er '.AccessKeyId' <<<"$credentials")
  secret_key=$(jq -er '.SecretAccessKey' <<<"$credentials")
  session_token=$(jq -er '.SessionToken' <<<"$credentials")
  AWS_ACCESS_KEY_ID=$access_key \
    AWS_SECRET_ACCESS_KEY=$secret_key \
    AWS_SESSION_TOKEN=$session_token \
    aws --region "$AWS_REGION" "$@"
}

write_environment() {
  local path=$1
  : >"$path"
  chmod 600 "$path"
  for name in ORIGIN_CP9 CP9_ROOT CP9_CONFIG RUN_DIR SESSION_NAME SOURCE_ROOT \
    SOURCE_FILE PROBE_DIR CLEANUP_AUTHORITY; do
    printf 'export %s=%q\n' "$name" "${!name:-}" >>"$path"
  done
}

prepare() {
  local selected_case=$1 evidence_dir=$2 caller account owner subnet="" ttl key_name private_key key_name_yaml private_key_yaml
  local work_root cp9_root tiltfile config_path env_path head account_contract
  local certificate_arn dns_zone_id dns_zone_name dedicated_config=""
  case_defaults "$selected_case"

  key_name=${CP9_UI_API_AWS_KEY_PAIR:-}
  private_key=${CP9_UI_API_SSH_PRIVATE_KEY:-}
  if test -z "$key_name" || test -z "$private_key"; then
    echo 'CP9_UI_API_AWS_KEY_PAIR and CP9_UI_API_SSH_PRIVATE_KEY are required together; CP9 does not create or import EC2 keys' >&2
    return 1
  fi
  if test ! -f "$private_key" || test ! -r "$private_key"; then
    echo "CP9_UI_API_SSH_PRIVATE_KEY must name a readable private-key file" >&2
    return 1
  fi
  key_name_yaml=$(jq -Rn --arg value "$key_name" '$value')
  private_key_yaml=$(jq -Rn --arg value "$private_key" '$value')

  ORIGIN_CP9=$(git rev-parse --show-toplevel)
  account_contract=$(jq -ce --arg account "$ACCOUNT_ID" '
    .include[] |
    select(.account_id == $account) |
    select(
      (.dns_zone_id | type) == "string" and (.dns_zone_id | test("^Z[A-Z0-9]+$")) and
      .dns_zone_name == ($account + ".cp9.dev.aws.tidbcloud.com") and
      (.certificate_arn | type) == "string"
    )
  ' "$ORIGIN_CP9/.github/cp9-cleanup-accounts.json") || {
    echo "account $ACCOUNT_ID lacks an initialized child DNS zone/certificate contract in .github/cp9-cleanup-accounts.json" >&2
    return 1
  }
  certificate_arn=$(jq -er '.certificate_arn' <<<"$account_contract")
  dns_zone_id=$(jq -er '.dns_zone_id' <<<"$account_contract")
  dns_zone_name=$(jq -er '.dns_zone_name' <<<"$account_contract")
  head=$(git rev-parse HEAD)
  mkdir -p "$evidence_dir"
  RUN_DIR=$(cd "$evidence_dir" && pwd)
  SESSION_NAME="c9${CASE_CODE}$(date -u +%m%d%H%M%S)$(printf '%02d' "$((RANDOM % 100))")"
  ttl=$(future_ttl)

  caller=$(aws sts get-caller-identity --profile "$AWS_PROFILE" --output json)
  account=$(jq -er '.Account' <<<"$caller")
  test "$account" = "$ACCOUNT_ID" || {
    echo "$AWS_PROFILE resolved account $account; expected $ACCOUNT_ID" >&2
    return 1
  }
  owner=$(jq -er '.Arn | split("/") | last | ascii_downcase' <<<"$caller")
  work_root=$(mktemp -d "${TMPDIR:-/tmp}/${SESSION_NAME}.XXXXXX")
  cp9_root="$work_root/controlplane9"
  config_path="$RUN_DIR/cp9.yaml"
  jq -n \
    --arg case "$selected_case" \
    --arg account "$ACCOUNT_ID" \
    --arg region "$AWS_REGION" \
    --arg profile "$AWS_PROFILE" \
    --arg session "$SESSION_NAME" \
    --arg owner "$owner" \
    --arg ttl "$ttl" \
    --arg key_pair "$key_name" \
    --arg ssh_private_key "$private_key" \
    --arg subnet "$subnet" \
    --arg origin_cp9 "$ORIGIN_CP9" \
    --arg cp9_root "$cp9_root" \
    --arg config "$config_path" \
    --arg work_root "$work_root" \
    --arg head "$head" \
    '{case:$case,account:$account,region:$region,profile:$profile,session:$session,
      owner:$owner,ttl:$ttl,key_pair:$key_pair,ssh_private_key:$ssh_private_key,subnet:$subnet,origin_cp9:$origin_cp9,
      cp9_root:$cp9_root,config:$config,work_root:$work_root,head:$head}' \
    >"$RUN_DIR/default-804-session.json"
  rollback_prepare() {
    local status=$?
    trap - ERR INT TERM
    cleanup "$RUN_DIR" || true
    set -e
    return "$status"
  }
  rollback_prepare_signal() {
    local status=$1
    trap - ERR INT TERM
    cleanup "$RUN_DIR" || true
    exit "$status"
  }
  trap rollback_prepare ERR
  trap 'rollback_prepare_signal 130' INT
  trap 'rollback_prepare_signal 143' TERM

  git -C "$ORIGIN_CP9" worktree add --detach "$cp9_root" "$head"
  if test -n "$REPOSITORIES"; then
    IFS=, read -r -a repositories <<<"$REPOSITORIES"
    for repository in "${repositories[@]}"; do
      clone_repository "$repository" "$work_root/$repository"
    done
  fi

  subnet=$(cp9_aws "$cp9_root" "$owner" "$SESSION_NAME" \
    ec2 describe-subnets \
    --filters Name=state,Values=available \
    --output json \
    | "$cp9_root/spec-tests/scripts/select-public-exposure-subnet.py")
  test -n "$subnet" && test "$subnet" != None || {
    echo "account $ACCOUNT_ID has no ALB-compatible public subnet in $AWS_REGION" >&2
    return 1
  }
  jq --arg subnet "$subnet" '.subnet = $subnet' \
    "$RUN_DIR/default-804-session.json" >"$RUN_DIR/default-804-session.json.tmp"
  mv "$RUN_DIR/default-804-session.json.tmp" "$RUN_DIR/default-804-session.json"

  tiltfile="$cp9_root/personas/dev-full/runtime/.cp9-spec-$SESSION_NAME.Tiltfile"
  cp "$cp9_root/personas/dev-full/runtime/example.Tiltfile" "$tiltfile"
  python3 - "$tiltfile" "$SERVICES" <<'PY'
import pathlib, re, sys

path = pathlib.Path(sys.argv[1])
services = [service for service in sys.argv[2].split(",") if service]
text = path.read_text()
replacement = "source_build_services = [\n" + "".join(
    f'    "{service}",\n' for service in services
) + "]"
updated, count = re.subn(
    r"source_build_services = \[.*?\n\]",
    replacement,
    text,
    count=1,
    flags=re.DOTALL,
)
assert count == 1
path.write_text(updated)
PY

  if test "$selected_case" = gcp-regional; then
    dedicated_config='dedicated:
  gcp:
    enabled: true
    seed_project_id: pingcap-gardener
    dataplane_project_id: pingcap-gardener
    region: us-west1
    zone: us-west1-a'
  fi

  cat >"$config_path" <<YAML
apiVersion: cp9.tidbcloud.com/v1alpha2
kind: Session
name: $SESSION_NAME
persona: dev-full
$dedicated_config
host:
  mode: remote-ec2-k3d
  ttl: 10h
  ec2:
    instance_type: m7i.xlarge
    key_pair: $key_name_yaml
    ssh_user: ec2-user
    ssh_private_key: $private_key_yaml
    subnet_id: $subnet
    root_volume_size_gib: 160
  exposure:
    enabled: true
    zone_id: $dns_zone_id
    zone_name: $dns_zone_name
    certificate_arn: $certificate_arn
target:
  cloud: aws
  region: $AWS_REGION
  aws:
    iam:
      source_profile: $AWS_PROFILE
      role_arn: arn:aws:iam::$ACCOUNT_ID:role/cp9-developer
terraform:
  repo: $cp9_root
  ref: $head
  path: terraform/aws/cp9-substrate
  var_file: terraform/aws/cp9-substrate/premium.tfvars
tilt:
  mode: $TILT_MODE
  file: $tiltfile
services:
  github_app_client_id: Iv23liKDo1S8LDZ07CZU
  source_build_parallelism: 2
YAML

  CP9_ROOT=$cp9_root
  CP9_CONFIG=$config_path
  CLEANUP_AUTHORITY=${CLEANUP_AUTHORITY:-"selected execution instruction for this spec or tools check"}
  SOURCE_ROOT=
  SOURCE_FILE=
  PROBE_DIR=
  case "$selected_case" in
    source-development)
      SOURCE_ROOT="$work_root/config-center"
      PROBE_DIR="$SOURCE_ROOT/cmd/server"
      SOURCE_FILE="$SOURCE_ROOT/cmd/server/main.go"
      ;;
    parallel) SOURCE_ROOT="$work_root" ;;
    private) SOURCE_ROOT="$work_root/tidb-management-service" ;;
  esac
  env_path="$RUN_DIR/defaults.env"
  write_environment "$env_path"
  trap - ERR INT TERM
  printf 'source %q\n' "$env_path"
}

cleanup() {
  local evidence_dir=$1 metadata cp9_root config work_root origin_cp9 session
  local host_json remote_state_dir remote_tilt_log
  local destroy_status=not_started worktree_status=not_started
  metadata="$evidence_dir/default-804-session.json"
  test -f "$metadata" || return 0
  cp9_root=$(jq -er '.cp9_root' "$metadata")
  config=$(jq -er '.config' "$metadata")
  work_root=$(jq -er '.work_root' "$metadata")
  origin_cp9=$(jq -er '.origin_cp9' "$metadata")
  session=$(jq -er '.session' "$metadata")

  set +e
  host_json="$cp9_root/.cp9/$session/host.json"
  if test -x "$cp9_root/skills/get-started/scripts/session-ssh.sh" && test -f "$host_json"; then
    remote_state_dir=$(jq -r '.remote.remote_state_dir // empty' "$host_json")
    if test -n "$remote_state_dir"; then
      remote_tilt_log="$remote_state_dir/controller/session/tilt-up.log"
      (
        umask 077
        CP9_REPO_ROOT="$cp9_root" \
          "$cp9_root/skills/get-started/scripts/session-ssh.sh" \
          "$session" tail -n 5000 -- "$remote_tilt_log" \
          >"$evidence_dir/remote-tilt-up.log" \
          2>"$evidence_dir/remote-tilt-up.stderr"
      )
    fi
  fi
  if test -x "$cp9_root/bin/cp9" && test -d "$cp9_root/.cp9/$(jq -r '.session' "$metadata")"; then
    (cd "$cp9_root" && bin/cp9 destroy config="$config") >"$evidence_dir/destroy.log" 2>&1
    test $? -eq 0 && destroy_status=pass || destroy_status=fail
  else
    destroy_status=not_needed
  fi
  if test "$destroy_status" = fail; then
    worktree_status=retained
  else
    if test -d "$cp9_root"; then
      git -C "$origin_cp9" worktree remove --force "$cp9_root" \
        >"$evidence_dir/remove-worktree.log" 2>&1
      test $? -eq 0 && worktree_status=pass || worktree_status=fail
    else
      worktree_status=not_needed
    fi
  fi
  if test "$worktree_status" = pass || test "$worktree_status" = not_needed; then
    rm -rf "$work_root"
  fi
  jq -n \
    --arg destroy "$destroy_status" \
    --arg worktree "$worktree_status" \
    '{destroy:$destroy,worktree:$worktree}' \
    >"$evidence_dir/cleanup-receipt.json"
  test "$destroy_status" != fail && test "$worktree_status" != fail
}

test $# -ge 1 || usage
command=$1
shift
case "$command" in
  prepare)
    test $# -eq 2 || usage
    prepare "$1" "$2"
    ;;
  cleanup)
    test $# -eq 1 || usage
    cleanup "$1"
    ;;
  *) usage ;;
esac