#!/usr/bin/env bash

mounts=()
ssh_args=()
remote_port=$((10000 + RANDOM % 10000))
local_user=$(whoami)

usage()
{
  echo "sshp (\"SSH PLUS\") from lackadaisical."
  echo "Accepts all standard SSH options (see man ssh)."
  echo "Additionally, supports SSHFS mounts using a Docker-style syntax (-m)."
  echo "Usage: sshp -m <local>:<remote> [user@]host [ssh_options]"
  echo "Example: sshp -m /home/juli:/home/juli juli@juli.pyon"
  exit 1
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    -m)
      if [[ -z "$2" ]]
      then
        echo "Error: -m requires argument"
        exit 1
      fi
      mounts+=("$2")
      shift 2
      ;; 
    *)
      ssh_args+=("$1")
      shift
      ;; 
  esac
done

if [[ ${#ssh_args[@]} -eq 0 ]]
then
  usage
fi

mount_logic=""
unmount_logic=""

for map in "${mounts[@]}"; do
  local_path="${map%%:*}"
  remote_path="${map##*:}"

  if [[ "$local_path" != /* ]]
  then
    local_path="$PWD/$local_path"
  fi

  mount_logic+="
  echo '>> Preparing mount: ${remote_path}';
  if ! mkdir -p \"${remote_path}\" 2>/dev/null
  then
    sudo mkdir -p \"${remote_path}\"
  fi

  if ! sshfs -p ${remote_port} -o StrictHostKeyChecking=no,idmap=user ${local_user}@localhost:\"${local_path}\" \"${remote_path}\" 2>/dev/null
  then
    echo '   (User mount failed, attempting escalation...)';
    if ! sudo sshfs -p ${remote_port} -o StrictHostKeyChecking=no,idmap=user,allow_other ${local_user}@localhost:\"${local_path}\" \"${remote_path}\"
        then
      echo '   ! Mount failed completely. Check permissions or keys.';
    fi
  else
    echo '   > Mounted successfully.';
  fi
  "

  unmount_logic+="
  fusermount -u -z \"${remote_path}\" 2>/dev/null || sudo fusermount -u -z \"${remote_path}\" 2>/dev/null;
  "
done

remote_script="
  if ! command -v sshfs >/dev/null 2>&1
  then
    echo 'WARNING: \"sshfs\" not found on remote host.'
    echo '>> Skipping mounts, proceeding with shell only...';
    echo '----------------------------------------------';
  else
    ${mount_logic}
  fi

  ${SHELL:-bash};

  if command -v sshfs >/dev/null 2>&1
  then
    ${unmount_logic}
  fi
"

ssh -t -R ${remote_port}:localhost:22 "${ssh_args[@]}" "$remote_script"