#!/usr/bin/env bash
set -euo pipefail

usage() {
  cat <<'USAGE'
Copy files via scp/rsync/sftp with automatic optimization.

Usage:
  ssh_copy.sh [options] push HOST LOCAL_PATH REMOTE_PATH
  ssh_copy.sh [options] pull HOST REMOTE_PATH LOCAL_PATH

Options:
  -u, --user USER           Override SSH user (or set REMOTE_USER)
  -p, --port PORT           SSH port (default: REMOTE_PORT or 22)
  -i, --key PATH            Identity file (default: REMOTE_KEY)
      --connect-timeout SEC Connect timeout (default: REMOTE_CONNECT_TIMEOUT or 10)
  -r, --recursive           Copy directories recursively
      --accept-new          Set StrictHostKeyChecking=accept-new

  # Transfer method
  -m, --method {auto,scp,rsync,sftp}
                            Transfer method (default: auto)
      --tar                 Force tar+scp packaging (exclusive with --method)
      --tar-format {tar,tar.gz,tar.xz}
                            Tar format (default: tar.gz)
      --tar-threshold N     File count threshold to trigger tar (default: 20)

  # Compression
      --compress {auto,yes,no}
                            Enable compression (default: auto)
      --compress-level N    Compression level 1-9 (default: 6)

  # Rsync options
      --exclude PATTERN     Exclude pattern (can be repeated)
      --delete              Delete extraneous files in destination
      --whole-file          Force full file transfer (no delta)

  # Output
      --progress            Show transfer progress
      --stats               Show transfer statistics
      --dry-run             Print the command that would run
  -h, --help                Show help

Environment defaults:
  REMOTE_USER, REMOTE_PORT, REMOTE_KEY, REMOTE_CONNECT_TIMEOUT

Method selection guide:
  - Single file: scp (default)
  - Directory sync: rsync -a (default)
  - Many small files (>20): tar.gz + scp (default)
  - Large file (>100MB): rsync -z (default)
  - Exclude patterns: rsync --exclude (only option)

Examples:
  # Standard push/pull (auto-selects best method)
  ssh_copy.sh push my-server ./file.txt /tmp/
  ssh_copy.sh pull my-server /var/log/syslog ./syslog

  # Force specific method
  ssh_copy.sh --method rsync -r push my-server ./dir /tmp/
  ssh_copy.sh --method sftp push my-server ./file.txt /tmp/

  # Force tar packaging
  ssh_copy.sh --tar push my-server ./many-files/ /tmp/

  # Rsync with exclusions
  ssh_copy.sh --method rsync -r --exclude '.git' --exclude 'node_modules' \
    --delete push my-server ./project/ /tmp/project/

  # Dry run
  ssh_copy.sh --method tar --dry-run push my-server ./dir /tmp/
USAGE
}

fail() {
  echo "Error: $*" >&2
  exit 2
}

require_arg() {
  local value="${1:-}"
  local opt="${2:-option}"
  [[ -n "$value" ]] || fail "$opt requires a value"
  printf '%s' "$value"
}

quote_shell() {
  local value="$1"
  printf "'%s'" "${value//\'/\'\\\'\'}"
}

