#compdef socket

# Parameter used:
#
#  socket_ports
#    The associative array that maps a host name to a space-separated list of 
#    ports.


local state line expl
typeset -A options

_arguments -s \
  -{b,c,f,q,r,v,w} \
  -{s,l} \
  '-p:command:->command' \
  ':arg1:->arg1' \
  ':arg2:->arg2'

case "$state" in
command)
  compset -q
  _normal
  ;;

arg1)
  if (( $+options[-s] )); then
    _message 'port to listen'
  else
    _description expl 'host'
    compadd "$expl[@]" - ${(k)socket_ports} || _hosts "$expl[@]"
  fi
  ;;

arg2)
  if (( ! $+options[-s] )); then
    _description expl 'port to connect'
    if (( $+socket_ports )); then
      compadd "$expl[@]" - ${=socket_ports[$line[2]]};
    else
      _message 'port to connect';
    fi
  fi
  ;;
esac