#!/bin/zsh

#TODO list
# 1. Скрипт не определяет, что нужно остановиться
# 2. Нужно автоматически определить, если копирование уже запущено то предложить остановить и вывести процесс резервирования
# Нет проверки на верность введённых параметров
# смотреть нужно сразу на всей системе. Не будем запускать несколько параллельных резервирований.


backup_server="192.168.0.120" # Сервер для резервирования
backup_user="root" # Пользователь на сервере
inc_snapshot="manual-$(date +%Y%m%d)" # Новый создаваемый снисмок
incremental=false
silent=false
stop=false
progress=true
progress_only=false
check=true
LOGFILE="log_bak_${inc_snapshot}.log"
unset local_dataset
unset remote_dataset

echo -ne "Обработка аргументов...\033[2K\r"

while [[ "$#" -gt 0 ]]; do
    case $1 in
		-rs|--remote-server) backup_server="$2"; shift ;;
		-ru|--remote-user) backup_user="$2"; shift ;;
		-s|--snapshot) inc_snapshot="$2"; shift ;;
		-ld|--local-dataset) local_dataset="$2"; shift ;;
		-rd|--remote-dataset) remote_dataset="$2"; shift ;;
		-l|--log-file) LOGFILE="$2"; shift ;;
		-i|--incremental) incremental=true ;;
		--no-progress) progress=false ;;
		--progress-only) progress_only=true ;;
		--no-check) check=false ;;
		--stop) stop=true ;;
		-h|--help) echo "Использование: ./zfs_send.zsh [OPTIONS]";
		           echo "-rs | --remote-server <ip/dns>";
		           echo "-гu | --remote-user <user>";
		           echo "-s  | --snapshot <snapshot_name>";
		           echo "-ld | --local-dataset <local/dataset>";
		           echo "-rd | --remote-dataset <remote/dataset>";
		           echo "-l  | --log-file <path/filename>";
		           echo "      --no-progress";
		           echo "      --progress-only";
		           echo "      --no-check";
		           echo "      --stop";
		           echo "-i  | --incremental" ;
					exit 0;;
		*) echo "Неизвестные параметры: $1"; exit 1 ;;
	esac
	shift
done


# Функция для записи сообщения в файл и на экран
echo_log () {
    local text=$1
    echo $text
    echo $text >> "$LOGFILE"
}


# Функция для определения SSH-ключа через конфигурацию
find_ssh_key() {
    local key

    # Получаем IdentityFile из конфигурации
    identity_files=$(ssh -G "$backup_user@$backup_server" 2>/dev/null | awk '/^IdentityFile / {print $2}')

    # Проверяем каждый IdentityFile из конфига
    for file in $identity_files; do
        if [[ -f "$file" ]]; then
            echo "$file"
            return
        fi
    done

    # Если ничего не найдено, ищем стандартные ключи
    for candidate in ~/.ssh/id_rsa ~/.ssh/id_ecdsa ~/.ssh/id_ed25519; do
        if [[ -f "$candidate" ]]; then
            echo "$candidate"
            return
        fi
    done

    # Если ключей нет, создаем новый
	if [ ! -f ~/.ssh/id_ed25519 ]; then
	    ssh-keygen -t ed25519 -N "" -f ~/.ssh/id_ed25519
	    # echo "Создаём новый ключ..."
	    ssh-keygen -t ed25519 -N "" -f ~/.ssh/id_ed25519 >&2
	    echo "~/.ssh/id_ed25519"
	else
		echo "Ключей для этого сервера не найдено. Ключ по умолчанию ~/.ssh/id_ed25519 уже существует. Создание нового ключа прервано"
	fi
}

# Определение ключа
key_path=$(find_ssh_key)
key_pub_path="${key_path}.pub"

check_ssh_connection() {
	ssh -o "BatchMode=yes" -o "ConnectTimeout=5" "$backup_user@$backup_server" exit 2>/dev/null
	return $?
}


