#!/bin/bash

YELLOW='\033[0;33m'
GREEN='\033[0;32m'
RED='\033[0;31m'
NC='\033[0m'

format_size() {
    local bytes=$1
    local kb=1024
    local mb=$((1024 * 1024))
    local gb=$((1024 * 1024 * 1024))
    if [ "$bytes" -gt $((10 * gb)) ]; then
        echo "$(( (bytes + gb - 1) / gb )) GB"
    elif [ "$bytes" -gt $((10 * mb)) ]; then
        echo "$(( (bytes + mb - 1) / mb )) MB"
    else
        echo "$(( (bytes + kb - 1) / kb )) KB"
    fi
}

if [ "$USER" != "root" ] && [ "`whoami`" != "root" ]; then
    echo -e "${RED}This script needs sudo :(${NC}"
    exit 1
fi

psql_docker_name=""

positional_args=()
while [[ $# -gt 0 ]]; do
    case $1 in
        --psql-docker-name)
            psql_docker_name="$2"
            shift 2
            ;;
        *)
            positional_args+=("$1")
            shift
            ;;
    esac
done
set -- "${positional_args[@]}"

if [ "$#" -ne 10 ] && [ "$#" -ne 11 ]; then
    echo -e "${YELLOW}Usage: $0 [--psql-docker-name <container>] <host> <user> <password> <psql_db_name_from> <psql_db_name_to> <psql_port> <ch_def_user_password> <ch_db_name_from> <ch_db_name_to> <ch_port> [<license>]${NC}"
    echo -e "${YELLOW}  --psql-docker-name <container>  PostgreSQL runs in the specified Docker container (uses docker exec for pg_dump/pg_restore)${NC}"
    exit 1
fi

host=$1
user=$2
password=$3
psql_db_name_from=$4
psql_db_name_to=$5
psql_port=$6
ch_def_user_password=$7
ch_db_name_from=$8
ch_db_name_to=$9
ch_port=${10}
license=${11}

if [ "$psql_db_name_from" == "$psql_db_name_to" ]; then
    read -p "Warning! PostgreSQL migration \"$psql_db_name_from\" into itself. PostgreSQL database \"$psql_db_name_from\" will be overwritten. Continue? (y/N): " confirm
    if [[ ! $confirm =~ ^[Yy]$ ]]; then
        echo "Canceled"
        exit 1
    fi
fi

if [ "$ch_db_name_from" == "$ch_db_name_to" ]; then
    read -p "Warning! ClickHouse migration \"$ch_db_name_from\" into itself. ClickHouse database \"$ch_db_name_to\" will be overwritten. Continue? (y/N): " confirm
    if [[ ! $confirm =~ ^[Yy]$ ]]; then
        echo "Canceled"
        exit 1
    fi
fi
echo "Continue"

if [[ -z "$license" ]]; then
    echo "Migrate with license"
    license=1
else
    if [ $license == "true" ] || [ $license == 1 ]; then
        echo "Migrate with license"
    else
        echo "Migrate without license"
    fi
fi

if [ -n "$psql_docker_name" ]; then
    pg_size_bytes=$(docker exec -e PGPASSWORD="$password" "$psql_docker_name" psql -U "$user" -d "$psql_db_name_from" -t -A -c "SELECT pg_database_size('$psql_db_name_from');" | tr -d '[:space:]')
else
    pg_size_bytes=$(PGPASSWORD="$password" psql -U "$user" -d "$psql_db_name_from" -h "$host" -p "$psql_port" -t -A -c "SELECT pg_database_size('$psql_db_name_from');" | tr -d '[:space:]')
fi

if ! [[ "$pg_size_bytes" =~ ^[0-9]+$ ]]; then
    echo -e "${RED}ERROR:${NC} failed to fetch PostgreSQL DB size for '$psql_db_name_from' (check credentials, host, port, db name)"
    exit 1
fi

ch_db_exists=$(clickhouse-client --host="$host" --port="$ch_port" --user=default --password="$ch_def_user_password" --query="SELECT count() FROM system.databases WHERE name = '$ch_db_name_from';" | tr -d '[:space:]')

if ! [[ "$ch_db_exists" =~ ^[0-9]+$ ]]; then
    echo -e "${RED}ERROR:${NC} cannot reach ClickHouse (check credentials, host, port)"
    exit 1
fi

if [ "$ch_db_exists" -eq 0 ]; then
    echo -e "${RED}ERROR:${NC} ClickHouse database '$ch_db_name_from' does not exist"
    exit 1
fi

ch_size_bytes=$(clickhouse-client --host="$host" --port="$ch_port" --user=default --password="$ch_def_user_password" --query="SELECT toUInt64(coalesce(sum(total_bytes), 0)) FROM system.tables WHERE database = '$ch_db_name_from';" | tr -d '[:space:]')

if ! [[ "$ch_size_bytes" =~ ^[0-9]+$ ]]; then
    echo -e "${RED}ERROR:${NC} failed to fetch ClickHouse DB size for '$ch_db_name_from'"
    exit 1
fi

buffer_bytes=$((512 * 1024 * 1024))
required_bytes=$(( pg_size_bytes * 3 / 2 + ch_size_bytes * 2 + buffer_bytes ))
available_bytes=$(df -B1 / | awk 'NR==2 {print $4}')

echo -e "Source DB sizes: PostgreSQL '${psql_db_name_from}' = $(format_size "$pg_size_bytes"), ClickHouse '${ch_db_name_from}' = $(format_size "$ch_size_bytes")"
echo -e "Estimated required free space: $(format_size "$required_bytes"), available: $(format_size "$available_bytes")"

if [ "$available_bytes" -lt "$required_bytes" ]; then
    echo -e "${RED}ERROR:${NC} requires minimum $(format_size "$required_bytes") free space on the disk"
    exit 1
fi

./src/dump.sh "$host" "$user" "$password" "$psql_docker_name" "$psql_db_name_from" "$psql_port" "$ch_def_user_password" "$ch_db_name_from" "$ch_port" "$license"
exit_code_dump=$?

if [ $exit_code_dump -eq 0 ]; then
    echo -e "\t${GREEN}Dumping DB done <3${NC}"
else
    echo -e "${RED}ERROR:${NC} dumping DB error!"
    exit 1
fi

./src/restore.sh "$host" "$user" "$password" "$psql_docker_name" "$psql_db_name_to" "$psql_port" "$ch_def_user_password" "$ch_db_name_from" "$ch_db_name_to" "$ch_port" "$license"
exit_code_restore=$?

if [ $exit_code_restore -eq 0 ]; then
    echo -e "\t${GREEN}Restoring DB done <3${NC}"
else
    echo -e "${RED}ERROR:${NC} restoring DB error!"
    exit 1
fi

rm -rf dumps
rm -rf restores
rm -rf /backups