remote_path_expr() {
  local value="$1"
  if [[ "$value" == "~" ]]; then
    printf '%s' '$HOME'
  elif [[ "$value" == "~/"* ]]; then
    printf '%s' "\$HOME/$(quote_shell "${value#\~/}")"
  else
    quote_shell "$value"
  fi
}

# Default values
port="${REMOTE_PORT:-22}"
user="${REMOTE_USER:-}"
key="${REMOTE_KEY:-}"
connect_timeout="${REMOTE_CONNECT_TIMEOUT:-10}"

# Method selection
method="auto"
recursive=false
accept_new=false
dry_run=false

# Tar options
force_tar=false
tar_format="tar.gz"
tar_threshold=20

# Compression
compress="auto"
compress_level=6

# Rsync options
declare -a exclude_patterns=()
delete_mode=false
whole_file=false

# Output options
show_progress=false
show_stats=false

# Transfer source stats
source_kind=""
source_count=0
source_size=0

# Parse options
while [[ $# -gt 0 ]]; do
  case "$1" in
    -u|--user)
      user="$(require_arg "${2:-}" "$1")"
      shift 2
      ;;
    -p|--port)
      port="$(require_arg "${2:-}" "$1")"
      shift 2
      ;;
    -i|--key)
      key="$(require_arg "${2:-}" "$1")"
      shift 2
      ;;
    --connect-timeout)
      connect_timeout="$(require_arg "${2:-}" "$1")"
      shift 2
      ;;
    -r|--recursive)
      recursive=true
      shift
      ;;
    --accept-new)
      accept_new=true
      shift
      ;;
    -m|--method)
      method="$(require_arg "${2:-}" "$1")"
      shift 2
      ;;
    --tar)
      force_tar=true
      shift
      ;;
    --tar-format)
      tar_format="$(require_arg "${2:-}" "$1")"
      shift 2
      ;;
    --tar-threshold)
      tar_threshold="$(require_arg "${2:-}" "$1")"
      shift 2
      ;;
    --compress)
      compress="$(require_arg "${2:-}" "$1")"
      shift 2
      ;;
    --compress-level)
      compress_level="$(require_arg "${2:-}" "$1")"
      shift 2
      ;;
    --exclude)
      exclude_patterns+=("$(require_arg "${2:-}" "$1")")
      shift 2
      ;;
    --delete)
      delete_mode=true
      shift
      ;;
    --whole-file)
      whole_file=true
      shift
      ;;
    --progress)
      show_progress=true
      shift
      ;;
    --stats)
      show_stats=true
      shift
      ;;
    --dry-run)
      dry_run=true
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    -*)
      echo "Unknown option: $1" >&2
      usage >&2
      exit 2
      ;;
    *)
      break
      ;;
  esac
done

if [[ $# -lt 4 ]]; then
  usage >&2
  exit 2
fi

direction="$1"
shift

case "$direction" in
  push|pull)
    ;;
  *)
    fail "Invalid direction: $direction (expected push or pull)"
    ;;
esac

host="$1"
shift

if [[ "$method" != "auto" && "$method" != "scp" && "$method" != "rsync" && "$method" != "sftp" && "$method" != "tar" ]]; then
  fail "Invalid method: $method"
fi

case "$compress" in
  auto|yes|no)
    ;;
  *)
    fail "Invalid compress mode: $compress"
    ;;
esac

case "$tar_format" in
  tar|tar.gz|tgz|tar.xz)
    ;;
  *)
    fail "Invalid tar format: $tar_format"
    ;;
esac

if ! [[ "$tar_threshold" =~ ^[0-9]+$ ]]; then
  fail "Invalid tar threshold: $tar_threshold"
fi

if ! [[ "$compress_level" =~ ^[0-9]+$ ]] || [[ "$compress_level" -lt 1 || "$compress_level" -gt 9 ]]; then
  fail "Invalid compress level: $compress_level"
fi

if ! [[ "$connect_timeout" =~ ^[0-9]+$ ]]; then
  fail "Invalid connect timeout: $connect_timeout"
fi

if $force_tar && [[ "$method" != "auto" && "$method" != "tar" ]]; then
  fail "--tar cannot be combined with --method $method"
fi

if $force_tar; then
  method="tar"
fi

# Build destination host string
dest_host="$host"
if [[ -n "$user" ]]; then
  host_no_user="${host#*@}"
  dest_host="${user}@${host_no_user}"
fi

# Build common SSH options
ssh_opts=(
  -p "$port"
  -o "ConnectTimeout=${connect_timeout}"
  -o "ServerAliveInterval=30"
  -o "ServerAliveCountMax=3"
)

if [[ -n "$key" ]]; then
  ssh_opts+=(-i "$key" -o "IdentitiesOnly=yes")
fi

if $accept_new; then
  ssh_opts+=(-o "StrictHostKeyChecking=accept-new")
fi

# Build scp options
scp_opts=(-P "$port" -p)
if [[ -n "$key" ]]; then
  scp_opts+=(-i "$key" -o "IdentitiesOnly=yes")
fi
if $recursive; then
  scp_opts+=(-r)
fi
if $accept_new; then
  scp_opts+=(-o "StrictHostKeyChecking=accept-new")
fi
if $show_progress; then
  scp_opts+=(-v)
fi

# Get file count for a path
get_file_count() {
  local path="$1"
  if [[ -f "$path" ]]; then
    echo 1
  elif [[ -d "$path" ]]; then
    if $recursive; then
      find "$path" -type f 2>/dev/null | wc -l
    else
      find "$path" -maxdepth 1 -type f 2>/dev/null | wc -l
    fi
  else
    echo 0
  fi
}

