use re
use str
use moi/util/condition
use moi/util/file
use moi/util/ip
use moi/util/list
use moi/util/map
use moi/util/option
 
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 ]
options-scp = [ 3 4 6 B C p q r v c F i l o P S ]

fn -init-hosts {
  local:result = [&]
  cat /etc/hosts | peach $str:trim-space~ | peach [l]{
    local:l = $l
    if (and (not-eq $l '') (not-eq $l[0] '#')) {
      put $l
    }
  } | eawk [_ ip @domains]{
    if (ip:is-ip $ip) {
      put [$ip $domains]
    }
  } | peach [e]{
    local:ip local:domains = $@e
    if (ip:is-ipv6 $ip) {
      ip = (ip:short6 $ip)
    }
    each [d]{
      if (eq $d[0] '#') {
        break
      }
      put [$ip $d]
    } $domains
  } | each [e]{
    local:ip local:d = $@e
    local:ipd = (map:value-of $result $ip &default=[&])
    ipd[$d] = $nil
    result[$ip] = $ipd
  }
  put $result
}

fn -init-known-hosts {
  local:khosts = [&22=[&]]
  cat ~/.ssh/known_hosts | peach [l]{
    put [(str:split ' ' $l)]
  } | peach [e]{
    local:domains @local:id = $@e
    id = (str:join ' ' $id)
    str:split ',' $domains | peach [d]{
      local:d = $d
      local:port = 22
      if (eq $d[0] '[') {
        local:i = (str:index $d ']')
        port = $d[(+ $i 2):]
        d = $d[1:$i]
      }
      put [$id $d $port]
    }
  } | each [e]{
    local:id local:d local:port = $@e
    local:pdomains = (map:value-of $khosts $port &default=[&])
    local:idomains = (map:value-of $pdomains $id &default=[&])
    idomains[$d] = $nil
    pdomains[$id] = $idomains
    khosts[$port] = $pdomains
  }
  local:result = [&22=[&]]
  keys $khosts | each [port]{
    local:pdomains = (map:value-of $result $port &default=[&])
    keys $khosts[$port] | peach [id]{
      local:id = $id
      local:domains = [(keys $khosts[$port][$id])]
      if (list:contains-not $ip:is-ip~ $domains) {
        list:filter-not $ip:is-ip~ $domains | peach [d]{ pdomains[$d] = $nil }
      } else {
        peach [d]{ pdomains[$d] = $nil } $domains
      }
    }
    result[$port] = $pdomains
  }
  put $result
}


fn -init-history {
  local:history = [&]
  edit:command-history | peach [e]{
    put $e[cmd]
  } | peach [h]{
    if (re:match '^(ssh|scp)\ .+' $h) {
      history[$h] = $nil
    }
  }
  local:result = [&22=[&]]
  keys $history | peach [h]{
    local:h = $h
    put [(re:split '[ \t]+' $h)]
  } | peach [e]{
    local:cmd @local:argv = $@e
    local:port local:name local:domain = 22 '' ''
    local:c = (- (count $argv) 1)
    list:loop [i arg]{
      if (and (< $i $c) (or (and (eq $cmd ssh) (eq $arg -p)) (eq $arg -P))) {
        port = $argv[(+ $i 1)]
      } else {
        local:nd = [(str:split '@' $arg)]
        if (== (count $nd) 2) {
          name = $nd[0]
          domain @_ = (str:split ':' $nd[1])
          break
        }
      }
    } $argv
    if (and (not-eq $name '') (not-eq $domain '')) {
      put [$port $name $domain]
    }
  } | each [e]{
    local:port local:name local:domain = $@e
    local:pnames = (map:value-of $result $port &default=[&])
    local:domains = (map:value-of $pnames $name &default=[&])
    domains[$domain] = $nil
    pnames[$name] = $domains
    result[$port] = $pnames
  }
  put $result
}

combinaisons = [&]

fn init {
  local:hosts local:known-hosts local:history = [] [] []
  peach [f]{ $f } [
    { hosts = (-init-hosts) }
    { known-hosts = (-init-known-hosts) }
    { history = (-init-history) }
  ]
  combinaisons = $history
  keys $known-hosts | peach [port]{
    local:port = $port
    local:pdomains = $known-hosts[$port]
    keys $pdomains | peach [domain]{
      local:domain = $domain
      if (and (ip:is-ip $domain) (has-key $hosts $domain)) {
        del pdomains[$domain]
        keys $hosts[$domain] | peach [d]{
          local:d = $d
          pdomains[$d] = $nil
        }
      }
    }
    if (> (keys $pdomains | count) 0) {
      if (has-key $history $port) {
        local:phist = $history[$port]
        local:dhist = (keys $phist | peach [name]{ keys $phist[$name] } | map:to-set)
        local:dknown = [&]
        keys $pdomains | peach [domain]{
          local:domain = $domain
          if (not (has-key $dhist $domain)) {
            dknown[$domain] = $nil
          }
        }
        if (> (keys $dknown | count) 0) {
          combinaisons[$port] = (assoc $phist '' $dknown)
        } else {
          combinaisons[$port] = [&''=$dknown]
        }
      } else {
        combinaisons[$port] = [&''=$pdomains]
      }
    }
  }
}

fn -known-ports { keys $combinaisons }

fn -known-names [&port=22]{
  keys (map:value-of $combinaisons $port &default=[&]) | list:filter [n]{ not-eq $n '' }
}

fn -known-hosts [&port=22 &name=$false]{
  local:phosts = (map:value-of $combinaisons $port &default=[&])
  if $name {
    keys (map:value-of $phosts $name &default=[&])
  } else {
    keys (keys $phosts | peach [name]{
      local:name = $name
      keys $phosts[$name]
    } | map:to-set)
  }
}

fn -port [&cmd=ssh @argv]{
  o = (condition:set (eq $cmd 'ssh') '-p' '-P')
  margs = (option:map $argv)
  map:value-of $margs $o &default=[]
}

fn -complete-options [list @args]{ each [o]{ put -$o } $list }

fn -complete-args [cmd @argv]{
  local:port = (-port $cmd $@argv)
  if (list:empty $port) {
    port = 22
  } else {
    port = $port[-1]
  }
  local:larg = $argv[-1]
  local:name @local:host = (str:split '@' $larg)
  if (list:empty $host) {
    -known-names &port=$port | each [n]{ put $n'@' }
    if (eq $cmd 'scp') {
      file:match-files $larg
    }
  } else {
    host @_ = (str:split ':' $host[0])
    local:check = [h]{ has-prefix $h $host }
    local:khosts = [(-known-hosts &port=$port &name=$name)]
    if (not (list:contains $check $khosts)) {
      khosts = [(-known-hosts &port=$port)]
      if (not (list:contains $check $khosts)) {
        khosts = [(-known-hosts)]
      }
    }
    each [h]{ put $name'@'$h } $khosts
    all $khosts
  }
}

fn complete [@argv]{
  if (keys $combinaisons | list:empty) {
    init
  }
  local:cmd = $argv[0]
  local:is-ssh = (eq $cmd ssh)
  local:po = (condition:set $is-ssh -p -P)
  if (<= (count $argv) 2) {
    -complete-options (condition:set $is-ssh $options-ssh $options-scp)
    -complete-args $@argv
  } elif (eq $argv[-2] $po) {
    -known-ports
  } else {
    -complete-args $@argv
  }
}

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

init&