#!/bin/bash
#v12
# Перевірка: якщо скрипт запущено автономно, а не через 'source'
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
  cat <<EOF >&2
  Error: This is a library script and should be sourced, not executed directly.
  Usage:
    Env vars:
      \$SSH_DATA_DIR - dir for ssh specific files - conf, known_hosts, etc
      \$RUNDIR - dir for mbuf stats on local and remote hosts
      \$ZFS_ANALYZER - path to zfs_analyzer script (optional)
      \$ZFS_SYNC_NO_COMMON_SNAP - "fail" (default) | "destroy"
      \$TRANSFERLOG  - recyclable file for process & error log
    Argument vars:
      \$src_host - source remote host:port or "local"
      \$src_ds   - source dataset
      \$dst_host - destination remote host:port or "local"
      \$dst_ds   - destination dataset
      one of hosts MUST be "local"
    Funcs:
      check_reachable - call before use main function
      zfs_sync_dataset - main function
    Returns:
      return 1 if failed
EOF
  exit 1
fi

#TODO: add check for no child datasets. There is no support for recursive sync

# Поведінка при відсутності спільного снапшота:
# "fail" (default) | "destroy"
ZFS_SYNC_NO_COMMON_SNAP="${ZFS_SYNC_NO_COMMON_SNAP:-fail}"


# Перевірка обов'язкової змінної оточення
if [[ -z "${SSH_DATA_DIR:-}" ]]; then
  echo "Error: SSH_DATA_DIR is not set." >&2
  exit 1
fi

if [[ -z "${RUNDIR:-}" ]]; then
  echo "Error: RUNDIR is not set." >&2
  exit 1
fi

# Шлях до аналізатора (за замовчуванням шукаємо в тій же директорії)
if [[ -z "$ZFS_ANALYZER" ]]; then
  ZFS_ANALYZER="/usr/bin/zynalyz"
fi

SSH_CONF="$SSH_DATA_DIR/config"
[[ ! -d "$SSH_DATA_DIR" ]] && mkdir -p "$SSH_DATA_DIR" && chmod 700 "$SSH_DATA_DIR"

# Bootstrapping конфігу, якщо він відсутній
if [[ ! -f "$SSH_CONF" ]]; then
  cat > "$SSH_CONF" <<EOF
Host *
    ControlMaster auto
    ControlPersist 1h
    ControlPath $SSH_DATA_DIR/%r@%h-%p
    StrictHostKeyChecking accept-new
    BatchMode yes
    ConnectTimeout 5
    UserKnownHostsFile $SSH_DATA_DIR/known_hosts
    HashKnownHosts yes
    ServerAliveInterval 60
    ServerAliveCountMax 3
EOF
  chmod 600 "$SSH_CONF"
fi

mbuf_stats=$RUNDIR/mbuf_stats

mbuf_logger() {
  while IFS= read -r line; do
    if [[ "$line" =~ ^mbuffer:\ in ]]; then
      echo "$line" >> "$mbuf_stats"
    else
      echo "$line" >&2
    fi
  done
}

check_reachable() {
  ssh_src exit >/dev/null 2>&1 || { echo "Error:  host $src_host unreachable" >&2; return 1; }
  ssh_dst exit >/dev/null 2>&1 || { echo "Error:  host $dst_host unreachable" >&2; return 1; }
  # Перевірка наявності mbuffer на обох хостах
  ssh_src "command -v mbuffer" >/dev/null 2>&1 || { echo "Error: mbuffer is not installed on source ($src_host)" >&2; return 1; }
  ssh_dst "command -v mbuffer" >/dev/null 2>&1 || { echo "Error: mbuffer is not installed on destination ($dst_host)" >&2; return 1; }
}

_ssh_exec() {
  local host=$1
  shift
  if [[ $host = local ]]; then
    bash -c "$*"
  else
    /usr/bin/ssh -F "$SSH_CONF" ${host/:/ -p} "$@"
  fi
}
ssh_src() { _ssh_exec "$src_host" "$@"; }
ssh_dst() { _ssh_exec "$dst_host" "$@"; }
ssh_src_stream() { _ssh_exec "$src_host" "$(declare -f mbuf_logger); $@ | mbuffer -W600 -q"; }
ssh_dst_stream() { _ssh_exec "$dst_host" "$(declare -f mbuf_logger); mbuffer -W600 -q | $@"; }

get_resume_token() {
  ssh_dst zfs get receive_resume_token -Hovalue $1 2>/dev/null | sed '/^-/d'
}

Rv=PRwhv
#Rv=PRwhsv
#       -L, --large-block
#       -P, --parsable
#       -R, --replicate
#       -c, --compressed
#       -e, --embed
#       -h, --holds
#       -s, --skip-missing
#       -v, --verbose
#       -p, --props
#       -w, --raw / -Lec