setup_ssh_key() {
	key_content=$(cat "$key_pub_path")
	echo "Установка SSH-ключа на $backup_user@$backup_server..."

	# Передача публичного ключа вручную
    ssh $backup_user@$backup_server "mkdir -p ~/.ssh && echo '$key_content' >> ~/.ssh/authorized_keys && chmod 600 ~/.ssh/authorized_keys && chmod 700 ~/.ssh"

	echo "
host $backup_server
hostname $backup_server
user $backup_user
identityFile $key_path
identitiesOnly yes
" >> ~/.ssh/config

	if [[ $? -ne 0 ]]; then
		echo "Ошибка: Не удалось установить SSH-ключ. Попробуйте прописать его на сервер вручную."
		echo "----------"
		echo $key_content
		echo "----------"
		exit 1
	else
		echo "SSH-ключ успешно установлен!"
	fi
}




# !
# Если нужно вывести только прогресс, то пропускаем все этапы подготовки и резервирования
if [[ $progress_only = false ]]; then




# Завершение запущенного процесса резервирования
if [[ $stop = true ]]; then
	pids=$(ps -t $(tty | sed 's|/dev/||') -o pid,args | grep -e zfs_send.zsh -e "zfs send" | grep -v -- '--stop') # Список PIDs с командой запуска 

	if [[ $pids = "" ]]; then
		echo "Нет запущенного процесса резервирования в этом терминале."
		read "srch?Икать во всех процессах системы? (y/N) "
		if [[ $srch = y ]]; then
			pids=$(ps -o pid,args | grep -e zfs_send.zsh -e "zfs send" | grep -v -- '--stop' )
		fi
	fi

	# Подтверждение
	if [[ $pids = "" ]]; then
		echo "Нет запущенного процесса резервирования"
	else
		echo "Будут остановлены следующие процессы:"
		echo $pids
		read "work? Продолжать (y/N)"
		if [[ $work = y ]]; then
			pids=$(echo $pids | awk '{print $1}')
			kill $pids
			echo "ok"
		else
			echo "Операция прервана"
		fi
	fi

	exit 0
fi


# Проверка соединения
echo "Резервирование на   $backup_user@$backup_server"
if check_ssh_connection; then
	echo "Connection OK"
else
    echo "SSH-ключ не обнаружен на сервере. Начинаем настройку..."
    setup_ssh_key
    if check_ssh_connection; then
		echo "Connection OK"
	else
        echo "Ошибка: SSH-ключ не работает. Проверьте настройки."
        exit 1
    fi
fi


# ----------
echo "Снимок файловой системы   $inc_snapshot"

case $incremental in
	true) echo "Режим инкрементной копии";;
	false) echo "Режим полной копии";;
esac

echo "Дополнительные опции: -h|--help"

# --------- Собираем данные ----------

echo -ne "Собираю данные...\033[2K\r"
# Получить список датасетов на локальной машине
local_datasets=$(zfs list -o name | grep -v -e ix-app -e boot-pool -e system-data -e jails -e NAME)
local_datasets=(${=local_datasets}) # Конвертирование в массив
# Получить список датасетов на удалённой машине
backup_datasets=$(ssh $backup_user@$backup_server zfs list -o name | grep -v -e ix-app -e boot-pool -e system-data -e jails -e NAME)
backup_datasets=(${=backup_datasets}) # Конвертирование в массив

declare -A work_datasets # Словарь ключ = резервируемый локальный датасет, значение = имя датасета на удалённом сервере

# Если заданы local_dataset & remote_dataset
if [[ -v local_dataset && -v remote_dataset ]]; then
	work_datasets[$local_dataset]=$remote_dataset
