#!/usr/bin/env zsh

# TODO: write proper docs for these configuration options.
#
# zstyle ':z4h:ssh:my_host' enable              'yes'
# zstyle ':z4h:ssh:*'       send-extra-files     '~/foo' '"$ZDOTDIR"/bar'
# zstyle ':z4h:ssh:*'       retrieve-extra-files '~/foo' '"$ZDOTDIR"/bar'
# zstyle ':z4h:ssh:*'       ssh-command          command ssh
# zstyle ':z4h:ssh:*'       retrieve-history     $ZDOTDIR/.zsh_history.remote
# zstyle ':z4h:ssh:*'       term                 tmux-256color
#
# z4h-ssh-configure() {
#   z4h_ssh_prelude+=(
#     "export BLAH=${(q)BLAH}"
#   )
#   z4h_ssh_send_files+=(
#     ~/foo   '~/foo'
#     $ZDOTDIR/bar '"$ZDOTDIR"/bar'
#   )
#   z4h_ssh_setup+=(
#     'echo "setting up"'
#   )
#   z4h_ssh_run=(
#     'echo "starting z4h"'
#     $z4h_ssh_launch_commands
#   )
#   z4h_ssh_teardown+=(
#     'echo "tearing down"'
#   )
#   z4h_ssh_retrieve_files+=(
#     '~/foo'          ~/foo
#     '"$ZDOTDIR"/bar' $ZDOTDIR/bar
#   )
# }

eval "$_z4h_opt"
-z4h-check-core-params || return

if (( _z4h_dangerous_root )); then
  print -Pru2 -- "%F{3}z4h%f: refusing to %Bssh%b as %F{1}root%f"
  return 1
fi

local -i must_passthrough i
local -a pos
for ((i = 1; i <= $#; ++i)); do
  case $*[i] in
    --) (( ++i <= $# )) && pos+=({$i..$#}); break;;
    -[OG]) must_passthrough=1; ((++i));;
    -*) [[ bcDEeFIiJLlmOopQRSWw == *${${*[i]}[-1]}* ]] && ((++i));;
    *)  pos+=($i);;
  esac
done

local z4h_min_version
z4h_min_version=${$(<$Z4H/zsh4humans/version)%$'\r'} || return
if [[ $z4h_min_version != <1-> ]]; then
  print -Pru2 -- '%F{3}z4h%f: invalid file content: %F{1}%U$Z4H/zsh4humans/version%f%u'
  return 1
fi
local -r z4h_min_version

local -r z4h_ssh_client=${${(%):-%m}:-unknown}

