#compdef dcsctl

# --- helpers ---------------------------------------------------------------

_dcsctl__docker_contexts() {
  local -a ctxs
  ctxs=(${(f)"$(command docker context ls --format '{{.Name}}' 2>/dev/null)"})
  _describe -t docker-contexts 'docker contexts' ctxs
}

_dcsctl__current_context() {
  command docker context show 2>/dev/null
}

_dcsctl__control_plane_base() {
  # Match dcsctl logic: base is $DCS_ROOT if set, else ~/.dcs
  print -r -- "${DCS_ROOT:-$HOME/.dcs}"
}

_dcsctl__projects() {
  local ctx base
  ctx="$(_dcsctl__current_context)"
  [[ -z "$ctx" ]] && return 0

  base="$(_dcsctl__control_plane_base)/$ctx/compose"
  [[ -d "$base" ]] || return 0

  local -a ignored projects
  ignored=(shared '@Recycle' '@Recently-Snapshot' '.DS_Store')

  projects=()
  local dir name
  for dir in "$base"/*; do
    [[ -d "$dir" ]] || continue
    name="${dir:t}"
    (( ${ignored[(I)$name]} )) && continue
    projects+=("$name")
  done

  _describe -t dcs-projects 'projects' projects
}

_dcsctl__secret_names_from_dir() {
  local ctx proj base projdir secretdir
  ctx="$(_dcsctl__current_context)"
  [[ -z "$ctx" ]] && return 1

  proj="${words[4]}"   # dcsctl secret add <project> <name>
  [[ -z "$proj" ]] && return 1

  base="$(_dcsctl__control_plane_base)/$ctx/compose"
  projdir="$base/$proj"
  secretdir="$projdir/secrets"
  [[ -d "$secretdir" ]] || return 1

  local -a names
  names=()
  local f n
  for f in "$secretdir"/*.txt; do
    [[ -f "$f" ]] || continue
    n="${f:t}"
    n="${n%.txt}"
    names+=("$n")
  done

  (( ${#names} == 0 )) && return 1
  _describe -t dcs-secrets 'secret names' names
}

# List service subdirectories in a swarm-mode project.
# Takes one argument: the project name.
_dcsctl__swarm_services() {
  local ctx proj base projdir svcdir
  ctx="$(_dcsctl__current_context)"
  [[ -z "$ctx" ]] && return 1

  proj="$1"
  [[ -z "$proj" ]] && return 1

  base="$(_dcsctl__control_plane_base)/$ctx/compose"
  projdir="$base/$proj"
  svcdir="$projdir/services"
  [[ -d "$svcdir" ]] || return 1

  local -a svcs
  svcs=()
  local dir name
  for dir in "$svcdir"/*/; do
    [[ -d "$dir" ]] || continue
    name="${dir:t}"
    svcs+=("$name")
  done

  (( ${#svcs} == 0 )) && return 1
  _describe -t dcs-swarm-services 'services' svcs
}

# Detect whether a project is in swarm mode (has services/ dir).
_dcsctl__is_swarm() {
  local ctx proj base projdir
  ctx="$(_dcsctl__current_context)"
  [[ -z "$ctx" ]] && return 1

  proj="$1"
  [[ -z "$proj" ]] && return 1

  base="$(_dcsctl__control_plane_base)/$ctx/compose"
  projdir="$base/$proj"
  [[ -d "$projdir/services" ]]
}

# Fetch services for a given project by asking docker compose for the resolved config.
_dcsctl__services() {
  local ctx proj base projdir
  ctx="$(_dcsctl__current_context)"
  [[ -z "$ctx" ]] && return 1

  proj="${words[3]}"   # dcsctl run <project> ...
  [[ -z "$proj" ]] && return 1

  base="$(_dcsctl__control_plane_base)/$ctx/compose"
  projdir="$base/$proj"
  [[ -d "$projdir" ]] || return 1

  local env1 env2 env3
  env1="$projdir/.env"
  env2="$projdir/service.env"
  env3="$projdir/service.secrets.env"

  [[ -f "$env1" && -f "$env2" && -f "$env3" && -f "$projdir/docker-compose.yml" ]] || return 1

  local -a svcs
  svcs=(${(f)"$(
    command docker --context "$ctx" compose \
      --project-directory "$projdir" \
      --project-name "$proj" \
      --env-file "$env1" \
      --env-file "$env2" \
      --env-file "$env3" \
      config --services 2>/dev/null
  )"})

  (( ${#svcs} == 0 )) && return 1
  _describe -t dcs-services 'services' svcs
}

# Heuristic: decide if we are in a position where a service name is expected.
_dcsctl__maybe_complete_service() {
  local subcmd cur
  subcmd="${words[4]}"   # dcsctl run <project> <subcmd> ...
  cur="${words[CURRENT]}"

  [[ "$cur" == -* ]] && return 1

  case "$subcmd" in
    logs|ps|stop|start|restart|up|pull|build|rm|kill|pause|unpause|top|exec)
      _dcsctl__services
      return $?
      ;;
  esac

  return 1
}

# Delegate completion to docker after we reconstruct words as docker compose ...
_dcsctl__docker_compose_delegate() {
  local ctx proj base projdir
  ctx="$(_dcsctl__current_context)"
  [[ -z "$ctx" ]] && return 1

  proj="${words[3]}"
  [[ -z "$proj" ]] && return 1

  base="$(_dcsctl__control_plane_base)/$ctx/compose"
  projdir="$base/$proj"

  local -a dc_words
  dc_words=(docker --context "$ctx" compose --project-directory "$projdir" --project-name "$proj")

  if (( ${#words} >= 4 )); then
    dc_words+=("${words[@]:3}")  # from 4th element
  fi

  local -a save_words
  local save_cword
  save_words=("${words[@]}")
  save_cword=$CURRENT

  words=("${dc_words[@]}")
  CURRENT=${#words}

  autoload -Uz _docker
  _docker

  words=("${save_words[@]}")
  CURRENT=$save_cword
}

# Helper for edit --service completion: extract the project arg from the command
# line and complete swarm service names.
_dcsctl__swarm_services_for_edit() {
  local proj=""
  local i
  for (( i=3; i <= ${#words}; i++ )); do
    local w="${words[$i]}"
    case "$w" in
      -e|--env|-r|--runtime|-s|--secret) ;;
      --service) (( i++ )) ;;  # skip the next arg (the service value)
      -*) ;;
      *)
        if [[ -z "$proj" ]]; then
          proj="$w"
        fi
        ;;
    esac
  done

  [[ -z "$proj" ]] && return 1
  _dcsctl__swarm_services "$proj"
}

# --- main completion -------------------------------------------------------

_dcsctl() {
  local -a subcmds
  subcmds=(
    'new:Create a new project scaffold'
    'verify:Verify and enforce host permissions for appdata bind mounts'
    'run:Run docker compose for a project (compose mode)'
    'up:Alias for run <project> up -d (compose mode)'
    'down:Bring down a project (compose or swarm)'
    'stop:Scale all services to 0 replicas (swarm projects only)'
    'deploy:Merge service fragments and deploy as Docker Swarm stack'
    'ls:Alias for docker compose ls'
    'dir:Print the project directory, or a service directory in a swarm project'
    'edit:Edit project files in $EDITOR'
    'secret:Manage project secrets'
    'service:Manage services within a swarm project'
    'import:Import an existing docker-compose.yml into a DCS project'
    'context:Manage DCS control-plane contexts'
    'reload:Down, rebuild and up all containers within a DCS project'
    'rename:Rename and update a DCS project'
    'help:Show help for a command'
  )

  if (( CURRENT == 2 )); then
    _describe -t commands 'dcsctl commands' subcmds
    return
  fi

  local cmd
  cmd="${words[2]}"

  case "$cmd" in
    new)
      _arguments -C \
        '(-v --verify)'{-v,--verify}'[Run verification after creating the project]' \
        '--verbose[Verbose output for verification (used with --verify)]' \
        '--compose[Create flat single-compose layout (legacy)]' \
        '--service[Initial service name for swarm layout]:service name:' \
        '1:project name:' \
        && return
      ;;

    verify)
      _arguments -C \
        '--verbose[Show current vs expected ownership/mode while verifying (QNAP only)]' \
        '1:project name:_dcsctl__projects' \
        && return
      ;;

    run)
      if (( CURRENT == 3 )); then
        _dcsctl__projects
        return
      fi

      if (( CURRENT >= 5 )); then
        _dcsctl__maybe_complete_service && return
      fi

      _dcsctl__docker_compose_delegate
      return
      ;;

    up)
      _arguments -C \
        '(-v --verify)'{-v,--verify}'[Run `dcsctl verify` before running docker compose]' \
        '1:project name:_dcsctl__projects' \
        && return
      ;;

    down)
      _arguments -C \
        '1:project name:_dcsctl__projects' \
        && return
      ;;

    stop)
      _arguments -C \
        '1:project name:_dcsctl__projects' \
        && return
      ;;

    deploy)
      _arguments -C \
        '(-v --verify)'{-v,--verify}'[Run verification before deploying]' \
        '--verbose[Verbose verification output]' \
        '1:project name:_dcsctl__projects' \
        && return
      ;;

    ls)
      return
      ;;

    dir)
      if (( CURRENT == 3 )); then
        _dcsctl__projects
        return
      fi
      if (( CURRENT == 4 )); then
        local proj="${words[3]}"
        _dcsctl__swarm_services "$proj" && return
        return 0
      fi
      ;;

    edit)
      _arguments -C \
        '(-e --env -r --runtime -s --secret)'{-e,--env}'[Edit .env]' \
        '(-e --env -r --runtime -s --secret)'{-r,--runtime}'[Edit service.env]' \
        '(-e --env -r --runtime -s --secret)'{-s,--secret}'[Edit service.secrets.env]' \
        '--service[Target a specific service (swarm mode)]:service name:_dcsctl__swarm_services_for_edit' \
        '1:project name:_dcsctl__projects' \
        && return
      ;;

    secret)
      if (( CURRENT == 3 )); then
        _describe -t secret-subcmds 'secret commands' \
          'add:Create/open secrets/<name>.txt in $EDITOR'
        return
      fi

      local sub
      sub="${words[3]}"
      case "$sub" in
        add)
          if (( CURRENT == 4 )); then
            _dcsctl__projects
            return
          fi
          if (( CURRENT == 5 )); then
            _dcsctl__secret_names_from_dir && return
            return 0
          fi
          ;;
      esac
      ;;

    service)
      if (( CURRENT == 3 )); then
        local -a service_subcmds
        service_subcmds=(
          'add:Add a new service to a swarm project'
          'ls:List services in a swarm project'
        )
        _describe -t service-subcmds 'service commands' service_subcmds
        return
      fi

      local sub
      sub="${words[3]}"
      case "$sub" in
        add)
          if (( CURRENT == 4 )); then
            _dcsctl__projects
            return
          fi
          # CURRENT == 5: free-form service name, no completions needed
          ;;
        ls)
          if (( CURRENT == 4 )); then
            _dcsctl__projects
            return
          fi
          ;;
      esac
      ;;

    import)
      _arguments -C \
        '--env-file[Additional env file to ingest]:env file:_files' \
        '--strict[Fail if compose violates template contract after import]' \
        '--swarm[Import as swarm project, splitting services into per-service fragments]' \
        '1:project name:_dcsctl__projects' \
        '2:compose path:_files' \
        && return
      ;;

    context)
      if (( CURRENT == 3 )); then
        _describe -t context-subcmds 'context commands' \
          'init:Initialize ~/.dcs/<context> with env.system and compose/'
        return
      fi

      local sub
      sub="${words[3]}"
      case "$sub" in
        init)
          if (( CURRENT == 4 )); then
            _dcsctl__docker_contexts
            return
          fi
          ;;
      esac
      ;;

    reload)
      _arguments -C \
        '1:project name:_dcsctl__projects' \
        && return
      ;;

    rename)
      _arguments -C \
        '1:project name:_dcsctl__projects' \
        '2:new project name:' \
        && return
      ;;

    help)
      if (( CURRENT == 3 )); then
        _describe -t commands 'dcsctl commands' subcmds
        return
      fi
      ;;
  esac

  _describe -t commands 'dcsctl commands' subcmds
}

_dcsctl "$@"