#compdef pwt

# pwt - Power Worktrees completion
# Install: add to fpath and run compinit

_pwt_commands() {
    local commands=(
        'init:Initialize project (from current repo or URL)'
        'create:Create new worktree'
        'add:Create new worktree (alias for create)'
        'track:Create worktree tracking an existing remote branch'
        'adopt:Register existing worktree and run setup'
        'setup:Register current worktree and run setup'
        'list:List worktrees and status'
        'ls:List worktrees and status'
        'tree:Visual tree view of worktrees'
        'status:Interactive TUI dashboard (like htop)'
        'cd:Navigate to worktree'
        'use:Set worktree as current'
        'current:Show current worktree path'
        '-:Go to previous worktree'
        'info:Show worktree details'
        'show:Show worktree details'
        'remove:Remove worktree'
        'rm:Remove worktree'
        'server:Start development server'
        's:Start development server'
        'gateway:Stable gateway to a worktree server'
        'servers:Show project server status'
        'run:Run command in worktree'
        'for-each:Run command in all worktrees'
        'editor:Open worktree in editor'
        'e:Open worktree in editor'
        'ai:Start AI tool in worktree'
        'open:Open worktree in Finder'
        'diff:Show diff between worktrees'
        'copy:Copy files between worktrees'
        'repair:Repair broken worktree'
        'fix:Repair broken worktree'
        'auto-remove:Remove merged worktrees'
        'cleanup:Remove merged worktrees'
        'restore:Recover backed up changes from trash'
        'fix-port:Resolve port conflict'
        'doctor:Check system health and configuration'
        'meta:Manage worktree metadata'
        'm:Manage worktree metadata (alias)'
        'alias:Set project short alias'
        'steps:List available Pwtfile steps'
        'step:Run a single Pwtfile step'
        'project:Manage project configs'
        'config:Configure current project'
        'port:Get port for worktree'
        'plugin:Manage pwt plugins'
        'claude-setup:Configure Claude Code integration'
        'setup-shell:Install shell integration'
        'shell-init:Output shell integration code'
        'jobs:Manage background jobs'
        'help:Show help'
    )
    _describe 'command' commands
}

