#!/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 file
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

# Function to retry SSH connection to bastion host
retry_ssh_connection() {
  local bastion_ip=$1
  for attempt in {1..5}; do
    log_message "Attempting to connect to bastion host (attempt $attempt)..."
    if ssh -tt -o ServerAliveInterval=60 -o ServerAliveCountMax=5 -o StrictHostKeyChecking=no -i "$key_file" ec2-user@"$bastion_ip" "echo 'Bastion is ready'"; then
      return
    fi
    sleep 30
  done
  log_message "Failed to connect to bastion host after multiple attempts. Exiting."
  exit 1
}

# Function to check readiness of private nodes via bastion
wait_for_private_nodes() {
  local bastion_ip=$1
  shift
  local private_nodes=("$@")
  
  log_message "Checking private nodes readiness via bastion host $bastion_ip..."
  ssh -tt -o ServerAliveInterval=60 -o ServerAliveCountMax=5 -o StrictHostKeyChecking=no -i "$key_file" ec2-user@"$bastion_ip" <<EOF
    set -e
    if ! command -v nc &>/dev/null; then
      echo "Installing nc (Netcat)..."
      if [ -f /etc/redhat-release ]; then
        sudo yum install -y nc
      elif [ -f /etc/debian_version ]; then
        sudo apt-get update && sudo apt-get install -y netcat
      else
        echo "Unsupported OS for installing nc. Exiting."
        exit 1
      fi
    fi

    for attempt in {1..10}; do
      all_ready=true
      for node in ${private_nodes[@]}; do
        echo "Checking node \$node readiness (attempt \$attempt)..."
        if ! nc -z -w 5 \$node 22; then
          echo "Node \$node is not ready."
          all_ready=false
        else
          echo "Node \$node is ready."
        fi
      done

      if \$all_ready; then
        echo "All nodes are ready."
        exit 0
      fi

      echo "Retrying node readiness check in 30 seconds..."
      sleep 30
    done

    echo "Some nodes are not ready after multiple attempts."
    exit 1
EOF
}


# Prompt for the number of clusters
read -p "Enter the number of TiDB clusters: " num_clusters

# Validate input
if ! [[ $num_clusters =~ ^[0-9]+$ ]] || [[ $num_clusters -le 0 ]]; then
  log_message "Error: Number of clusters must be a positive integer."
  exit 1
fi

# Export environment variables for Terraform
export TF_VAR_num_clusters=$num_clusters

# Log into AWS
log_message "Logging into AWS..."
aws sso login --profile DBaaS-DevUser-Role-219248915861 || {
  log_message "AWS login failed. Exiting."
  exit 1
}

# Initialize and apply Terraform
log_message "Initializing Terraform..."
terraform init || {
  log_message "Terraform initialization failed. Exiting."
  exit 1
}

terraform apply -auto-approve || {
  log_message "Terraform apply failed. Exiting."
  exit 1
}

# Retrieve cluster IPs
log_message "Retrieving cluster IPs from Terraform output..."
bastion_ips=($(terraform output -json bastion_public_ips | jq -r '.[]')) || {
  log_message "Failed to retrieve bastion IPs."
  exit 1
}

pd_ips_list=($(terraform output -json pd_private_ips | jq -r '.[]')) || {
  log_message "Failed to retrieve PD IPs."
  exit 1
}

tidb_ips_list=($(terraform output -json tidb_private_ips | jq -r '.[]')) || {
  log_message "Failed to retrieve TiDB IPs."
  exit 1
}

tikv_ips_list=($(terraform output -json tikv_private_ips | jq -r '.[]')) || {
  log_message "Failed to retrieve TiKV IPs."
  exit 1
}

tiflash_ips=($(terraform output -json tiflash_private_ips | jq -r '.[]')) || {
  log_message "Failed to retrieve TiFlash IPs."
  exit 1
}

sleep 60

# Wait for bastion host to be ready
bastion_ip=${bastion_ips[0]}  # Assuming single bastion host
retry_ssh_connection "$bastion_ip"

# Check readiness of private nodes via bastion
all_private_nodes=("${pd_ips_list[@]}" "${tidb_ips_list[@]}" "${tikv_ips_list[@]}" "${tiflash_ips[@]}")
wait_for_private_nodes "$bastion_ip" "${all_private_nodes[@]}"

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

  ssh -tt -o ServerAliveInterval=60 -o ServerAliveCountMax=5 -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"
  }
}

# Path to store cluster credentials
credentials_file="cluster_credentials.json"

# Ensure the credentials file exists and restrict permissions
if [[ ! -f "$credentials_file" ]]; then
  echo "[]" > "$credentials_file"
  chmod 600 "$credentials_file"
fi

# Deploy and validate 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]}
  local cluster_output_log="cluster_output_${cluster_idx}.log"

  ssh -tt -o ServerAliveInterval=60 -o ServerAliveCountMax=5 -o StrictHostKeyChecking=no -i "$key_file" ec2-user@"$bastion_ip" <<EOF
    set -e
    # Configure SSH
    mkdir -p ~/.ssh
    if [ ! -f ~/.ssh/id_rsa ]; then
      cp ~/tidb-demo.pem ~/.ssh/id_rsa
      chmod 400 ~/.ssh/id_rsa
    fi

    if [ ! -f ~/.ssh/config ]; then
      echo -e "Host *\n  IdentityFile ~/.ssh/id_rsa\n  StrictHostKeyChecking no" > ~/.ssh/config
      chmod 600 ~/.ssh/config
    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 TiDB cluster
    tiup cluster check ~/topology.yaml --apply --user ec2-user --identity_file ~/.ssh/id_rsa -y
    sleep 30
    tiup cluster deploy tidb-cluster-$cluster_idx $tiup_version ~/topology.yaml --identity_file ~/.ssh/id_rsa -y | tee ~/cluster_output_${cluster_idx}.log
    sleep 60
    echo "Starting cluster tidb-cluster-$cluster_idx..."
    tiup cluster start tidb-cluster-$cluster_idx --init | tee -a ~/cluster_output_${cluster_idx}.log
    sleep 10
    echo "Cluster tidb-cluster-$cluster_idx started successfully."
    exit
EOF

  # Download the log file to the local system
  scp -i "$key_file" ec2-user@"$bastion_ip":~/cluster_output_${cluster_idx}.log "./cluster_output_${cluster_idx}.log"

  # Extract and save the password
  local password=$(grep "The new password is:" "./cluster_output_${cluster_idx}.log" | awk -F': ' '{print $2}' | tr -d '\r' | tail -1)
  if [[ -n "$password" ]]; then
    log_message "Saving credentials for cluster tidb-cluster-$cluster_idx..."
    jq ". += [{\"cluster\": \"tidb-cluster-$cluster_idx\", \"password\": \"$password\"}]" "$credentials_file" > temp.json && mv temp.json "$credentials_file"
  else
    log_message "Warning: Password not found for cluster tidb-cluster-$cluster_idx."
  fi
}
  
# Main logic
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]}

  retry_ssh_connection "$bastion_ip"
  transfer_ssh_key "$bastion_ip"
  deploy_cluster $((cluster_idx + 1)) "$bastion_ip" "${pd_ips[@]}" "${tidb_ips[@]}" "${tikv_ips[@]}" $tiflash_ip || {
    log_message "Cluster deployment failed for cluster $((cluster_idx + 1)). Exiting."
    exit 1
  }
done

log_message "All clusters deployed and configured successfully."