local z4h_ssh_host
if (( $#pos == 1 )); then
  local user_host=$*[pos[1]]
  z4h_ssh_host=${${user_host##*@}%%:*}
fi
local -r z4h_ssh_host

[[ -n $z4h_ssh_host ]] || must_passthrough=1

local -i z4h_ssh_enable=$(( !must_passthrough ))
zstyle -t :z4h:ssh:$z4h_ssh_host enable || z4h_ssh_enable=0

local -i mkdir_control_master=0
local default_ssh_command=(command ssh)
if (( z4h_ssh_enable )); then
  mkdir_control_master=1
  default_ssh_command+=(
    -o ControlMaster=auto
    -o ControlPersist=5
    -o ControlPath='~/.ssh/s/%C')
fi

local -a z4h_ssh_command
if ! zstyle -a :z4h:ssh:$z4h_ssh_host ssh-command z4h_ssh_command; then
  z4h_ssh_command=($default_ssh_command)
fi

local term
zstyle -s :z4h:ssh:$z4h_ssh_host term term || term=${TERM:/tmux-256color/screen-256color}

local -A z4h_ssh_send_files z4h_ssh_retrieve_files
local -a z4h_ssh_prelude z4h_ssh_setup z4h_ssh_run z4h_ssh_teardown
local -aU z4h_retrieve_history

if (( !must_passthrough )); then
  z4h_ssh_prelude=(
    '"export" ZDOTDIR="$HOME"'
    'if command -v "locale" >"/dev/null" 2>&1; then
      "export" LC_ALL="C"
    fi')

  z4h_ssh_send_files=(
    $ZDOTDIR/.zshenv   '"$ZDOTDIR"/.zshenv'
    $ZDOTDIR/.zprofile '"$ZDOTDIR"/.zprofile'
    $ZDOTDIR/.zshrc    '"$ZDOTDIR"/.zshrc'
    $ZDOTDIR/.zlogin   '"$ZDOTDIR"/.zlogin'
    $ZDOTDIR/.zlogout  '"$ZDOTDIR"/.zlogout')

  local file
  for file in $ZDOTDIR/.p10k{,-ascii}{,-8color}.zsh(N) $ZDOTDIR/.zsh_history.*:$z4h_ssh_host(N); do
    z4h_ssh_send_files[$file]='"$ZDOTDIR"/'${(q)file:t}
  done

  local -a extra_files
  if zstyle -a :z4h:ssh:$z4h_ssh_host send-extra-files extra_files; then
    local src dst
    for dst in $extra_files; do
      eval "src=$dst"
      z4h_ssh_send_files[$src]=$dst
    done
  fi

  z4h_ssh_run=(
    'if "[" "-f" "$ZDOTDIR"/.zshenv "-a" "-r" "$ZDOTDIR"/.zshenv "]"; then
      "." "$ZDOTDIR"/.zshenv
    else
      >&2 "printf" "\\033[33mz4h\\033[0m: not a readable file: \\033[31m%s\033[0m\n" "$ZDOTDIR"/.zshenv
      "false"
    fi')

  if zstyle -a :z4h:ssh:$z4h_ssh_host retrieve-extra-files extra_files; then
    local src dst
    for src in $extra_files; do
      eval "dst=$src"
      z4h_ssh_retrieve_files[$src]=$dst
    done
  fi

  zstyle -a :z4h:ssh:$z4h_ssh_host retrieve-history z4h_retrieve_history || z4h_retrieve_history=()
fi

local configure
if zstyle -s :z4h:ssh:$z4h_ssh_host configure configure; then
  eval $configure || return
elif (( $+functions[z4h-ssh-configure] )); then
  z4h-ssh-configure || return
fi

if (( ! $#z4h_ssh_command )); then
  print -Pru2 -- '%F{3}z4h%f: empty %F{1}z4h_ssh_command%f'
  return 1
fi

if [[ $mkdir_control_master == 1 &&
      ${(pj:\0:)z4h_ssh_command} == ${(pj:\0:)default_ssh_command} &&
      ! -d ~/.ssh/s && -n ~(#qNU) ]]; then
  zf_mkdir -pm 700 ~/.ssh/s || return
  {
    >~/.ssh/s/README <<\END
This directory has been created by `z4h ssh`. It stores control sockets
for SSH connections. See ControlMaster, ControlPath and ControlPersist
in `man ssh_config`. This directory must not be writable by anyone other
than the current user.
END
  } || return
  if [[ -e ~/.ssh/control-master/README ]]; then
    {
      >>~/.ssh/s/README <<\END

You might also have ~/.ssh/control-master with the same file in it. This
is the old directory that was used by zsh4humans for the same purpose
before 2021-11-14. If you don't reference that directory explicitly from
~/.ssh/config, you can safely delete it.
END
    } || return
  fi
fi

if (( must_passthrough || !z4h_ssh_enable )); then
  TERM=${term:-$TERM} "${z4h_ssh_command[@]}" "$@"
  return
fi

if (( $#z4h_retrieve_history )); then
  local local_hist_tmp=$Z4H/tmp/ssh-history.tmp.$sysparams[pid]
  z4h_ssh_retrieve_files[\$HISTFILE]=$local_hist_tmp
  zf_rm -f -- $local_hist_tmp || return
else
  local local_hist_tmp=
fi

local file
for file in "${(@kv)z4h_ssh_send_files}"; do
  if [[ -z $file ]]; then
    print -Pru2 -- '%F{3}z4h%f: empty element(s) in %F{1}z4h_ssh_send_files%f'
    return 1
  fi
  if [[ $file == */ ]]; then
    print -Pru2 -- "%F{3}z4h%f: element(s) of %Bz4h_ssh_send_files%b end with %B/%b: %F{1}${file//\%/%%}%f"
    return 1
  fi
done
for file in "${(@kv)z4h_ssh_retrieve_files}"; do
  if [[ -z $file ]]; then
    print -Pru2 -- '%F{3}z4h%f: empty element(s) in %F{1}z4h_ssh_retrieve_files%f'
    return 1
  fi
  if [[ $file == */ ]]; then
    print -Pru2 -- "%F{3}z4h%f: element(s) of %Bz4h_ssh_retrieve_files%b end with %B/%b: %F{1}${file//\%/%%}%f"
    return 1
  fi
done
for file in "${(@)z4h_retrieve_history}"; do
  if [[ -z $file ]]; then
    print -Pru2 -- '%F{3}z4h%f: empty element(s) in %F{1}z4h_retrieve_history%f'
    return 1
  fi
  if [[ $file == */ ]]; then
    print -Pru2 -- "%F{3}z4h%f: element(s) of %Bz4h_retrieve_history%b end with %B/%b: %F{1}${file//\%/%%}%f"
    return 1
  fi
  if [[ -e $file ]]; then
    if [[ ! ( -f $file && -r $file && -w $file ) ]]; then
      print -Pru2 -- "%F{3}z4h%f: element of %Bz4h_retrieve_history%b is not a readable & writable file: %F{1}${file//\%/%%}%f"
      return 1
    fi
  elif [[ -d ${file:h} ]]; then
    if [[ ! -w ${file:h} ]]; then
      print -Pru2 -- "%F{3}z4h%f: element of %Bz4h_retrieve_history%b is in a non-writable directory: %F{1}${file//\%/%%}%f"
      return 1
    fi
  else
    zf_mkdir -p -- ${file:h} || return
  fi
done

if (( $#z4h_ssh_retrieve_files && ! $+commands[base64] )); then
  print -Pru2 -- '%F{3}z4h%f: command not found: %F{1}base64%f'
  return 1
fi

local tmpdir
if (( $+commands[mktemp] )); then
  tmpdir=$(command mktemp -d -- $Z4H/tmp/ssh.XXXXXXXXXX) || return
else
  tmpdir=$Z4H/tmp/ssh.tmp.$sysparams[pid]
  zf_rm -rf -- $tmpdir || return
  zf_mkdir -- $tmpdir || return
fi

{
  local -i i=0
  local src dst
  local indices=() send_to=()
  for src dst in ${(kv)z4h_ssh_send_files}; do
    (( ++i ))
    send_to+=($dst)
    [[ -e $src ]] || continue
    local target=${src:A}
    if [[ -z $target(#qN.) && -z $target(#qN/) ]]; then
      print -Pru2 -- "%F{3}z4h%f: unsupported file type: %F{1}${src//\%/%%}%f"
      return 1
    fi
    if [[ ${tmpdir:A} == $target(|/*) ]]; then
      print -Pru2 -- "%F{3}z4h%f: cannot send file: %F{1}${src//\%/%%}%f"
      return 1
    fi
    zf_ln -s -- $target $tmpdir/$i || return
    indices+=($i)
  done

  local -a retrieve_from retrieve_to
  local from to
  for from to in ${(kv)z4h_ssh_retrieve_files}; do
    retrieve_from+=($from)
    retrieve_to+=($to)
  done

  local dump_marker=${(%):-%n}.$sysparams[pid].$EPOCHSECONDS.$RANDOM

  local script
  script=${"$(<$Z4H/zsh4humans/sc/ssh-bootstrap)"//$'\r'} || return
  script=${script//'^TERM^'/${(q)term}}
  script=${script//'^MIN_VERSION^'/${(q)z4h_min_version}}
  script=${script//'^SSH_HOST^'/${(q)z4h_ssh_host}}
  script=${script//'^SSH_CLIENT^'/${(q)z4h_ssh_client}}
  script=${script//'^SSH_ARGS^'/${(q)${(j: :)@}}}
  script=${script//'^PRELUDE^'/${(F)z4h_ssh_prelude}}
  script=${script//'^SEND_TO^'/${(j: :)send_to}}
  script=${script//'^SETUP^'/${(F)z4h_ssh_setup}}
  script=${script//'^RUN^'/${(F)z4h_ssh_run}}
  script=${script//'^TEARDOWN^'/${(F)z4h_ssh_teardown}}
  if (( $#retrieve_from )); then
    script=${script//'^EMPTY_RETRIEVE_FROM^'/"'false'"}
  else
    script=${script//'^EMPTY_RETRIEVE_FROM^'/"'true'"}
  fi
  script=${script//'^RETRIEVE_FROM^'/${(j: :)retrieve_from}}
  script=${script//'^DUMP_MARKER^'/${(q)dump_marker}}
  script=${script//'^CAN_SAVE_RESTORE_SCREEN^'/${_z4h_can_save_restore_screen}}

  script=${script//'^DUMP_POS^'/${(r:8:: :)${#script}}}

  print -r -- $script >$tmpdir/script || return

  local tar_v tar_c_opt tar_x_opt
  if tar_v=$(command tar --version 2>/dev/null) && [[ $tar_v == *'GNU tar'* ]]; then
    tar_c_opt=(--owner=0 --group=0)
    tar_x_opt=(--warning=no-unknown-keyword --warning=no-timestamp --no-same-owner)
  fi

  if (( $#indices )); then
    command tar -C $tmpdir $tar_c_opt -czhf - -- $indices >>$tmpdir/script || return
  fi

  local args=("$@")
  args[pos[1],pos[1]-1]=('-T')
  local remote_script=.z4h-ssh.${(%):-%n}.$sysparams[pid].$EPOCHSECONDS.$RANDOM

  # Tricky corner cases where this command must work:
  #
  #   1. The remote shell is csh (default on FreeBSD).
  #   2. There is no /tmp on the remote host (e.g., Termux).
  #   3. TMPDIR is not set.
  #   4. TMPDIR has spaces in it.
  #
  # The next command (the one that invokes /bin/sh) must also work in these
  # cases. It should also propagate the exit status of /bin/sh.
  local cmd="test -w /tmp && cat >/tmp/$remote_script && echo 1 && exit"
  cmd+=" || "
  cmd+="test ! -e /tmp/$remote_script && cat >~/$remote_script && echo 2"
  local loc
  loc=$("${z4h_ssh_command[@]}" "${args[@]}" $cmd <$tmpdir/script) || return
} always {
  zf_rm -rf -- $tmpdir
}

case ${loc//[[:space:]]} in
  1) remote_script="/tmp/$remote_script";;
  2) remote_script="~/$remote_script";;
  *)
    print -Pru2 -- "%F{3}z4h%f: failed to upload bootstrap script"
    return 1
  ;;
esac

args[pos[1]]='-t'

local stty
if [[ -v commands[stty] && -v _z4h_tty_fd ]]; then
  stty=$(command stty -g <&$_z4h_tty_fd 2>/dev/null) || stty=
fi

{ (  # subshell to avoid TTOU

local -i bypass=0

local -i pid=$sysparams[pid]

{
  setopt no_multi_os
  "${z4h_ssh_command[@]}" "${args[@]}" "sh $remote_script" 2>&1 1>&3 |
     LC_ALL=C command grep -vxE '(Shared c|C)onnection to .* closed\.(.)?' >&2
   return $pipestatus[1]
} 3>&1 | {
  if (( $+commands[mktemp] )); then
    tmpdir=$(command mktemp -d -- $Z4H/tmp/ssh.XXXXXXXXXX) || return
  else
    tmpdir=$Z4H/tmp/ssh.tmp.$sysparams[pid]
    zf_rm -rf -- $tmpdir || return
    zf_mkdir -- $tmpdir || return
  fi

  unsetopt multibyte
  local LC_ALL=C

  unset _z4h_saved_screen

  {
    local buf=
    local mark=$'\001z4h.'$dump_marker
    while true; do
      [[ -n $buf ]] || sysread 'buf[$#buf+1]' || return $(( $? != 5 ))
      if [[ $buf != *$mark[1]* ]]; then
        print -rn -- $buf
        buf=
        continue
      fi
      while true; do
        print -rn -- ${buf%%$mark[1]*}
        buf=$mark[1]${buf#*$mark[1]}
        local -i prefix=$(($#buf < $#mark ? $#buf : $#mark))
        (( prefix )) || continue
        [[ ${buf:0:$prefix} == ${mark:0:$prefix} ]] && break
        print -rn -- $buf[1]
        buf[1]=""
        continue 2
      done
      while (( $#buf < $#mark )) && [[ $mark == $buf* ]]; do
        # What should we do if the output ends with a proper prefix of mark?
        # Print it or not? Return an error or not? We choose to not print and return
        # success iff we've reached eof.
        sysread -s $(($#mark - $#buf)) 'buf[$#buf+1]' && continue
        return $(( $? != 5 ))
      done
      if [[ $buf != $mark* ]]; then
        print -rn -- $buf[1]
        buf[1]=""
        continue
      fi
      buf[1,$#mark]=""
      while (( $#buf < 16 )); do
        sysread -s $((16 - $#buf)) 'buf[$#buf+1]' && continue
        return $(( $? != 5 ))
      done

      {
        case ${buf[1,16]%% #} in
          bypass)
            bypass=1
            break
          ;;
          save-screen)
            (( _z4h_can_save_restore_screen )) || return
            local _z4h_saved_screen=
            -z4h-save-screen || return
            _z4h_saved_screen+=x
            continue
          ;;
          restore-screen)
            [[ -n $_z4h_saved_screen ]] || return
            _z4h_saved_screen[-1]=
            -z4h-restore-screen || return
            unset _z4h_saved_screen
            continue
          ;;
          <->)
            local -i len=buf[1,16]
          ;;
          *)
            return 1
          ;;
        esac
      } always {
        buf=${buf:16}
      }

      (( len )) || continue
      if [[ -d $tmpdir ]]; then
        local dump_file=$tmpdir/dump.base64
      else
        local dump_file=/dev/null
      fi
      {
        local -i n=$((len < $#buf ? len : $#buf))
        print -rn -- $buf[1,n] || return
        (( len -= n ))
        buf[1,n]=""
        while (( len )); do
          sysread -s $((len > 65636 ? 65636 : len)) -o 1 -c n || return $(( $? != 5 ))
          (( len -= n, 1 ))
        done
      } >$dump_file || return
      if [[ $dump_file != /dev/null ]]; then
        if base64 -d <<<'Cg==' &>/dev/null; then
          local base64_opt=-d
        else
          local base64_opt=-D
        fi
        <$tmpdir/dump.base64 command base64 $base64_opt |
          command tar -C $tmpdir $tar_x_opt -xzf - || return
        local -i i
        for i in {1..$#retrieve_to}; do
          local src=$tmpdir/$i
          local dst=$retrieve_to[i]
          [[ -e $src ]] || continue
          if [[ -e $dst ]]; then
            zf_rm -rf -- $dst || return
          fi
          if ! command mv -f -- $src $dst 2>/dev/null; then
            command cp -rf -- $src $dst || return
          fi
        done
        if [[ -s $local_hist_tmp ]]; then
          local local_hist
          for local_hist in $z4h_retrieve_history; do
            local TMPPREFIX=$local_hist
            () {
              () { fc -pa -- $1 $HISTSIZE $SAVEHIST } $1 && zf_mv -f -- $1 $local_hist
            } =(command cat -- $local_hist(N) $local_hist_tmp) || return
          done
        fi
      fi
    done
  } always {
    local -i err=$?
    zf_rm -rf -- $tmpdir $local_hist_tmp
    if (( err )); then
      kill -- -$pid 2>/dev/null
    fi
  }
}

if (( bypass )); then
  setopt no_multi_os
  {
    TERM=${term:-$TERM} "${z4h_ssh_command[@]}" "${args[@]}" 2>&1 1>&3 |
      LC_ALL=C command grep -vxE '(Shared c|C)onnection to .* closed\.(.)?' >&2
    return $pipestatus[1]
  } 3>&1
fi

 ) } always {
  [[ -n $stty ]] && command stty $stty <&$_z4h_tty_fd 2>/dev/null
}