#compdef wallet

# Zsh completion script for the wallet CLI
# See instructions in ../README.md

_wallet() {
    local -a commands
    local -a subcommands
    local curcontext="$curcontext" state line
    typeset -A opt_args

    _arguments -C \
        '(-c --continuous-run)'{-c,--continuous-run}'[Continuous run flag]' \
        '--auth[Basic authentication in the format user or user\:password]:auth:' \
        '1: :->command' \
        '*:: :->args'

    case $state in
        command)
            commands=(
                'auth-transfer:Authenticated transfer subcommand'
                'chain-info:Generic chain info subcommand'
                'account:Account view and sync subcommand'
                'pinata:Pinata program interaction subcommand'
                'token:Token program interaction subcommand'
                'amm:AMM program interaction subcommand'
                'check-health:Check the wallet can connect to the node and builtin local programs match the remote versions'
                'config:Command to setup config, get and set config fields'
                'restore-keys:Restoring keys from given password at given depth'
                'deploy-program:Deploy a program'
                'help:Print help message or the help of the given subcommand(s)'
            )
            _describe -t commands 'wallet commands' commands
            ;;
        args)
            case $line[1] in
                auth-transfer)
                    _wallet_auth_transfer
                    ;;
                chain-info)
                    _wallet_chain_info
                    ;;
                account)
                    _wallet_account
                    ;;
                pinata)
                    _wallet_pinata
                    ;;
                token)
                    _wallet_token
                    ;;
                amm)
                    _wallet_amm
                    ;;
                config)
                    _wallet_config
                    ;;
                restore-keys)
                    _wallet_restore_keys
                    ;;
                deploy-program)
                    _wallet_deploy_program
                    ;;
                help)
                    _wallet_help
                    ;;
            esac
            ;;
    esac
}

# auth-transfer subcommand
_wallet_auth_transfer() {
    local -a subcommands
    
    _arguments -C \
        '1: :->subcommand' \
        '*:: :->args'

    case $state in
        subcommand)
            subcommands=(
                'init:Initialize account under authenticated transfer program'
                'send:Send native tokens from one account to another with variable privacy'
                'help:Print this message or the help of the given subcommand(s)'
            )
            _describe -t subcommands 'auth-transfer subcommands' subcommands
            ;;
        args)
            case $line[1] in
                init)
                    _arguments \
                        '--account-id[Account ID to initialize]:account_id:_wallet_account_ids'
                    ;;
                send)
                    _arguments \
                        '--from[Source account ID]:from_account:_wallet_account_ids' \
                        '--to[Destination account ID (for owned accounts)]:to_account:_wallet_account_ids' \
                        '--to-npk[Destination nullifier public key (for foreign private accounts)]:npk:' \
                        '--to-ipk[Destination viewing public key (for foreign private accounts)]:ipk:' \
                        '--amount[Amount of native tokens to send]:amount:'
                    ;;
            esac
            ;;
    esac
}

# chain-info subcommand
_wallet_chain_info() {
    local -a subcommands
    
    _arguments -C \
        '1: :->subcommand' \
        '*:: :->args'

    case $state in
        subcommand)
            subcommands=(
                'current-block-id:Get current block id from sequencer'
                'block:Get block at id from sequencer'
                'transaction:Get transaction at hash from sequencer'
                'help:Print this message or the help of the given subcommand(s)'
            )
            _describe -t subcommands 'chain-info subcommands' subcommands
            ;;
        args)
            case $line[1] in
                block)
                    _arguments \
                        '--id[Block ID to retrieve]:block_id:'
                    ;;
                transaction)
                    _arguments \
                        '--hash[Transaction hash to retrieve]:tx_hash:'
                    ;;
            esac
            ;;
    esac
}

# account subcommand
_wallet_account() {
    local -a subcommands
    
    _arguments -C \
        '1: :->subcommand' \
        '*:: :->args'

    case $state in
        subcommand)
            subcommands=(
                'get:Get account data'
                'list:List all accounts'
                'ls:List all accounts (alias for list)'
                'new:Produce new public or private account'
                'sync-private:Sync private accounts'
                'help:Print this message or the help of the given subcommand(s)'
            )
            _describe -t subcommands 'account subcommands' subcommands
            ;;
        args)
            case $line[1] in
                get)
                    _arguments \
                        '(-r --raw)'{-r,--raw}'[Get raw account data]' \
                        '(-k --keys)'{-k,--keys}'[Display keys (pk for public accounts, npk/ipk for private accounts)]' \
                        '(-a --account-id)'{-a,--account-id}'[Account ID to query]:account_id:_wallet_account_ids'
                    ;;
                list|ls)
                    _arguments \
                        '(-l --long)'{-l,--long}'[Display detailed account information]'
                    ;;
                new)
                    _arguments -C \
                        '1: :->account_type' \
                        '*:: :->new_args'
                    case $state in
                        account_type)
                            compadd public private
                            ;;
                        new_args)
                            _arguments \
                                '--cci[Chain index of a parent node]:chain_index:'
                            ;;
                    esac
                    ;;
            esac
            ;;
    esac
}