# Get total size in bytes
get_total_size() {
  local path="$1"
  if [[ -f "$path" ]]; then
    stat -c%s "$path" 2>/dev/null || echo 0
  elif [[ -d "$path" ]]; then
    if $recursive; then
      du -sb "$path" 2>/dev/null | cut -f1 || echo 0
    else
      find "$path" -maxdepth 1 -type f -exec stat -c%s {} + 2>/dev/null | awk '{s+=$1} END {print s+0}'
    fi
  else
    echo 0
  fi
}

probe_local_source_stats() {
  local path="$1"

  if [[ -f "$path" ]]; then
    source_kind="file"
    source_count=1
    source_size=$(stat -c%s "$path" 2>/dev/null || echo 0)
    return
  fi

  if [[ -d "$path" ]]; then
    source_kind="dir"
    source_count=$(find "$path" -maxdepth 1 -type f 2>/dev/null | wc -l | tr -d '[:space:]')
    source_size=$(du -sb "$path" 2>/dev/null | cut -f1 || echo 0)
    return
  fi

  fail "Local path not found: $path"
}

probe_remote_source_stats() {
  local path="$1"
  local quoted_path
  quoted_path=$(remote_path_expr "$path")
  local remote_cmd
  remote_cmd="if [ -f $quoted_path ]; then printf 'file 1 %s\n' \"\$(stat -c%s -- $quoted_path 2>/dev/null || echo 0)\"; elif [ -d $quoted_path ]; then count=\$(find $quoted_path -maxdepth 1 -type f 2>/dev/null | wc -l | tr -d '[:space:]'); size=\$(du -sb $quoted_path 2>/dev/null | cut -f1 || echo 0); printf 'dir %s %s\n' \"\$count\" \"\$size\"; else printf 'missing 0 0\n'; fi"

  read -r source_kind source_count source_size < <(
    ssh "${ssh_opts[@]}" "$dest_host" "$remote_cmd"
  )

  if [[ "$source_kind" == "missing" ]]; then
    fail "Remote path not found: $path"
  fi
}

# Determine if tar should be used
should_use_tar() {
  if $force_tar; then
    return 0
  fi

  if [[ "$source_kind" != "dir" ]]; then
    return 1
  fi

  if [[ "$source_count" -gt "$tar_threshold" ]]; then
    return 0
  fi

  return 1
}

# Auto-detect best transfer method
auto_detect_method() {
  if should_use_tar; then
    echo "tar"
    return
  fi

  # If rsync options specified, use rsync
  if [[ ${#exclude_patterns[@]} -gt 0 ]] || $delete_mode || $whole_file; then
    echo "rsync"
    return
  fi

  if [[ "$source_kind" == "unknown" ]]; then
    if $recursive || [[ "$source_path" == */ ]]; then
      echo "rsync"
    else
      echo "scp"
    fi
    return
  fi

  if [[ "$source_kind" == "file" ]]; then
    if [[ "$source_size" -gt 104857600 ]]; then
      echo "rsync"
    else
      echo "scp"
    fi
    return
  fi

  # Directory sync defaults to rsync
  if [[ "$source_kind" == "dir" ]]; then
    echo "rsync"
    return
  fi

  fail "Unable to determine source type for auto transfer"
}

# Determine compression flag
get_compress_flag() {
  case "$compress" in
    yes)
      echo true
      ;;
    no)
      echo false
      ;;
    auto)
      # Auto compress for large files (>100MB)
      if [[ "$source_size" -gt 104857600 ]]; then
        echo true
      else
        echo false
      fi
      ;;
  esac
}

# Get tar extension based on format
get_tar_extension() {
  case "$tar_format" in
    tar.gz|tgz)
      echo "tar.gz"
      ;;
    tar.xz)
      echo "tar.xz"
      ;;
    tar)
      echo "tar"
      ;;
    *)
      echo "tar.gz"
      ;;
  esac
}

# Get tar compression flags
get_tar_compress_flags() {
  case "$tar_format" in
    tar.gz|tgz)
      echo "-z"
      ;;
    tar.xz)
      echo "-J"
      ;;
    tar)
      echo ""
      ;;
    *)
      echo "-z"
      ;;
  esac
}

