use path
use re
use str
use ../mods/common
use ../mods/list
use ../mods/map
use ../mods/option

var options-ssh = [ 1 2 4 6 A D f g I i L l m o v a b C c e F k N n p q R s T t X x ]
var options-scp = [ 3 4 6 B C p q r v c F i l o P S ]

fn -kh {
  cat ~/.ssh/known_hosts | peach {|l|
    put [(str:split ' ' $l)]
  } | peach {|e|
    var domains @_ = $@e
    str:split ',' $domains
  }
}

fn -port {|cmd @argv|
  var o     = (common:cond (eq $cmd 'ssh') '-p' '-P')
  var margs = (option:map $argv)
  var p     = (map:value-of $margs $o &default=[])
  cond:cexec (list:is-empty $p) 22 { put $p[-1] }
}

fn -complete-names {
  var fp = $E:HOME/.config/elvish/private/sshnames
  if (path:is-regular $fp) {
    cat $fp | from-lines | each {|n| put (printf '%s@' $n) }
  }
}

fn -complete-domains {|name hosts|
  each {|h|
    put (printf '%s@%s' $name $h)
  } $hosts
}

fn -complete-remote-dir {|port address dir|
  var cmd = (printf ^
    "for f in '%s'*; do if [[ -d $f ]]; then echo $f/; else echo $f; fi; done" ^
    $dir)
  try {
    ssh -p $port $address $cmd | each {|f|
      put (printf '%s:%s' $address $f)
    }
  } except e { }
}

fn -complete-args {|hosts cmd @argv|
  var arg = $argv[-1]
  var i   = (str:index $arg @)
  if (< $i 0) {
    -complete-names
    all $hosts
    if (eq $cmd scp) {
      edit:complete-filename $cmd $@argv
    }
    return
  }
  var n h = $arg[..$i] $arg[(+ $i 1)..]
  if (eq $cmd scp) {
    set i = (str:index $h :)
    if (>= $i 0) {
      var d = $h[(+ $i 1)..]
      set h = $h[..$i]
      if (list:contains $h $hosts) {
        var p = (-port $cmd @argv)
        -complete-remote-dir $p $n@$h $d
      }
      return
    }
  }
  -complete-domains $n $hosts
}

fn complete {|@argv|
  var @hosts = (-kh)
  var cmd    = $argv[0]
  var is-ssh = (eq $cmd ssh)
  var po     = (common:cond $is-ssh -p -P)
  if (<= (count $argv) 2) {
    all (common:cond $is-ssh $options-ssh $options-scp)
    -complete-args $hosts $@argv
  } elif (eq $argv[-2] $po) {
    put 22
  } else {
    -complete-args $hosts $@argv
  }
}

set edit:completion:arg-completer[scp] = $complete~
set edit:completion:arg-completer[ssh] = $complete~