do_sync_resume() {
    ssh_src_stream zfs send -Pevt $resume_token \
    | ssh_dst_stream zfs recv -suvF $exclude_props $dst_ds
}


do_sync_full() {
  lastsnap=$(sed '$!d' <<< $src_snaps)
  [[ -z $lastsnap ]] && return
  echo do sync $src_ds@$lastsnap
  ssh_src_stream zfs send -$Rv $src_ds@$lastsnap \
    | ssh_dst_stream zfs recv -suvF $exclude_props $dst_ds
  local res=$?
  [[ $zfs_type = filesystem ]] && ssh_dst zfs set canmount=off $dst_ds
  return $res
}

do_sync_incremental() {
  lastsnap=$(sed '$!d' <<< $src_snaps)
  [[ $last_common_snap = $lastsnap ]] && return

  echo do sync $src_ds@$lastsnap
  ssh_src_stream zfs send -$Rv $src_ds@$lastsnap -I@$last_common_snap \
    | ssh_dst_stream zfs recv -uvF $exclude_props $dst_ds
}

_zfs_sync_dataset() {(
  trap - EXIT
  set +e -o pipefail
  echo "$src_ds->$dst_ds"

  zfs_type=$(ssh_src zfs get type -Hovalue $src_ds)
  exclude_props=
  [[ $zfs_type = filesystem ]] && exclude_props='-x canmount -x mountpoint -x ua.com.omix:autosnap'

  local old_resume_token=""
  while true; do
    resume_token=$(get_resume_token $dst_ds)
    [[ -z $resume_token ]] && break
    if [[ $old_resume_token = $resume_token ]]; then
      ssh_dst zfs recv -A $dst_ds
      resume_token=""
      break
    fi
    do_sync_resume
    old_resume_token=$resume_token
  done

  src_snaps=$(ssh_src zfs list -rtsnap -Ho name $src_ds | sed 's/.*@//;/omix_hourly/d')
  dst_snaps=$(ssh_dst zfs list -rtsnap -Ho name $dst_ds 2>/dev/null | sed 's/.*@//')

  if [[ -z $dst_snaps ]]; then 
    do_sync_full; return
  fi
  last_common_snap=$(grep -F -xf <(echo "$src_snaps") <(echo "$dst_snaps") | tail -n1)
  if [[ -z $last_common_snap ]]; then
    if [[ "$ZFS_SYNC_NO_COMMON_SNAP" == "destroy" ]]; then
      echo "Warning: no common snapshot, destroying $dst_ds and re-syncing" >&2
      ssh_dst zfs destroy -r $dst_ds
      do_sync_full; return
    else
      echo "Error: no common snapshot found between $src_ds and $dst_ds" >&2; return 1
    fi
  fi
  do_sync_incremental
)}

# Функція для показу статистики
show_transfer_stats() {
  echo "" >&2
  echo "=====================================" >&2
  if [[ -z "${TRANSFERLOG:-}" ]]; then
    echo "No transfer log available" >&2
  elif [[ ! -f "$TRANSFERLOG" ]]; then
    echo "Transfer log not found: $TRANSFERLOG" >&2
  elif [[ -f "$ZFS_ANALYZER" ]]; then
    bash "$ZFS_ANALYZER" "$TRANSFERLOG" >&2
  else
    echo "Transfer log: $TRANSFERLOG" >&2
    tail -n 10 "$TRANSFERLOG" >&2
  fi
  echo "=====================================" >&2
  echo "" >&2
}

zfs_sync_dataset() {
  ssh_src "mkdir -p '$RUNDIR'" && ssh_dst "mkdir -p '$RUNDIR'" || return 1

  _zfs_sync_dataset &>${TRANSFERLOG}
  local result=$?

  # Cleanup mbuf stats files
  ssh_src rm -f $mbuf_stats 2>/dev/null
  ssh_dst rm -f $mbuf_stats 2>/dev/null

  if [[ $result -ne 0 ]]; then
    # Exit code 2: dataset is a clone and destination lacks the origin snapshot
    if grep -q "cannot receive: local origin for clone" "$TRANSFERLOG" 2>/dev/null; then
      return 2
    fi
    return 1
  fi

  return 0
}