# Build rsync options
build_rsync_opts() {
  local rsync_opts=("-a" "-v")

  if $show_progress; then
    rsync_opts+=("--progress")
  fi

  if $show_stats; then
    rsync_opts+=("--stats")
  fi

  if [[ ${#exclude_patterns[@]} -gt 0 ]]; then
    for pattern in "${exclude_patterns[@]}"; do
      rsync_opts+=("--exclude=$pattern")
    done
  fi

  if $delete_mode; then
    rsync_opts+=("--delete")
  fi

  if $whole_file; then
    rsync_opts+=("--whole-file")
  fi

  # Compression for large files
  local use_compress
  use_compress=$(get_compress_flag)
  if $use_compress; then
    rsync_opts+=("-z" "--compress-level=$compress_level")
  fi

  echo "${rsync_opts[@]}"
}

# Build SSH command for rsync
build_rsync_ssh_cmd() {
  local ssh_cmd=("ssh")
  ssh_cmd+=(-p "$port")
  ssh_cmd+=(-o "ConnectTimeout=${connect_timeout}")
  ssh_cmd+=(-o "ServerAliveInterval=30")
  ssh_cmd+=(-o "ServerAliveCountMax=3")
  if [[ -n "$key" ]]; then
    ssh_cmd+=(-i "$key" -o "IdentitiesOnly=yes")
  fi
  if $accept_new; then
    ssh_cmd+=(-o "StrictHostKeyChecking=accept-new")
  fi
  echo "${ssh_cmd[@]}"
}

# Push with tar+scp
do_push_with_tar() {
  local local_path="$1"
  local remote_path="$2"
  local ext
  ext=$(get_tar_extension)
  local local_tarball
  local_tarball="$(mktemp "${TMPDIR:-/tmp}/_transfer_XXXXXX.${ext}")"
  local remote_tarball="/tmp/$(basename "$local_tarball")"
  local tar_compress
  tar_compress=$(get_tar_compress_flags)

  if $dry_run; then
    echo "tar ${tar_compress} -cf '$local_tarball' -C '$(dirname "$local_path")' '$(basename "$local_path")'"
    echo "scp ${scp_opts[*]} '$local_tarball' '${dest_host}:$remote_tarball'"
    echo "ssh ${ssh_opts[*]} '$dest_host' \"mkdir -p $(remote_path_expr "$remote_path") && tar ${tar_compress} -xf $(quote_shell "$remote_tarball") -C $(remote_path_expr "$remote_path") && rm -f $(quote_shell "$remote_tarball")\""
    echo "rm -f '$local_tarball'"
    return
  fi

  # Create tarball
  tar ${tar_compress} -cf "$local_tarball" -C "$(dirname "$local_path")" "$(basename "$local_path")"

  # Transfer via scp
  scp "${scp_opts[@]}" "$local_tarball" "${dest_host}:${remote_tarball}"

  # Remote: extract and cleanup
  ssh "${ssh_opts[@]}" "$dest_host" "mkdir -p $(remote_path_expr "$remote_path") && tar ${tar_compress} -xf $(quote_shell "$remote_tarball") -C $(remote_path_expr "$remote_path") && rm -f $(quote_shell "$remote_tarball")"

  # Local cleanup
  rm -f "$local_tarball"
}

# Pull with tar+scp
do_pull_with_tar() {
  local remote_path="$1"
  local local_path="$2"
  local ext
  ext=$(get_tar_extension)
  local local_tarball
  local_tarball="$(mktemp "${TMPDIR:-/tmp}/_transfer_XXXXXX.${ext}")"
  local remote_tarball="/tmp/$(basename "$local_tarball")"
  local tar_compress
  tar_compress=$(get_tar_compress_flags)
  local remote_dir
  remote_dir=$(dirname "$remote_path")
  local remote_base
  remote_base=$(basename "$remote_path")

  if $dry_run; then
    echo "ssh ${ssh_opts[*]} '$dest_host' \"cd $(remote_path_expr "$remote_dir") && tar ${tar_compress} -cf $(quote_shell "$remote_tarball") $(quote_shell "$remote_base")\""
    echo "scp ${scp_opts[*]} '${dest_host}:$remote_tarball' '$local_tarball'"
    echo "mkdir -p '$local_path' && tar ${tar_compress} -xf '$local_tarball' -C '$local_path'"
    echo "rm -f '$local_tarball'"
    echo "ssh ${ssh_opts[*]} '$dest_host' \"rm -f $(quote_shell "$remote_tarball")\""
    return
  fi

  # Remote: create tarball
  ssh "${ssh_opts[@]}" "$dest_host" "cd $(remote_path_expr "$remote_dir") && tar ${tar_compress} -cf $(quote_shell "$remote_tarball") $(quote_shell "$remote_base")"

  # Download tarball
  scp "${scp_opts[@]}" "${dest_host}:${remote_tarball}" "$local_tarball"

  # Local: extract and cleanup
  mkdir -p "$local_path"
  tar ${tar_compress} -xf "$local_tarball" -C "$local_path"
  rm -f "$local_tarball"

  # Remote cleanup
  ssh "${ssh_opts[@]}" "$dest_host" "rm -f $(quote_shell "$remote_tarball")"
}

# Standard scp transfer
do_scp() {
  local local_path="$1"
  local remote_path="$2"

  if $dry_run; then
    if [[ "$direction" == "push" ]]; then
      echo "scp ${scp_opts[*]} '$local_path' '${dest_host}:$remote_path'"
    else
      echo "scp ${scp_opts[*]} '${dest_host}:$local_path' '$remote_path'"
    fi
    return
  fi

  case "$direction" in
    push)
      scp "${scp_opts[@]}" "$local_path" "${dest_host}:${remote_path}"
      ;;
    pull)
      scp "${scp_opts[@]}" "${dest_host}:${local_path}" "$remote_path"
      ;;
  esac
}

# Rsync transfer
do_rsync() {
  local local_path="$1"
  local remote_path="$2"
  local rsync_opts
  rsync_opts=$(build_rsync_opts)
  local rsync_ssh
  rsync_ssh=$(build_rsync_ssh_cmd)

  if $dry_run; then
    if [[ "$direction" == "push" ]]; then
      echo "rsync ${rsync_opts} -e '${rsync_ssh}' '$local_path' '${dest_host}:$remote_path'"
    else
      echo "rsync ${rsync_opts} -e '${rsync_ssh}' '${dest_host}:$local_path' '$remote_path'"
    fi
    return
  fi

  case "$direction" in
    push)
      rsync ${rsync_opts} -e "${rsync_ssh}" "$local_path" "${dest_host}:${remote_path}"
      ;;
    pull)
      rsync ${rsync_opts} -e "${rsync_ssh}" "${dest_host}:${local_path}" "$remote_path"
      ;;
  esac
}