_pwt_worktrees() {
    local project="${1:-}"
    local worktrees
    # Find the pwt binary (prefer binary over shell function for reliability)
    local pwt_bin="${commands[pwt]:-$(whence -p pwt)}"
    [[ -z "$pwt_bin" ]] && pwt_bin="pwt"

    # Use pwt list --names for accurate project detection
    # Pass --project if specified (for "pwt <project> cd <tab>" syntax)
    local raw_output
    if [[ -n "$project" ]]; then
        raw_output=$("$pwt_bin" --project "$project" list --names 2>/dev/null)
    else
        raw_output=$("$pwt_bin" list --names 2>/dev/null)
    fi

    # Parse output and strip trailing slashes from each worktree name
    # The --names output includes trailing / for display, but trailing / in completions
    # triggers zsh directory completion behavior which breaks worktree completion
    local line
    worktrees=()
    for line in ${(f)raw_output}; do
        # Strip trailing slash
        worktrees+=("${line%/}")
    done

    # Fallback to just @ if pwt fails (e.g., not in a project)
    [[ ${#worktrees[@]} -eq 0 ]] && worktrees=("@")
    _describe 'worktree' worktrees
}

_pwt_projects() {
    local projects
    if [[ -d ~/.pwt/projects ]]; then
        projects=(${(f)"$(ls ~/.pwt/projects/ 2>/dev/null)"})
        _describe 'project' projects
    fi
}

# Check if input matches a project (exact, partial, or alias)
_pwt_is_project() {
    local input="$1"
    # Exact match
    [[ -d ~/.pwt/projects/$input ]] && return 0
    # Partial match (e.g., "ticket" matches "ticket-app")
    local matches=(~/.pwt/projects/${input}*(N))
    [[ ${#matches[@]} -eq 1 ]] && return 0
    # Check aliases in all project configs
    for config in ~/.pwt/projects/*/config.json(N); do
        # Check "alias" (string) field
        local alias_val=$(jq -r '.alias // empty' "$config" 2>/dev/null)
        [[ "$alias_val" == "$input" ]] && return 0
        # Check "aliases" (array) field
        jq -e --arg a "$input" '.aliases // [] | index($a) != null' "$config" >/dev/null 2>&1 && return 0
    done
    return 1
}

# Resolve partial project name or alias to full name
_pwt_resolve_project() {
    local input="$1"
    # Exact match
    [[ -d ~/.pwt/projects/$input ]] && echo "$input" && return
    # Partial match
    local matches=(~/.pwt/projects/${input}*(N))
    if [[ ${#matches[@]} -eq 1 ]]; then
        basename "${matches[1]}"
        return
    fi
    # Check aliases in all project configs
    for config in ~/.pwt/projects/*/config.json(N); do
        local project=$(basename "$(dirname "$config")")
        # Check "alias" (string) field
        local alias_val=$(jq -r '.alias // empty' "$config" 2>/dev/null)
        [[ "$alias_val" == "$input" ]] && echo "$project" && return
        # Check "aliases" (array) field
        if jq -e --arg a "$input" '.aliases // [] | index($a) != null' "$config" >/dev/null 2>&1; then
            echo "$project"
            return
        fi
    done
}

_pwt_branches() {
    local branches
    branches=(${(f)"$(git branch -a 2>/dev/null | sed 's/^[* ]*//' | sed 's|remotes/||')"})
    _describe 'branch' branches
}

_pwt_meta_actions() {
    local actions=(
        'list:List all metadata'
        'show:Show metadata for worktree'
        'set:Update metadata field'
        'import:Import existing worktrees'
    )
    _describe 'action' actions
}

_pwt_project_actions() {
    local actions=(
        'list:List all projects'
        'init:Initialize new project'
        'show:Show project config'
        'set:Update project config'
        'path:Print config directory'
    )
    _describe 'action' actions
}

_pwt_config_keys() {
    local keys=(
        'main_app:Main app path'
        'worktrees_dir:Worktrees directory'
        'branch_prefix:Branch prefix (e.g. user/)'
        'base_port:Base port for allocation'
        'gateway_port:Stable gateway proxy port'
        'gateway_host:Gateway public URL host'
    )
    _describe 'key' keys
}

_pwt_gateway_actions() {
    local actions=(
        'init:Configure gateway port'
        'up:Start gateway proxy daemon'
        'down:Stop gateway proxy'
        'start:Alias for up'
        'stop:Alias for down'
        'restart:Restart gateway proxy'
        'status:Show gateway status'
        'use:Point gateway at a worktree'
        'url:Print gateway URL'
        'logs:Show gateway logs'
        'help:Show help'
    )
    _describe 'action' actions
}

_pwt_meta_fields() {
    local fields=(
        'base:Base branch'
        'description:Worktree description'
        'branch:Git branch'
        'port:Port number'
    )
    _describe 'field' fields
}

_pwt_remove_flags() {
    local flags=(
        '--with-branch:Also delete the branch (if merged)'
        '--force-branch:Force delete the branch (even if not merged)'
        '--kill-port:Kill processes using the port (opt-in)'
        '--kill-sidekiq:Kill Sidekiq processes (opt-in)'
        '--kill-all:Kill both port and Sidekiq processes'
        '-y:Skip confirmation prompts'
        '--yes:Skip confirmation prompts'
    )
    _describe 'flag' flags
}

_pwt_backups() {
    local backups
    if [[ -d ~/.pwt/trash ]]; then
        # Get worktree names from JSON metadata (without timestamp suffix for easier typing)
        backups=(${(f)"$(find ~/.pwt/trash -maxdepth 1 -name '*.json' -exec jq -r '.worktree // empty' {} \; 2>/dev/null | sort -u)"})
        [[ ${#backups[@]} -eq 0 ]] && backups=(${(f)"$(find ~/.pwt/trash -maxdepth 1 -name '*.patch' -exec basename {} .patch \; 2>/dev/null | sed -E 's/_[0-9]{8}_[0-9]{6}$//' | sort -u)"})
        _describe 'backup' backups
    fi
}

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

    _arguments -C \
        '--project[Specify project]:project:_pwt_projects' \
        '--bg[Run in background]' \
        '--no-input[Close stdin, set PWT_AGENT=1]' \
        '--quiet[Suppress non-essential output]' \
        '--verbose[Show detailed output]' \
        '1: :->command' \
        '*: :->args'

    case $state in
        command)
            # Check if first arg is a project name (exact or partial)
            if _pwt_is_project "$words[2]"; then
                local _proj="$(_pwt_resolve_project "$words[2]")"
                # After project: offer commands AND worktrees (for exec syntax)
                _alternative \
                    'commands:command:_pwt_commands' \
                    "worktrees:worktree:_pwt_worktrees $_proj"
            else
                _alternative \
                    'projects:project:_pwt_projects' \
                    'commands:command:_pwt_commands'
            fi
            ;;
        args)
            local cmd="${words[2]}"

            # Handle project as first arg (exact or partial match)
            if _pwt_is_project "$cmd"; then
                local _project="$(_pwt_resolve_project "$cmd")"
                cmd="${words[3]}"
                case $cmd in
                    create)
                        case $CURRENT in
                            4) _pwt_branches ;;
                            5) _pwt_branches ;;
                            *) _message 'description' ;;
                        esac
                        ;;
                    track)
                        case $CURRENT in
                            4) _pwt_branches ;;
                            *) _arguments '--name[Override worktree name]:name:' '-e[Open editor]' '--editor[Open editor]' '-a[Start AI]' '--ai[Start AI]' '-n[Dry run]' '--dry-run[Dry run]' '--clone[Use clone mode]' ;;
                        esac
                        ;;
                    adopt|setup)
                        _files -/
                        ;;
                    remove|rm)
                        _alternative \
                            "worktrees:worktree:_pwt_worktrees $_project" \
                            'flags:flag:_pwt_remove_flags'
                        ;;
                    cd|use|server|s|info|show|fix-port|port|editor|e|ai|open|current|repair|fix|step)
                        _pwt_worktrees "$_project"
                        ;;
                    gateway)
                        case $CURRENT in
                            4) _pwt_gateway_actions ;;
                            5)
                                case ${words[4]} in
                                    use) _pwt_worktrees "$_project" ;;
                                    init) _arguments '--port[Gateway port]:port:' '--host[Gateway public host]:host:' '-H[Gateway public host]:host:' ;;
                                esac
                                ;;
                            *) _arguments '--port[Gateway port]:port:' '--host[Gateway public host]:host:' '-H[Gateway public host]:host:' '--json[JSON output]' '--all[Show all]' '-a[Show all]' ;;
                        esac
                        ;;
                    run|for-each)
                        case $CURRENT in
                            4) _pwt_worktrees "$_project" ;;
                            *) _message 'command' ;;
                        esac
                        ;;
                    diff)
                        _pwt_worktrees "$_project"
                        ;;
                    copy|cp)
                        case $CURRENT in
                            4|5) _pwt_worktrees "$_project" ;;
                            *) _message 'file patterns' ;;
                        esac
                        ;;
                    servers)
                        _arguments '--all[Show stopped worktrees]' '-a[Show stopped worktrees]' '--json[JSON output]'
                        ;;
                    tree|status|doctor|list|ls)
                        # These commands have optional flags, no required args
                        ;;
                    meta|m)
                        case $CURRENT in
                            4) _pwt_meta_actions ;;
                            5)
                                case ${words[4]} in
                                    show|set) _pwt_worktrees "$_project" ;;
                                esac
                                ;;
                            6)
                                case ${words[4]} in
                                    set) _pwt_meta_fields ;;
                                esac
                                ;;
                        esac
                        ;;
                    project)
                        case $CURRENT in
                            4) _pwt_project_actions ;;
                            5)
                                case ${words[4]} in
                                    show|set|path) _pwt_projects ;;
                                    init) _message 'project name' ;;
                                esac
                                ;;
                            6)
                                case ${words[4]} in
                                    set) _pwt_config_keys ;;
                                esac
                                ;;
                        esac
                        ;;
                    config)
                        case $CURRENT in
                            4) _pwt_config_keys ;;
                            *) _message 'value' ;;
                        esac
                        ;;
                    restore)
                        case $CURRENT in
                            4) _pwt_backups ;;
                            5) _pwt_worktrees "$_project" ;;
                        esac
                        ;;
                    jobs)
                        local jobs_actions=(
                            'list:List all jobs'
                            'logs:View job output'
                            'stop:Stop a running job'
                            'clean:Remove stale job entries'
                            'help:Show help'
                        )
                        _describe 'action' jobs_actions
                        ;;
                esac
            else
                case $cmd in
                    create)
                        case $CURRENT in
                            3) _pwt_branches ;;
                            4) _pwt_branches ;;
                            *) _message 'description' ;;
                        esac
                        ;;
                    track)
                        case $CURRENT in
                            3) _pwt_branches ;;
                            *) _arguments '--name[Override worktree name]:name:' '-e[Open editor]' '--editor[Open editor]' '-a[Start AI]' '--ai[Start AI]' '-n[Dry run]' '--dry-run[Dry run]' '--clone[Use clone mode]' ;;
                        esac
                        ;;
                    adopt|setup)
                        _files -/
                        ;;
                    remove|rm)
                        _alternative \
                            'worktrees:worktree:_pwt_worktrees' \
                            'flags:flag:_pwt_remove_flags'
                        ;;
                    cd|use|server|s|info|show|fix-port|repair|fix|port|editor|e|ai|open|current)
                        _pwt_worktrees
                        ;;
                    gateway)
                        case $CURRENT in
                            3) _pwt_gateway_actions ;;
                            4)
                                case ${words[3]} in
                                    use) _pwt_worktrees ;;
                                    init) _arguments '--port[Gateway port]:port:' '--host[Gateway public host]:host:' '-H[Gateway public host]:host:' ;;
                                esac
                                ;;
                            *) _arguments '--port[Gateway port]:port:' '--host[Gateway public host]:host:' '-H[Gateway public host]:host:' '--json[JSON output]' ;;
                        esac
                        ;;
                    run|for-each)
                        case $CURRENT in
                            3) _pwt_worktrees ;;
                            *) _message 'command' ;;
                        esac
                        ;;
                    diff)
                        _pwt_worktrees
                        ;;
                    copy|cp)
                        case $CURRENT in
                            3|4) _pwt_worktrees ;;
                            *) _message 'file patterns' ;;
                        esac
                        ;;
                    servers)
                        _arguments '--all[Show stopped worktrees]' '-a[Show stopped worktrees]' '--json[JSON output]'
                        ;;
                    tree|status|doctor)
                        # These commands have optional flags, no required args
                        ;;
                    plugin)
                        local plugin_actions=(
                            'list:List installed plugins'
                            'install:Install a plugin'
                            'remove:Remove a plugin'
                            'create:Create a new plugin'
                            'path:Show plugins directory'
                            'help:Show plugin help'
                        )
                        _describe 'action' plugin_actions
                        ;;
                    claude-setup)
                        local claude_actions=(
                            'install:Install status line script'
                            'vars:Show available variables'
                            'format:Show or set format'
                            'preview:Preview status line'
                            'test:Run tests'
                            'toggle:Toggle on/off'
                            'help:Show help'
                        )
                        _describe 'action' claude_actions
                        ;;
                    jobs)
                        local jobs_actions=(
                            'list:List all jobs'
                            'logs:View job output'
                            'stop:Stop a running job'
                            'clean:Remove stale job entries'
                            'help:Show help'
                        )
                        _describe 'action' jobs_actions
                        ;;
                    help|-h|--help)
                        local help_topics=(
                            'concepts:What is pwt, worktree vs clone'
                            'commands:Full command reference'
                            'navigation:cd, use, multi-project'
                            'pwtfile:Syntax, hooks, helpers, examples'
                            'all:Everything (good for LLMs)'
                        )
                        _describe 'topic' help_topics
                        ;;
                    auto-remove|cleanup)
                        _pwt_branches
                        ;;
                    restore)
                        case $CURRENT in
                            3) _pwt_backups ;;
                            4) _pwt_worktrees ;;
                        esac
                        ;;
                    meta|m)
                        case $CURRENT in
                            3) _pwt_meta_actions ;;
                            4)
                                case ${words[3]} in
                                    show|set) _pwt_worktrees ;;
                                esac
                                ;;
                            5)
                                case ${words[3]} in
                                    set) _pwt_meta_fields ;;
                                esac
                                ;;
                        esac
                        ;;
                    project)
                        case $CURRENT in
                            3) _pwt_project_actions ;;
                            4)
                                case ${words[3]} in
                                    show|set|path) _pwt_projects ;;
                                    init) _message 'project name' ;;
                                esac
                                ;;
                            5)
                                case ${words[3]} in
                                    set) _pwt_config_keys ;;
                                esac
                                ;;
                        esac
                        ;;
                    config)
                        case $CURRENT in
                            3) _pwt_config_keys ;;
                            *) _message 'value' ;;
                        esac
                        ;;
                esac
            fi
            ;;
    esac
}

_pwt "$@"