# pinata subcommand
_wallet_pinata() {
    local -a subcommands
    
    _arguments -C \
        '1: :->subcommand' \
        '*:: :->args'

    case $state in
        subcommand)
            subcommands=(
                'claim:Claim tokens from the Piñata faucet'
                'help:Print this message or the help of the given subcommand(s)'
            )
            _describe -t subcommands 'pinata subcommands' subcommands
            ;;
        args)
            case $line[1] in
                claim)
                    _arguments \
                        '--to[Destination account ID to receive claimed tokens]:to_account:_wallet_account_ids'
                    ;;
            esac
            ;;
    esac
}

# token subcommand
_wallet_token() {
    local -a subcommands

    _arguments -C \
        '1: :->subcommand' \
        '*:: :->args'

    case $state in
        subcommand)
            subcommands=(
                'new:Produce a new token'
                'send:Send tokens from one account to another with variable privacy'
                'burn:Burn tokens on holder, modify definition'
                'mint:Mint tokens on holder, modify definition'
                'help:Print this message or the help of the given subcommand(s)'
            )
            _describe -t subcommands 'token subcommands' subcommands
            ;;
        args)
            case $line[1] in
                new)
                    _arguments \
                        '--name[Token name]:name:' \
                        '--total-supply[Total supply of tokens to mint]:total_supply:' \
                        '--definition-account-id[Account ID for token definition]:definition_account:_wallet_account_ids' \
                        '--supply-account-id[Account ID to receive initial supply]:supply_account:_wallet_account_ids'
                    ;;
                send)
                    _arguments \
                        '--from[Source holding account ID]:from_account:_wallet_account_ids' \
                        '--to[Destination holding account ID (for owned accounts)]:to_account:_wallet_account_ids' \
                        '--to-npk[Destination nullifier public key (for foreign private accounts)]:npk:' \
                        '--to-ipk[Destination viewing public key (for foreign private accounts)]:ipk:' \
                        '--amount[Amount of tokens to send]:amount:'
                    ;;
                burn)
                    _arguments \
                        '--definition[Definition account ID]:definition_account:_wallet_account_ids' \
                        '--holder[Holder account ID]:holder_account:_wallet_account_ids' \
                        '--amount[Amount of tokens to burn]:amount:'
                    ;;
                mint)
                    _arguments \
                        '--definition[Definition account ID]:definition_account:_wallet_account_ids' \
                        '--holder[Holder account ID (for owned accounts)]:holder_account:_wallet_account_ids' \
                        '--holder-npk[Holder nullifier public key (for foreign private accounts)]:npk:' \
                        '--holder-ipk[Holder viewing public key (for foreign private accounts)]:ipk:' \
                        '--amount[Amount of tokens to mint]:amount:'
                    ;;
            esac
            ;;
    esac
}