# SFTP transfer (interactive/ scripted)
do_sftp() {
  local local_path="$1"
  local remote_path="$2"

  if $dry_run; then
    if [[ "$direction" == "push" ]]; then
      echo "sftp ${ssh_opts[*]} '${dest_host}' <<< 'put $local_path $remote_path'"
    else
      echo "sftp ${ssh_opts[*]} '${dest_host}' <<< 'get $local_path $remote_path'"
    fi
    return
  fi

  case "$direction" in
    push)
      sftp "${ssh_opts[@]}" "$dest_host" <<< "put '$local_path' '$remote_path'"
      ;;
    pull)
      sftp "${ssh_opts[@]}" "$dest_host" <<< "get '$local_path' '$remote_path'"
      ;;
  esac
}

# Main logic
local_path="$1"
remote_path="$2"
source_path="$local_path"

need_source_stats=false
if [[ "$method" == "auto" ]]; then
  need_source_stats=true
elif [[ "$method" == "rsync" && "$compress" == "auto" ]]; then
  need_source_stats=true
fi

if $need_source_stats; then
  if $dry_run && [[ "$direction" == "pull" ]]; then
    source_kind="unknown"
    source_count=0
    source_size=0
  elif [[ "$direction" == "push" ]]; then
    probe_local_source_stats "$local_path"
  else
    probe_remote_source_stats "$local_path"
  fi
fi

if [[ "$method" == "auto" ]]; then
  method=$(auto_detect_method)
fi

# Validate method
case "$method" in
  scp|rsync|sftp|tar)
    ;;
  *)
    echo "Invalid method: $method" >&2
    exit 2
    ;;
esac

# Execute transfer
case "$method" in
  scp)
    do_scp "$local_path" "$remote_path"
    ;;
  rsync)
    do_rsync "$local_path" "$remote_path"
    ;;
  sftp)
    do_sftp "$local_path" "$remote_path"
    ;;
  tar)
    if [[ "$direction" == "push" ]]; then
      do_push_with_tar "$local_path" "$remote_path"
    else
      do_pull_with_tar "$local_path" "$remote_path"
    fi
    ;;
esac