else	# Если нет, то вычислить
	if [[ $incremental = true ]]; then
		# Высчитать что на что резервировать
		for ((i = 1; i <= ${#local_datasets}; i++)); do
			for ((j = 1; j <= ${#backup_datasets}; j++)); do
				a=$(sed -r "s/(.*)\/(.*)/\2/" <<<"$local_datasets[i]") # Определить нижний по иерархии
				b=$(sed -r "s/(.*)\/(.*)/\2/" <<<"$backup_datasets[j]") # тоже для второго
				if [[ "$a" == "$b" ]]; then
					work_datasets[${local_datasets[i]}]=${backup_datasets[j]}
				fi
			done
		done
	else # Если не инкрементная копия, то запросить меcто для резервирования
		# Вывести на экран запрос
		column_widths="| %-4s%-35s | %-4s%-35s\n"
		i_max=$(( ${#local_datasets} > ${#backup_datasets} ? ${#local_datasets} : ${#backup_datasets} )) # Максимальное количество элементов
		printf $column_widths "№"    "Локальные датасеты" "№"    "На $backup_server"
		printf $column_widths "----" "------------------" "----" "-------------------"
		# echo "№ - Локальные\t\t№ - Удалённые"
		for ((i = 1; i <= $i_max; i++)); do
		if [[ -v local_datasets[i] ]]; then i1=$i; else; i1=""; fi
		if [[ -v backup_datasets[i] ]]; then i2=$i; else; i2=""; fi
		printf $column_widths "$i1" "${local_datasets[i]}" "$i2" "${backup_datasets[i]}"
		done

		wrong_enter=true
		while $wrong_enter = true ; do
			read "loc_ds_indx_list?Локальные индексы через пробел (0=выход; all=режим исключения): "
			if [[ $loc_ds_indx_list = "all" ]]; then
				loc_ds_indx_list=({1..${#local_datasets}})
				break
			fi
			if [[ $loc_ds_indx_list = *[[:digit:]]* ]]; then
				loc_ds_indx_list=(${=loc_ds_indx_list}) # Convert to list
				for i in $loc_ds_indx_list; do
					if [[ $i = 0 ]]; then exit; fi
					if (( $i < 1 )) || (( $i > ${#local_datasets} )); then
						echo "$i ! Выход за диапазон"
						wrong_enter=true
						break
					else
						wrong_enter=false
					fi
				done
			fi
		done

		echo "Введите номер датасета в котором создать кописю исходного (- исключить, 0 выйти)"
		for i in $loc_ds_indx_list; do
			read "bak_ds_indx?$i - ${local_datasets[i]} -> "
			case $bak_ds_indx in
				0) exit;;
				-) ;;
				*) if (( $bak_ds_indx >= 1 )) && (( $bak_ds_indx <= ${#backup_datasets} )); then
					loc_name_ds=$(sed -r "s/(.*)\/(.*)/\2/" <<<"$local_datasets[$i]") # Определить нижний по иерархии
					bak_ds="$bak_ds/$loc_name_ds" # Добавить имя создаваемого датасета
					work_datasets[$local_datasets[$i]]=$bak_ds
				else; exit 1; fi ;;
			esac
		done
	fi # Конец блока запроса места резервирования
fi # Конец блока наличия переданных параметров датасетов в скрипт


# Словарь со списком последких снапшотов для датасетов
declare -A last_loc_snaps
echo -ne "Проверяю на наличие $inc_snapshot...\033[2K\r"
column_widths="%-45s * %-45s\n"
for loc_ds bak_ds in ${(kv)work_datasets}; do
	# Запросить последние снимки
	loc_snap="$(zfs list -t snapshot -o name ${loc_ds} 2>/dev/null | grep 'manual' | tail -n1 | egrep -o '@.+' )"
	loc_snap=${loc_snap:1} # удалить @ в начале строки
	if [[ $incremental = true ]]; then
		bak_snap=$(ssh $backup_user@$backup_server zfs list -t snapshot -o name ${bak_ds} 2>/dev/null | grep 'manual' | tail -n1 | egrep -o '@.+')
		bak_snap=${bak_snap:1} # удалить @ в начале строки
	fi
	# echo "$loc_ds @ $loc_snap -> $bak_ds @ $bak_snap"
	# Если последний снимок локального датасета совпадает с сегодняшним
	if [[ "$loc_snap" = "$inc_snapshot" ]]; then
		if [[ $incremental = true ]] && [[ "$loc_snap" = "$bak_snap" ]]; then
			# Если последние снимки на обеих машинах совпадают в инкрементном режиме
			printf $column_widths "${loc_ds}@${loc_snap}" "существует на обеих машинах. Исключено из списка резервируемых"
			unset "work_datasets[$loc_ds]"
		else
			# Если на удалённой машине нет такого снимка, то предложить удалить.
			printf $column_widths "${loc_ds}@${loc_snap}" "существует на исходной машине, но отсутвует на $backup_server"
			read "act?ok: резервировать | re: пересоздать (удаляет сразу!!!) | по умолчанию исключить из работы : "
			case $act in
				re)	zfs destroy "${loc_ds}@${loc_snap}";
					loc_snap="$(zfs list -t snapshot -o name ${loc_ds} | grep 'manual' | tail -n1 |  egrep -o '@.+' )";
					loc_snap="${loc_snap:1}";
					last_loc_snaps[$loc_ds]=$loc_snap ;;
				ok)	last_loc_snaps[$loc_ds]=$loc_snap ;;
				*)	echo "$loc_ds Исключено из списка";
					unset "work_datasets[$loc_ds]" ;;
			esac
		fi
	else
		# добавить в список снапшотов
		if [[ $loc_snap != "" ]]; then
			last_loc_snaps[$loc_ds]=$loc_snap
		else
			if [[ $incremental = true ]]; then
				printf $column_widths "${loc_ds}" "Не имеет снимков. Не возможно сделать инкрементный резерв."
				unset "work_datasets[$loc_ds]"
			fi
		fi
	fi
done

# --------- Запрашиваем подтверждение пользователя ----------

echo ""
echo "Список резервируемых датасетов:"
column_widths="%-4s %-35s %-20s -> %-35s\n"

if [[ $incremental = true ]]; then
	i=0; for loc_ds bak_ds in ${(kv)work_datasets}; do
		(( i++ ))
		printf $column_widths $i $loc_ds "${last_loc_snaps[$loc_ds]}" $bak_ds
	done
	printf "%-4s %-35s %-20s\n" "" "Инкрементный снимок:" "$inc_snapshot (*)"
else
	i=0; for loc_ds bak_ds in ${(kv)work_datasets}; do
		(( i++ ))
		printf $column_widths $i $loc_ds "$inc_snapshot (*)" "$bak_ds (*)"
	done
fi

# --------- Выбор датасетов ----------

if [[ $check = true ]]; then
	read 'work?Утвердите данные. "y" в работу | "1 2.." выбрать через пробел | НЕТ по умолчанию : '
	if [[ $work = *[[:digit:]]* ]]; then
		indx_list=(${=work})
		# Если нет индекса в списке, то удалить из словаря для резервирования
		i=0; for loc_ds bak_ds in ${(kv)work_datasets}; do
			i=$((i+1))
			if [[ ${indx_list[(r)$i]} != $i ]]; then # Если нет $i в списке $indx_list
				unset "work_datasets[$loc_ds]"
			else
				echo "$i - $loc_ds \t ${last_loc_snaps[$loc_ds]} \t -> \t${bak_ds}"
			fi
		done
	fi

	# Если был выведен список, то запросить запуск в работу нового списка
	if [[ ${#indx_list} > 0 ]]; then
		work="N"
		read "work?В работу (y/N) "
	fi
else
	work="y"
fi

# --------- Резервирование ----------

if [[ $work = "y" || $work = "Y" ]]; then
	TS0=$(date +%s)
	echo $TS0 > /dev/shm/backup_time_start  # Сохраняем отметку о времени начала
	echo $TS0 > /dev/shm/backup_time_circle # Сохраняем отметку о времени круга
	echo "Результат работы записываю в файл $LOGFILE"

	# Записать лог список резервируемых датасетов
	echo "--- $(date +'%Y.%m.%d %H:%M.%S') ---" >> "$LOGFILE"
	echo "Список резервируемых датасетов:" >> "$LOGFILE"
	n_tasks=0
	for loc_ds in ${(k)work_datasets}; do
		(( n_tasks++ ))
		echo "$key \t ${last_loc_snaps[$loc_ds]} \t -> \t${work_datasets[$loc_ds]}" >> "$LOGFILE"
	done
	echo $n_tasks > /dev/shm/backup_n_tasks

	echo_log
	echo_log "--- Резервирую данные ---"
	i_task=0 # Выполненных зачач
	if [[ $progress == true ]]; then echo $i_task > /dev/shm/backup_i_task; fi # Сохраняем количество готовых задач в память
	for loc_ds bak_ds in ${(kv)work_datasets}; do
		TS1=$(date +%s)
		if [[ $progress == true ]]; then echo $TS1 > /dev/shm/backup_time_circle; fi # Сохраняем отметку о времени начала следующего задания
		echo_log " * snapshot ${loc_ds}@${inc_snapshot}"
		zfs snapshot ${loc_ds}@${inc_snapshot}; #FIXME выдаёт ошибку если мы оставили существующий снимок
		echo_log " * Start sending ${loc_ds} at $(date +'%Y.%m.%d %H:%M.%S')"
		case $incremental in
			true) zfs send -V -i ${loc_ds}@${last_loc_snaps[$loc_ds]} ${loc_ds}@${inc_snapshot} | ssh $backup_user@$backup_server zfs receive ${bak_ds}@${inc_snapshot} ;;
			false) zfs send -V ${loc_ds}@${inc_snapshot} | ssh $backup_user@$backup_server zfs receive ${bak_ds} ;;
		esac
		TS2=$(date +%s)
		(( i_task++ ))
		if [[ $progress == true ]]; then echo $i_task > /dev/shm/backup_i_task; fi
		echo_log "--- готово за: $(date -d@$(($TS2-$TS1)) -u '+%H:%M.%S') ---"
	done & # Запустить резервирование в фоне
fi



# !
# Конец блока условия progress_only
fi




# Вывод прогресса
# Это отдельный блок программы, который не имеет прямого доступа к переменным процесса в фоне.
# Обмен переменными проивится ерез временный файл
if [[ $progress == true ]]; then
	echo ""

	pids=$(ps -t $(tty | sed 's|/dev/||') -o pid,args | grep -e zfs_send.zsh -e "zfs send" | grep -v -- 'grep' | grep -v -- 'progress-only' ) # Список PIDs с командой запуска

	if [[ $pids = "" ]]; then
    	echo "Нет запущенного процесса резервирования в этом терминале."
	    read "srch?Икать во всех процессах системы? (y/N) "
	    if [[ $srch = y ]]; then
	        pids=$(ps a -o pid,args | grep -e zfs_send.zsh -e "zfs send" | grep -v -- 'grep' | grep -v -- 'progress-only' )
	    fi
	fi
echo  $pids
	# Выводить прогресс запрашивая список запущенных процессов
	if [[ $pids = "" ]]; then
	    echo "Нет запущенного процесса резервирования"
	else
		echo ""
		# инициализируем переменнные
		TS0=$(< /dev/shm/backup_time_start)
		TS1=$(< /dev/shm/backup_time_circle)

		last_percent=0
		last_volume=0
		last_volume_time=$TS1
		speed=0
		reset_interval=10 # интревал сбрасывания скорости до нуля, в случае если нет именений
		last_update_time=$TS1  # Время последнего обновления

		time_changed_percent_value=$TS1

		n_tasks=$(< /dev/shm/backup_n_tasks)
		while [[ $i_task -lt $n_tasks ]]; do # выводим статус пока не завершены все задачи
			i_task=$(< /dev/shm/backup_i_task)
			progress=$(ps -u | grep "sending" | grep -v "grep" | sed -r "s/(.*) zfs: (.*)/\2/")

            # Извлекаем объем и процент
			progress_part=$(echo "$progress" | grep -oP '\(\K.*(?=\))') # Извлекаем только часть в скобках

			if [[ -n $progress_part ]]; then
				# Разбиваем часть на компоненты
				# Процент определим по объёму
				# percent=$(echo "$progress_part" | awk -F': |/' '{print $1}' | tr -d '%')
				current_volume_str=$(echo "$progress_part" | awk -F': |/' '{print $2}')
				total_volume_str=$(echo "$progress_part" | awk -F': |/' '{print $3}')

                # Определяем единицу для интервала сброса
                unit=${current_volume_str: -1}
                case $unit in
                    G) reset_interval=100;;
                    M) reset_interval=10;;
                    *) reset_interval=5;;
                esac

                # Преобразуем в байты
                current=$(echo "$current_volume_str" | awk '{sub(/[^0-9.]/,""); print $1}')
				case $unit in
					G) current_volume=$(( current * 1073741824 )) ;;
					M) current_volume=$(( current * 1048576 )) ;;
					K) current_volume=$(( current * 1024 )) ;;
					*) current_volume=0 ;;
				esac

				total=$(echo "$total_volume_str" | awk '{sub(/[^0-9.]/,""); print $1}')
				case $unit in
					G) total_volume=$(( total * 1073741824 )) ;;
					M) total_volume=$(( total * 1048576 )) ;;
					K) total_volume=$(( total * 1024 )) ;;
					*) total_volume=0 ;;
				esac

				percent=$(( current_volume * 100 / total_volume ))
            else
                percent=0
                current_volume=0
            fi

			# Вычисляем скорость
			now=$(date +%s)

			# Если объем изменился
			elapsed_from_vol_change=$((now - last_volume_time))
			if [[ $current_volume -ne $last_volume && elapsed_from_vol_change -ne 0 ]]; then
				speed=$(awk "BEGIN { printf \"%.1f\", ($current_volume - $last_volume) / $elapsed_from_vol_change }")
				last_volume=$current_volume
				last_volume_time=$now
				last_update_time=$now  # Обновляем время последнего изменения
			else
				# Если объем не изменился, проверяем время с последнего обновления
				time_since_update=$((now - last_update_time))
				if (( time_since_update >= reset_interval )); then
					speed=0  # Обнуляем скорость, если нет изменений в течение интервала
					last_update_time=$now  # Обновляем время
				fi
			fi

            # Конвертируем скорость в человекочитаемый формат
			if (( speed > 1073741824 )); then
				speed_human=$(awk "BEGIN { printf \"%.1fG\", $speed / 1073741824 }")
			elif (( speed > 1048576 )); then
				speed_human=$(awk "BEGIN { printf \"%.1fM\", $speed / 1048576 }")
			elif (( speed > 1024 )); then
				speed_human=$(awk "BEGIN { printf \"%.1fK\", $speed / 1024 }")
			else
				speed_human="${speed}B"
			fi


			# Начинаем вывод
			echo -n "\033[2K\r$progress  $speed_human/s "

			if (( percent > 0 )); then
				now=$(date +%s)
				if (( percent != last_percent )); then
					time_changed_percent_value=$now
					last_percent=$percent
					TS1=$(< /dev/shm/backup_time_circle)
					elapsed_total=$(( $(date +%s) - TS1 ))
				fi
				time_part_of_percent=$(( now - time_changed_percent_value )) # Время с последнего изменения процента
				# Используем awk для более точных расчётов
				# estimated_total  - Всего времени на задачу
				# estimated_remain - Осталось времени на задачу
				read estimated_total estimated_remain <<< $(awk -v p="$percent" -v t="$elapsed_total" -v tp="$time_part_of_percent" '
				BEGIN {
					total = (p > 0 ? (100 / p * t) : 0);
					remain = (total > 0 ? (total - t - tp) : 0);
					printf("%d %d", total, remain)
				}')
				if (( estimated_remain >= 86400 )) || (( estimated_total >= 86400 )); then # Если осталось более чем сутки, то отобразить дни
					estimated_remain_days=$(( estimated_remain / 86400 )); estimated_remain_time=$(( estimated_remain % 86400 ))
					estimated_total_days=$(( estimated_total / 86400 )); estimated_total_time=$(( estimated_total % 86400 ))
					echo -n "  $estimated_remain_days д. $(date -d@$estimated_remain_time -u '+%H:%M.%S') ост. / $estimated_total_days д. $(date -d@$estimated_total_time -u '+%H:%M.%S') "
				else
					echo -n "  $(date -d@$estimated_remain -u '+%H:%M.%S') ост. / $(date -d@$estimated_total -u '+%H:%M.%S') "
				fi
			else
				for ((i = 1; i < 4 ; i++)); do echo -n "."; sleep 1; done
			fi

			sleep 1

		done

		if [[ -f /dev/shm/backup_time_start ]]; then
			TS0=$(< /dev/shm/backup_time_start)
			echo_log "\n--- Все завершено за $(date -d@$(( $(date +%s) - $TS0 )) -u '+%H:%M.%S') ---\n"
		else
			echo_log "\n--- Все задания завершены ---\n"
		fi

	    # Удалить временные файлы состояний, если они существуют
		[ -f /dev/shm/backup_i_task ] && rm /dev/shm/backup_i_task
		[ -f /dev/shm//dev/shm/backup_time_start ] && rm /dev/shm/backup_time_start
		[ -f /dev/shm//dev/shm/backup_time_circle ] && rm /dev/shm/backup_time_circle
		# Обнуляем переменные после завершения задачи, чтобы избежать ошибок при повторном запуске
		unset last_volume last_volume_time last_update_time last_percent time_changed_percent_value

	fi

fi