#compdef sqruff

_sqruff() {
  local curcontext="$curcontext" state line
  typeset -A opt_args

  local -a dialects=(
    ansi athena bigquery clickhouse databricks db2 duckdb mysql oracle
    postgres redshift snowflake sparksql sqlite trino tsql
  )

  local -a common=(
    '--config=[path to a configuration file]:config file:_files'
    "--dialect=[override the dialect]:dialect:(${dialects})"
    '--parsing-errors[show parse errors]'
    '(-h --help)'{-h,--help}'[print help]'
  )

  _arguments -C \
    $common \
    '(-V --version)'{-V,--version}'[print version]' \
    '1: :->command' \
    '*:: :->args'

  case $state in
    command)
      local -a commands=(
        'lint:Lint SQL files via passing a list of files or using stdin'
        'fix:Fix SQL files via passing a list of files or using stdin'
        'lsp:Run an LSP server'
        'info:Print information about sqruff and the current environment'
        'rules:Explain the available rules'
        'dialects:List available dialects'
        'templaters:List available templaters'
        'help:Print this message or the help of the given subcommand(s)'
      )
      _describe -t commands 'sqruff command' commands
      ;;
    args)
      case $line[1] in
        lint|fix)
          _arguments \
            $common \
            '(-f --format)'{-f+,--format=}'[output format for the results]:format:(human github-annotation-native json)' \
            '*:path:_files'
          ;;
        lsp|info|rules|dialects|templaters)
          _arguments $common
          ;;
        help)
          local -a sub=(lint fix lsp info rules dialects templaters help)
          _describe -t commands 'command' sub
          ;;
      esac
      ;;
  esac
}

_sqruff "$@"