# amm subcommand
_wallet_amm() {
    local -a subcommands
    
    _arguments -C \
        '1: :->subcommand' \
        '*:: :->args'

    case $state in
        subcommand)
            subcommands=(
                'new:Create a new liquidity pool'
                'swap:Swap tokens using the AMM'
                'add-liquidity:Add liquidity to an existing pool'
                'remove-liquidity:Remove liquidity from a pool'
                'help:Print this message or the help of the given subcommand(s)'
            )
            _describe -t subcommands 'amm subcommands' subcommands
            ;;
        args)
            case $line[1] in
                new)
                    _arguments \
                        '--user-holding-a[User token A holding account ID]:holding_a:_wallet_account_ids' \
                        '--user-holding-b[User token B holding account ID]:holding_b:_wallet_account_ids' \
                        '--user-holding-lp[User LP token holding account ID]:holding_lp:_wallet_account_ids' \
                        '--balance-a[Amount of token A to deposit]:balance_a:' \
                        '--balance-b[Amount of token B to deposit]:balance_b:'
                    ;;
                swap)
                    _arguments \
                        '--user-holding-a[User token A holding account ID]:holding_a:_wallet_account_ids' \
                        '--user-holding-b[User token B holding account ID]:holding_b:_wallet_account_ids' \
                        '--amount-in[Amount of tokens to swap]:amount_in:' \
                        '--min-amount-out[Minimum tokens expected in return]:min_amount_out:' \
                        '--token-definition[Definition ID of the token being provided]:token_def:'
                    ;;
                add-liquidity)
                    _arguments \
                        '--user-holding-a[User token A holding account ID]:holding_a:_wallet_account_ids' \
                        '--user-holding-b[User token B holding account ID]:holding_b:_wallet_account_ids' \
                        '--user-holding-lp[User LP token holding account ID]:holding_lp:_wallet_account_ids' \
                        '--max-amount-a[Maximum amount of token A to deposit]:max_amount_a:' \
                        '--max-amount-b[Maximum amount of token B to deposit]:max_amount_b:' \
                        '--min-amount-lp[Minimum LP tokens to receive]:min_amount_lp:'
                    ;;
                remove-liquidity)
                    _arguments \
                        '--user-holding-a[User token A holding account ID]:holding_a:_wallet_account_ids' \
                        '--user-holding-b[User token B holding account ID]:holding_b:_wallet_account_ids' \
                        '--user-holding-lp[User LP token holding account ID]:holding_lp:_wallet_account_ids' \
                        '--balance-lp[Amount of LP tokens to burn]:balance_lp:' \
                        '--min-amount-a[Minimum token A to receive]:min_amount_a:' \
                        '--min-amount-b[Minimum token B to receive]:min_amount_b:'
                    ;;
            esac
            ;;
    esac
}

# config subcommand
_wallet_config() {
    local -a subcommands
    local -a config_keys

    config_keys=(
        'all'
        'override_rust_log'
        'sequencer_addr'
        'seq_poll_timeout_millis'
        'seq_tx_poll_max_blocks'
        'seq_poll_max_retries'
        'seq_block_poll_max_amount'
        'initial_accounts'
        'basic_auth'
    )

    _arguments -C \
        '1: :->subcommand' \
        '*:: :->args'

    case $state in
        subcommand)
            subcommands=(
                'get:Getter of config fields'
                'set:Setter of config fields'
                'description:Prints description of corresponding field'
                'help:Print this message or the help of the given subcommand(s)'
            )
            _describe -t subcommands 'config subcommands' subcommands
            ;;
        args)
            case $line[1] in
                get|description)
                    compadd -a config_keys
                    ;;
                set)
                    _arguments \
                        '1:key:compadd -a config_keys' \
                        '2:value:'
                    ;;
            esac
            ;;
    esac
}

# restore-keys subcommand
_wallet_restore_keys() {
    _arguments \
        '(-d --depth)'{-d,--depth}'[How deep in tree accounts may be]:depth:'
}

# deploy-program subcommand
_wallet_deploy_program() {
    _arguments \
        '1:binary filepath:_files'
}

# help subcommand
_wallet_help() {
    local -a commands
    commands=(
        'auth-transfer:Authenticated transfer subcommand'
        'chain-info:Generic chain info subcommand'
        'account:Account view and sync subcommand'
        'pinata:Pinata program interaction subcommand'
        'token:Token program interaction subcommand'
        'amm:AMM program interaction subcommand'
        'check-health:Check the wallet can connect to the node'
        'config:Command to setup config, get and set config fields'
        'restore-keys:Restoring keys from given password at given depth'
        'deploy-program:Deploy a program'
    )
    _describe -t commands 'wallet commands' commands
}

# Helper function to complete account IDs
# Uses `wallet account list` to get available accounts
# Only includes accounts with /N prefix (where N is a number)
_wallet_account_ids() {
    local -a accounts
    local line
    
    # Try to get accounts from wallet account list command
    # Filter to lines starting with /N (numbered accounts) and extract the account ID
    if command -v wallet &>/dev/null; then
        while IFS= read -r line; do
            # Remove trailing comma if present and add to array
            [[ -n "$line" ]] && accounts+=("${line%,}")
        done < <(wallet account list 2>/dev/null | grep '^/[0-9]' | awk '{print $2}')
    fi
    
    # Provide type prefixes as fallback if command fails or returns nothing
    if (( ${#accounts} == 0 )); then
        compadd -S '' -- 'Public/' 'Private/'
        return
    fi
    
    _multi_parts / accounts
}

_wallet "$@"