#compdef kpf

_kpf() {
  local context state state_descr line
  typeset -A opt_args

  _arguments -S -C \
    '(-n --namespace)'{-n,--namespace}'[Kubernetes namespace to use]:namespace:_kpf_namespaces' \
    '(-A --all)'{-A,--all}'[Show all services across all namespaces]' \
    '(-l --all-ports)'{-l,--all-ports}'[Include ports from pods, deployments, daemonsets, etc.]' \
    '(-c --check)'{-c,--check}'[Check and display endpoint status in service selection table]' \
    '(-d --debug)'{-d,--debug}'[Enable debug output for troubleshooting]' \
    '(-t --debug-terminal)'{-t,--debug-terminal}'[Enable debug output for troubleshooting display issues]' \
    '--run-http-health-checks[Enable HTTP connectivity health checks (disabled by default)]' \
    '(-0)-0[Listen on all interfaces (0.0.0.0) instead of localhost]' \
    '(-pn --prompt-namespace)'{-pn,--prompt-namespace}'[Interactively select a namespace before service selection]' \
    '--auto-reconnect[Automatically reconnect on failure]' \
    '--auto-select-free-port[Automatically select a free local port]' \
    '--capture-usage[Capture usage statistics]' \
    '--multiline-command[Display command in multiple lines]' \
    '--reconnect-attempts[Number of reconnection attempts]:attempts:' \
    '--reconnect-delay[Delay between reconnection attempts in seconds]:delay:' \
    '--show-context[Show Kubernetes context in output]' \
    '--show-direct-command[Show the direct kubectl command]' \
    '--usage-folder[Folder to store usage statistics]:folder:_directories' \
    '(-v --version)'{-v,--version}'[Show version]' \
    '(-h --help)'{-h,--help}'[Show help]' \
    '1:service:->services' \
    '2:port:->ports'

  case $state in
    services)
      _kpf_services
      ;;
    ports)
      _kpf_ports
      ;;
  esac
}

_kpf_namespaces() {
  local -a namespaces
  namespaces=("${(@f)$(kubectl get namespaces -o jsonpath='{range .items[*]}{.metadata.name}{"\n"}{end}' 2>/dev/null)}")
  
  if (( ${#namespaces[@]} > 0 )); then
    _describe 'namespace' namespaces
  fi
}

_kpf_get_ns_opts() {
  # This function expects 'ns_opts' to be defined in the caller's scope
  ns_opts=()
  if (( ${words[(I)-A|--all]} )); then
    ns_opts=(-A)
    return
  fi

  local ns=${opt_args[-n]:-${opt_args[--namespace]}}
  if [[ -z "$ns" ]]; then
    local ns_idx=${words[(I)-n|--namespace]}
    if (( ns_idx && ns_idx < CURRENT )); then
      ns=${words[ns_idx+1]}
    fi
  fi
  
  if [[ -n "$ns" && "$ns" != "$PREFIX" ]]; then
    ns_opts=(-n "$ns")
  fi
}

_kpf_services() {
  local -a ns_opts
  _kpf_get_ns_opts

  local -a services
  services=("${(@f)$(kubectl get services "${ns_opts[@]}" -o jsonpath='{range .items[*]}{.metadata.name}{"\n"}{end}' 2>/dev/null)}")

  local -a resources
  for s in $services; do
    [[ -n "$s" ]] && resources+=("svc/$s")
  done

  if (( ${#resources[@]} > 0 )); then
    _describe 'service' resources
  fi
}

_kpf_ports() {
  # The service name is the first positional argument
  # When using _arguments -C, 'line' contains positional args
  local service="${line[1]}"
  
  # Strip svc/ prefix if present
  service="${service#svc/}"

  # Return early if no service specified yet
  [[ -z "$service" ]] && return 0

  local -a ns_opts
  _kpf_get_ns_opts

  # Get ports for the service
  local -a port_descriptions
  local port_data
  port_data=$(kubectl get service "$service" "${ns_opts[@]}" -o jsonpath='{range .spec.ports[*]}{.port}{"/"}{.protocol}{" "}{.name}{"\n"}{end}' 2>/dev/null)

  # Parse and format ports
  while IFS= read -r line_out; do
    [[ -z "$line_out" ]] && continue
    local port_proto="${line_out%% *}"
    local port_name="${line_out#* }"
    local port="${port_proto%%/*}"

    # Format as port:port for kubectl port-forward syntax
    # Use backslash to escape the colon for _describe
    if [[ -n "$port_name" && "$port_name" != "$port_proto" ]]; then
      port_descriptions+=("${port}\:${port}:${port_name} (${port_proto})")
    else
      port_descriptions+=("${port}\:${port}:${port_proto}")
    fi
  done <<< "$port_data"

  # Offer completions with descriptions
  if (( ${#port_descriptions[@]} > 0 )); then
    _describe 'port' port_descriptions
  fi
}

_kpf "$@"


