#!/bin/bash

# Configuration
KEYSPACE="dev_keyspace_1"
LOCAL_DIR="data/dumps"
NODES=("cass-dev-01" "cass-dev-02" "cass-dev-03")
SSH_USER="alfredo.oliviero"
DATA_PATH="/data"
SNAPSHOT_PATH="${DATA_PATH}/${KEYSPACE}"
NODETOOL="/home/alfredo.oliviero/apache-cassandra-4.1.3/bin/nodetool"
CQLSH="/home/alfredo.oliviero/apache-cassandra-4.1.3/bin/cqlsh"
CQLSH_IP="10.1.28.100"
DUMP_TAG="dump_docker"

NODE=cass-dev-01
NODE_NAME=node1

# Create directory for dumps
mkdir -p $LOCAL_DIR

# Function to log messages
log() {
    local MESSAGE="$1"
    echo -e "$MESSAGE" | tee -a "$LOCAL_DIR/dump.log"
}

# Function to extract table name from table directory
get_table_name() {
    local TABLE_PATH="$1"
    local TABLE_DIR=$(dirname $(dirname $TABLE_PATH))
    local TABLE_DIR_NAME=$(basename $TABLE_DIR)
    local TABLE_NAME=$(echo $TABLE_DIR_NAME | sed 's/-[a-f0-9]\{32\}$//')
    echo $TABLE_NAME
}

# Function to create snapshot and copy it locally
snapshot_and_copy() {
    NODE=$1
    NODE_NAME=$2
    rm -rf $LOCAL_DIR/$NODE_NAME

    log "Removing old snapshots on $NODE"
    ssh $SSH_USER@$NODE "sudo $NODETOOL clearsnapshot -t $DUMP_TAG -- $KEYSPACE"

    log "Creating snapshot on $NODE"
    ssh $SSH_USER@$NODE "sudo $NODETOOL snapshot -t $DUMP_TAG $KEYSPACE"

    # Find the snapshot path and copy it if it exists
    TABLES=$(ssh $SSH_USER@$NODE "find $SNAPSHOT_PATH -name $DUMP_TAG")
    for TABLE_PATH in $TABLES; do
        TABLE_NAME=$(get_table_name $TABLE_PATH)
        LOCAL_TABLE_DIR="${LOCAL_DIR}/${NODE_NAME}/${TABLE_NAME}"

        log ">> table path $TABLE_PATH\n>> table name $TABLE_NAME\n>> local table dir $LOCAL_TABLE_DIR; "

        mkdir -p $LOCAL_TABLE_DIR
        
        log "Copying snapshot from $NODE:$TABLE_PATH to $LOCAL_TABLE_DIR/$DUMP_TAG"
        rsync -C -r $SSH_USER@$NODE:$TABLE_PATH/ $LOCAL_TABLE_DIR/
    done
}

# Function to handle script interruption
cleanup() {
    log "Script interrupted. Cleaning up..."
    # Add any additional cleanup commands here
    exit 1
}

# Set trap to catch signals and run cleanup
trap cleanup SIGINT SIGTERM

log "Starting snapshot creation for keyspace $KEYSPACE"

# Export keyspace schema from the first node
log "Exporting keyspace schema for $KEYSPACE from ${NODES[0]}"
mkdir -p ${LOCAL_DIR}/schema

ssh $SSH_USER@${NODES[0]} "sudo $CQLSH $CQLSH_IP -e 'DESCRIBE KEYSPACE $KEYSPACE;'" > "${LOCAL_DIR}/schema/${KEYSPACE}_schema.cql"

# Perform snapshot and copy for each node in parallel
NODE_INDEX=1
for NODE in "${NODES[@]}"; do
    snapshot_and_copy $NODE "node$NODE_INDEX" &
    NODE_INDEX=$((NODE_INDEX + 1))
done

# Wait for all background jobs to finish
wait

log "Backup completed."

# Display sizes of dumps
log "Total size of each node dump directory:"
for NODE_INDEX in {1..3}; do
    du -sh "${LOCAL_DIR}/node${NODE_INDEX}"
done

log "Total size of all dump directories:"
du -sh ${LOCAL_DIR}/*