#!/bin/bash

# Exit on any command failure
set -e

# Configuration
log_file="cluster_setup.log"
key_file="./tidb-demo.pem"
tiup_version="v8.5.0"

# Log startup
echo "Cluster setup started at $(date)" > "$log_file"

# Function to log messages
log_message() {
  echo "$1" | tee -a "$log_file"
}

# Function to handle errors
handle_error() {
  log_message "An error occurred on line $1. Check $log_file for details."
  exit 1
}

# Trap errors to call handle_error
trap 'handle_error $LINENO' ERR

# Parse cluster outputs
parse_cluster_outputs() {
  bastion_ips=($(terraform output -json bastion_public_ips | jq -r '.[]'))
  pd_ips_list=($(terraform output -json pd_private_ips | jq -r '.[]'))
  tidb_ips_list=($(terraform output -json tidb_private_ips | jq -r '.[]'))
  tikv_ips_list=($(terraform output -json tikv_private_ips | jq -r '.[]'))
  tiflash_ips=($(terraform output -json tiflash_private_ips | jq -r '.[]'))
}

# Print cluster information
print_cluster_info() {
  local cluster_number=$1
  local bastion_ip=$2
  local pd_ips=("${@:3:3}")
  local tidb_ips=("${@:6:2}")
  local tikv_ips=("${@:8:3}")
  local tiflash_ip=${11}

  log_message "Cluster $cluster_number Information:"
  log_message "Bastion IP: $bastion_ip"
  log_message "PD Node IPs: ${pd_ips[*]}"
  log_message "TiDB Node IPs: ${tidb_ips[*]}"
  log_message "TiKV Node IPs: ${tikv_ips[*]}"
  log_message "TiFlash Node IP: $tiflash_ip"
  log_message
}

# Transfer SSH key if not exists
transfer_ssh_key() {
  local bastion_ip=$1

  ssh -o StrictHostKeyChecking=no -i "$key_file" ec2-user@"$bastion_ip" "test -f ~/tidb-demo.pem" || {
    log_message "Transferring SSH identity file to bastion..."
    scp -o StrictHostKeyChecking=no -i "$key_file" "$key_file" ec2-user@"$bastion_ip":~/tidb-demo.pem
    ssh -o StrictHostKeyChecking=no -i "$key_file" ec2-user@"$bastion_ip" "chmod 400 ~/tidb-demo.pem"
  }
}

# Configure and deploy TiDB cluster
deploy_cluster() {
  local cluster_idx=$1
  local bastion_ip=$2
  local pd_ips=("${@:3:3}")
  local tidb_ips=("${@:6:2}")
  local tikv_ips=("${@:8:3}")
  local tiflash_ip=${11}

  # Assign monitoring, grafana, and alertmanager to the first PD node
  local monitor_ip=${pd_ips[0]}

  ssh -tt -o StrictHostKeyChecking=no -i "$key_file" ec2-user@"$bastion_ip" <<EOF
    set -e

    # Configure SSH
    mkdir -p ~/.ssh
    if [ ! -f ~/.ssh/id_rsa ]; then
      echo "Copying SSH key to ~/.ssh/id_rsa..."
      cp ~/tidb-demo.pem ~/.ssh/id_rsa
      chmod 400 ~/.ssh/id_rsa
    else
      echo "SSH key ~/.ssh/id_rsa already exists. Skipping copy."
    fi

    if [ ! -f ~/.ssh/config ]; then
      echo -e "Host *\n  IdentityFile ~/.ssh/id_rsa\n  StrictHostKeyChecking no" > ~/.ssh/config
      chmod 600 ~/.ssh/config
    else
      echo "SSH configuration file ~/.ssh/config already exists. Skipping configuration."
    fi

    # Create topology file
    cat <<TOPOLOGY > ~/topology.yaml
global:
  user: "ec2-user"
  ssh_port: 22
  deploy_dir: "/tidb-deploy"
  data_dir: "/tidb-data"
pd_servers:
$(for ip in "${pd_ips[@]}"; do echo "  - host: $ip"; done)
tidb_servers:
$(for ip in "${tidb_ips[@]}"; do echo "  - host: $ip"; done)
tikv_servers:
$(for ip in "${tikv_ips[@]}"; do echo "  - host: $ip"; done)
tiflash_servers:
  - host: $tiflash_ip
monitoring_servers:
  - host: $monitor_ip
grafana_servers:
  - host: $monitor_ip
alertmanager_servers:
  - host: $monitor_ip
TOPOLOGY

    # Deploy and start TiDB cluster with retries
    for i in {1..3}; do
      tiup cluster check ~/./topology.yaml --apply --user ec2-user --identity_file ~/.ssh/id_rsa
      tiup cluster deploy tidb-cluster-$cluster_idx $tiup_version ~/topology.yaml --identity_file ~/.ssh/id_rsa -y && break || sleep 10
      echo "Retrying deployment ($i/3)..."
    done

    tiup cluster start tidb-cluster-$cluster_idx || {
      echo "Failed to start TiDB cluster. Check logs for details." | tee -a ~/cluster_start.log
      exit 1
    }
    exit
EOF
}

# Main logic
parse_cluster_outputs

num_clusters=${#bastion_ips[@]}
log_message "Detected $num_clusters clusters."

for cluster_idx in $(seq 0 $((num_clusters - 1))); do
  bastion_ip=${bastion_ips[$cluster_idx]}
  pd_ips=(${pd_ips_list[@]:$((cluster_idx * 3)):3})
  tidb_ips=(${tidb_ips_list[@]:$((cluster_idx * 2)):2})
  tikv_ips=(${tikv_ips_list[@]:$((cluster_idx * 3)):3})
  tiflash_ip=${tiflash_ips[$cluster_idx]}

  print_cluster_info $((cluster_idx + 1)) $bastion_ip "${pd_ips[@]}" "${tidb_ips[@]}" "${tikv_ips[@]}" $tiflash_ip
  transfer_ssh_key "$bastion_ip"
  deploy_cluster $((cluster_idx + 1)) "$bastion_ip" "${pd_ips[@]}" "${tidb_ips[@]}" "${tikv_ips[@]}" $tiflash_ip
done

log_message "All clusters deployed and configured successfully."