#compdef mpytool

# ZSH completion for mpytool
# Install: copy to ~/.zsh/completions/_mpytool and add to fpath

# Cache for remote paths (per port+dir)
typeset -gA _mpytool_cache
typeset -gA _mpytool_cache_time

_mpytool_commands() {
    # Get commands dynamically from mpytool
    local -a commands
    commands=(${(f)"$(mpytool _commands 2>/dev/null)"})
    _describe -t commands 'mpytool command' commands
}

_mpytool_options() {
    # Get options for a command dynamically from mpytool
    # $1 = command name
    local -a options
    options=(${(f)"$(mpytool _options "$1" 2>/dev/null)"})
    [[ ${#options[@]} -gt 0 ]] && _describe -t options "$1 option" options
}

_mpytool_detect_ports() {
    # Use mpytool ports for consistent detection across all platforms
    mpytool ports 2>/dev/null
}

_mpytool_ports() {
    # Complete serial ports with descriptions
    local -a ports descs
    local line
    while IFS=$'\t' read -r port desc; do
        [[ -z "$port" ]] && continue
        ports+=("$port")
        descs+=("$port -- $desc")
    done <<< "$(_mpytool_detect_ports)"
    [[ ${#ports[@]} -gt 0 ]] && compadd -l -d descs -a ports
}

_mpytool_get_port() {
    # Use global _mpytool_current_port if set (from main function)
    if [[ -n "$_mpytool_current_port" ]]; then
        echo "$_mpytool_current_port"
        return
    fi

    # Fallback: auto-detect only if single port available
    local -a ports
    local line
    while IFS=$'\t' read -r port desc; do
        [[ -n "$port" ]] && ports+=("$port")
    done <<< "$(_mpytool_detect_ports)"
    if [[ ${#ports[@]} -eq 1 ]]; then
        echo "${ports[1]}"
    fi
}

_mpytool_complete_remote() {
    # $1 = "dirs" for directories only (optional)
    local dirs_only="$1"

    # Use compset to handle : and path prefixes (like _files does)
    # This sets IPREFIX so compadd -S '/' works correctly
    compset -P ':'
    compset -P '*/'
    # IPREFIX now contains the directory part (e.g., ":", ":/lib/", ":lib/sub/")
    # PREFIX contains the partial filename being completed

    local dir_query="${IPREFIX}"

    # Get port
    local port=$(_mpytool_get_port)
    if [[ -z "$port" ]]; then
        compadd -x 'no device connected'
        return
    fi

    # Cache per (port, dir_query), 10 seconds
    local cache_key="${port//\//_}_${dir_query}"
    local now=$(date +%s)
    local cache_age=$((now - ${_mpytool_cache_time[$cache_key]:-0}))

    local -a entries=()
    if [[ $cache_age -lt 10 && -n "${_mpytool_cache[$cache_key]}" ]]; then
        entries=(${(f)_mpytool_cache[$cache_key]})
    else
        local output
        output=$(mpytool -p "$port" _paths "$dir_query" 2>/dev/null)
        if [[ $? -ne 0 ]]; then
            compadd -x 'cannot read device'
            return
        fi
        entries=("${(@f)output}")
        _mpytool_cache[$cache_key]="${(F)entries}"
        _mpytool_cache_time[$cache_key]=$now
    fi

    # Add completions - compset already handles the prefix via IPREFIX
    # -S '/' -q for directories: removable suffix (bold, Tab enters, Space removes)
    # -d for display: show names with / in menu to distinguish from files
    local -a dirs=() dirs_display=() files=()
    local e
    for e in "${entries[@]}"; do
        [[ -z "$e" ]] && continue
        if [[ "$e" == */ ]]; then
            dirs+=("${e%/}")
            dirs_display+=("$e")
        else
            files+=("$e")
        fi
    done

    [[ ${#dirs[@]} -gt 0 ]] && compadd -S '/' -q -d dirs_display -- "${dirs[@]}"
    [[ -z "$dirs_only" && ${#files[@]} -gt 0 ]] && compadd -- "${files[@]}"
}

_mpytool_clear_cache() {
    # Call this to clear the remote file cache
    _mpytool_cache=()
    _mpytool_cache_time=()
}

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

    _arguments -C \
        '-V[show version]' \
        '--version[show version]' \
        '-p[serial port]:port:_mpytool_ports' \
        '--port[serial port]:port:_mpytool_ports' \
        '-a[network address]:address:' \
        '--address[network address]:address:' \
        '-b[baud rate]:baud:(9600 19200 38400 57600 115200 230400 460800 921600)' \
        '--baud[baud rate]:baud:(9600 19200 38400 57600 115200 230400 460800 921600)' \
        '-d[debug level]' \
        '--debug[debug level]' \
        '-v[verbose output]' \
        '--verbose[verbose output]' \
        '-q[quiet mode]' \
        '--quiet[quiet mode]' \
        '-f[force copy]' \
        '--force[force copy]' \
        '-e[exclude directory]:dir:_files -/' \
        '--exclude-dir[exclude directory]:dir:_files -/' \
        '*::command:->commands'

    # Set port for helper functions to use
    _mpytool_current_port="${opt_args[-p]:-${opt_args[--port]:-}}"

    case "$state" in
        commands)
            # Find the last -- separator and work relative to it
            local cmd_start=1
            local i
            for ((i=1; i < CURRENT; i++)); do
                if [[ "${words[i]}" == "--" ]]; then
                    cmd_start=$((i + 1))
                fi
            done

            local cmd="${words[cmd_start]}"
            local pos=$((CURRENT - cmd_start + 1))

            # If at command position (right after -- or at start), complete commands
            if [[ $CURRENT -eq $cmd_start ]]; then
                _mpytool_commands
                return
            fi

            # If current word is --, offer commands for next position
            if [[ "${words[CURRENT]}" == "--" ]]; then
                return
            fi

            # Command-specific completions (pos is position within current command group)
            # pos=2 is first arg after command, pos=3 is second, etc.
            local nargs=$((pos - 2))  # number of already completed args
            case "$cmd" in
                ls|tree)
                    # Optional remote path, -- anytime
                    if [[ $pos -eq 2 ]]; then
                        if [[ "$words[CURRENT]" == :* ]]; then
                            _mpytool_complete_remote
                        else
                            compadd -S '' ':'
                        fi
                    fi
                    compadd -- '--'
                    ;;
                cat)
                    # 1+ remote files required, -- after at least one
                    if [[ "$words[CURRENT]" == :* ]]; then
                        _mpytool_complete_remote
                    else
                        compadd -S '' ':'
                    fi
                    [[ $nargs -ge 1 ]] && compadd -- '--'
                    ;;
                mkdir|rm)
                    # 1+ remote paths required, -- after at least one
                    if [[ "$words[CURRENT]" == :* ]]; then
                        _mpytool_complete_remote
                    else
                        compadd -S '' ':'
                    fi
                    [[ $nargs -ge 1 ]] && compadd -- '--'
                    ;;
                mv)
                    # 2+ remote paths required, -- after at least two
                    if [[ "$words[CURRENT]" == :* ]]; then
                        _mpytool_complete_remote
                    else
                        compadd -S '' ':'
                    fi
                    [[ $nargs -ge 2 ]] && compadd -- '--'
                    ;;
                cd)
                    # Exactly 1 remote dir, -- only after it
                    if [[ $pos -eq 2 ]]; then
                        if [[ "$words[CURRENT]" == :* ]]; then
                            _mpytool_complete_remote 'dirs'
                        else
                            compadd -S '' ':'
                        fi
                    else
                        compadd -- '--'
                    fi
                    ;;
                path)
                    # path [-f|-a|-d] [paths...]
                    _mpytool_options path
                    # Remote paths (: prefix)
                    if [[ "$words[CURRENT]" == :* ]]; then
                        _mpytool_complete_remote
                    else
                        compadd -S '' ':'
                    fi
                    # -- after at least one path
                    [[ $nargs -ge 1 ]] && compadd -- '--'
                    ;;
                cp)
                    # n local + 1 remote OR n remote + 1 local
                    # Count local and remote args (excluding current word)
                    local n_local=0 n_remote=0
                    for ((i=cmd_start+1; i < CURRENT; i++)); do
                        if [[ "${words[i]}" == :* ]]; then
                            ((n_remote++))
                        else
                            ((n_local++))
                        fi
                    done
                    if [[ $n_local -ge 1 && $n_remote -ge 1 ]]; then
                        # Complete, only offer --
                        compadd -- '--'
                    else
                        # Need more args, offer files and :
                        if [[ "$words[CURRENT]" == :* ]]; then
                            _mpytool_complete_remote
                        else
                            _files
                            compadd -S '' ':'
                        fi
                    fi
                    ;;
                exec)
                    # 1 code string, -- after it
                    [[ $nargs -ge 1 ]] && compadd -- '--'
                    ;;
                run)
                    # 1 local .py file, -- after it
                    if [[ $pos -eq 2 ]]; then
                        _files -g '*.py'
                    fi
                    [[ $nargs -ge 1 ]] && compadd -- '--'
                    ;;
                edit)
                    # --editor CMD + 1 remote file, -- after file
                    _mpytool_options edit
                    if [[ "$words[CURRENT]" == :* ]]; then
                        _mpytool_complete_remote
                    else
                        compadd -S '' ':'
                    fi
                    [[ $nargs -ge 1 ]] && compadd -- '--'
                    ;;
                rtc)
                    # rtc [-s|-l|-u] [datetime]
                    _mpytool_options rtc
                    compadd -- '--'
                    ;;
                pwd)
                    # No arguments, -- immediately
                    compadd -- '--'
                    ;;
                reset)
                    # Reset flags or --
                    _mpytool_options reset
                    compadd -- '--'
                    ;;
                stop|info)
                    # No arguments, -- immediately
                    compadd -- '--'
                    ;;
                repl|monitor)
                    # Interactive, no chaining
                    ;;
                flash)
                    # flash [read|write|erase|ota] [label] [file]
                    local flash_subcmd="${words[cmd_start+1]}"
                    if [[ $pos -eq 2 ]]; then
                        local -a subcmds=(
                            'read:read flash/partition to file'
                            'write:write file to flash/partition'
                            'erase:erase flash/partition'
                            'ota:OTA firmware update'
                        )
                        _describe -t subcmds 'flash subcommand' subcmds
                    elif [[ "$flash_subcmd" == "ota" ]]; then
                        # ota: 1 firmware file, -- after it
                        if [[ $pos -eq 3 ]]; then
                            _files -g "*.bin"
                        else
                            compadd -- '--'
                        fi
                    elif [[ $pos -eq 3 && "$flash_subcmd" == "erase" ]]; then
                        compadd -- '--full'
                        compadd -- '--'
                    elif [[ $pos -eq 3 ]]; then
                        # read/write: label (device-specific, no completion)
                        :
                    elif [[ $pos -eq 4 && "$flash_subcmd" != "erase" ]]; then
                        _files
                    elif [[ "$flash_subcmd" == "erase" ]]; then
                        compadd -- '--full'
                        compadd -- '--'
                    fi
                    # -- after complete subcommand
                    if [[ "$flash_subcmd" == "erase" && $pos -ge 3 ]]; then
                        :  # already handled above
                    elif [[ "$flash_subcmd" == "ota" && $pos -ge 4 ]]; then
                        compadd -- '--'
                    elif [[ "$flash_subcmd" != "erase" && "$flash_subcmd" != "ota" && $pos -ge 5 ]]; then
                        compadd -- '--'
                    fi
                    ;;
                sleep)
                    # 1 number, -- after it
                    [[ $nargs -ge 1 ]] && compadd -- '--'
                    ;;
                mount)
                    # -m and -w flags + 1 local dir + optional :mount_point, -- for chaining
                    _mpytool_options mount
                    # Then local dir and mount point
                    if [[ "$words[CURRENT]" == :* ]]; then
                        compadd -S '' ':'
                    else
                        _files -/
                    fi
                    compadd -- '--'
                    ;;
                ln)
                    # n local sources + 1 remote dest (: prefix, absolute path)
                    local n_local=0 n_remote=0
                    for ((i=cmd_start+1; i < CURRENT; i++)); do
                        if [[ "${words[i]}" == :* ]]; then
                            ((n_remote++))
                        else
                            ((n_local++))
                        fi
                    done
                    if [[ $n_local -ge 1 && $n_remote -ge 1 ]]; then
                        # Complete, only offer --
                        compadd -- '--'
                    else
                        if [[ "$words[CURRENT]" == :* ]]; then
                            _mpytool_complete_remote
                        else
                            _files
                            compadd -S '' ':'
                        fi
                    fi
                    ;;
                speedtest|ports)
                    # No arguments, -- immediately
                    compadd -- '--'
                    ;;
                --)
                    _mpytool_commands
                    ;;
                *)
                    _mpytool_commands
                    ;;
            esac
            ;;
    esac
}

# Only run if called as completion (not when sourced)
if [[ -n "$compstate" ]]; then
    _mpytool "$@"
fi