# Sync all VM disks using vm_disks[] and storage_map[] from parse_vmconf/parse_storage.
#
# Callback hooks (define in calling script, all optional):
#   disk_unsupported_cb  — called when storage is not in storage_map
#                           default: log warning and skip
#                           available: $dev, $stg_vol
#   disk_paths_cb        — must set src_ds and dst_ds
#                           default: src_ds=dst_ds=src_zpath/volname
#                           available: $dev, $storeid, $volname, $src_zpath
#   disk_pre_sync_cb     — called before zfs_sync_dataset
#   disk_post_sync_cb    — called after successful sync
#
sync_vm_disks() {
  for dev in $(echo "${!vm_disks[@]}" | tr ' ' '\n' | sort); do
    local stg_vol="${vm_disks[$dev]%%,*}"
    local storeid="${stg_vol%%:*}"
    local volname="${stg_vol#*:}"

    local src_zpath="${storage_map[$storeid]:-}"
    if [[ -z "$src_zpath" ]]; then
      if declare -F disk_unsupported_cb &>/dev/null; then
        disk_unsupported_cb
      else
        log "WARNING: $dev: storage '$stg_vol' is not supported, skipping"
      fi
      continue
    fi
    if [[ "$src_zpath" == "#shared" ]]; then
      log "Skip $dev ($storeid is shared storage)"; continue
    fi

    # Build dataset paths (can be overridden by callback)
    src_ds="${src_zpath}/${volname}"
    dst_ds="$src_ds"
    if declare -F disk_paths_cb &>/dev/null; then
      disk_paths_cb
    fi

    log "Syncing $dev: $src_ds → $dst_ds"

    if declare -F disk_pre_sync_cb &>/dev/null; then
      disk_pre_sync_cb
    fi

    local sync_rc=0; zfs_sync_dataset || sync_rc=$?
    if [[ "$SAVELOGS" == "true" ]]; then
      cp "$TRANSFERLOG" "${SAVELOGS_TMPDIR}/transfer_${dev}.log" 2>/dev/null || true
    fi
    if [[ $sync_rc -eq 2 ]]; then
      local origin=$(ssh_src "zfs get -Hovalue origin $src_ds")
      error "Failed: $dev is a clone of $origin — manually transfer the origin dataset first, then retry. Log: $TRANSFERLOG"
    elif [[ $sync_rc -ne 0 ]]; then
      error "Failed: $dev; Log: $TRANSFERLOG"
    fi

    if declare -F disk_post_sync_cb &>/dev/null; then
      disk_post_sync_cb
    fi

    local summary=$(bash "$ZFS_ANALYZER" "$TRANSFERLOG" --summary 2>/dev/null)
    log "✓ $dev synced: $summary"
  done
}

parse_storage(){
  unset storage_map storage_nodes
  while read -r line; do
    declare -gA $line
  done < <(awk -v allnodes="$1" '
/^[a-z]+:/ { 
    type = $1; sub(/:$/, "", type); 
    id = $2;
    pool = ""; nodes = ""; shared = 0;
    next;
}

/^[[:space:]]+pool[[:space:]]/   { pool = $2; }
/^[[:space:]]+nodes[[:space:]]/  { nodes = $2; }
/^[[:space:]]+shared[[:space:]]1/ { shared = 1; }

/^$/ {
        if (type == "zfspool") {
        }else if (type ~ /^(nfs|pbs|cifs|glusterfs|cephfs|rbd)$/ || shared == 1) {
            pool="#shared";
        }else{next;}

        printf "storage_map[%s]=%s\n", id, pool;
        printf "storage_nodes[%s]=%s\n", id, (nodes != "") ? nodes : allnodes;

}')
}

parse_vmconf() {
    unset vm_disks vm_name vm_onboot
    declare -gA vm_disks
    declare -g vm_name="unknown"
    declare -g vm_onboot=0

    while read -r line; do
        [[ -n "$line" ]] && declare -g "$line"
    done < <(awk '
        /^\[/ { exit }
        /^(name|hostname):/ { 
            val = $0; sub(/^(name|hostname):[[:space:]]*/, "", val);
            printf "vm_name=%s\n", val; 
            next 
        }

        # onboot (KVM та LXC однаковий)
        /^onboot:/ { 
            printf "vm_onboot=%s\n", ($2 == "1") ? 1 : 0; 
            next 
        }
        /^(unused[0-9]+|ide[0-9]+|sata[0-9]+|scsi[0-9]+|virtio[0-9]+|rootfs|mp[0-9]+):/ {
           if ($0 ~ /: none/) next;
           dev = $1; sub(/:$/, "", dev);

            # Витягуємо частину сторедж:волюм
            split($2, parts, ",");
            stg_vol = parts[1];

            # Шукаємо тільки потрібні прапорці
            opts = "";
            if ($0 ~ /backup=0/) opts = opts ",backup=0";
            if ($0 ~ /replicate=0/) opts = opts ",replicate=0";

            printf "vm_disks[%s]=%s%s\n", dev, stg_vol, opts;
        }
    ')
}

# USAGE:
#parse_storage "$(ls /etc/pve/nodes -m | sed 's/ //g')" < <(cat  /etc/pve/storage.cfg)
#declare -p storage_map storage_nodes

#config_file=/etc/pve/qemu-server/113.conf
#vmid=$(basename "$config_file" .conf)
#parse_vmconf < <(cat $config_file)

#declare -p vmid vm_name vm_onboot vm_disks
