#!/bin/bash
# SPDX-License-Identifier: MIT

_10mtctl_completion() {
    local cur prev opts
    COMPREPLY=()
    cur="${COMP_WORDS[COMP_CWORD]}"
    prev="${COMP_WORDS[COMP_CWORD-1]}"

    # Available commands
    opts="list info name cleanup nuke compose console provision start stop snapshot restore-snapshot"

    # Commands that require VM IDs (everything except list)
    vm_commands="info name cleanup nuke compose console provision start stop snapshot restore-snapshot"

    # Function to get VM IDs from current user's workspace
    _get_vm_ids() {
        local workdir="/var/tmp/10mt/users/${USER}/workspace"

        if [ -d "${workdir}" ] && [ -r "${workdir}" ]; then
            find "${workdir}" -maxdepth 1 -type d -printf '%f\n' 2>/dev/null | \
                grep -v "^workspace$" | \
                grep -v "^\.$" | \
                sort
        fi
    }

    # Function to check if a VM is running
    _is_vm_running() {
        local vm_path="/var/tmp/10mt/users/${USER}/workspace/$1"

        [ ! -e "${vm_path}/id" ] && return 1

        local vm_id
        vm_id="$(cat "${vm_path}/id" 2>/dev/null)" || return 1

        local running_vms
        running_vms=$(virsh --connect qemu:///system list --state-running --name 2>/dev/null | tr '\n' ' ')

        printf '%s' " ${running_vms} " | grep -q " ${vm_id} "
    }

    # Function to get running VMs only
    _get_running_vms() {
        local vm_ids
        vm_ids=$(_get_vm_ids)

        for vm in $vm_ids; do
            if _is_vm_running "$vm"; then
                echo "$vm"
            fi
        done
    }

    # Function to get stopped VMs only
    _get_stopped_vms() {
        local vm_ids
        vm_ids=$(_get_vm_ids)

        for vm in $vm_ids; do
            if ! _is_vm_running "$vm"; then
                echo "$vm"
            fi
        done
    }

    # Handle sudo prefix - adjust word positions
    local cmd_offset=0
    if [[ "${COMP_WORDS[0]}" == "sudo" ]]; then
        cmd_offset=1
    fi

    local actual_cmd_pos=$((1 + cmd_offset))
    local actual_arg_pos=$((2 + cmd_offset))

    # If we're completing the command (10mtctl command)
    if [ ${COMP_CWORD} -eq ${actual_cmd_pos} ]; then
        COMPREPLY=($(compgen -W "${opts}" -- ${cur}))
        return 0
    fi

    # If we're completing the VM ID argument
    if [ ${COMP_CWORD} -eq ${actual_arg_pos} ]; then
        local command="${COMP_WORDS[${actual_cmd_pos}]}"
        if [[ " ${vm_commands} " =~ " ${command} " ]]; then
            local vm_ids
            case "${command}" in
                console|stop)
                    # Commands that require running VMs
                    vm_ids=$(_get_running_vms)
                    ;;
                start)
                    # Commands that require stopped VMs
                    vm_ids=$(_get_stopped_vms)
                    ;;
                *)
                    # Commands that work with any VM state
                    vm_ids=$(_get_vm_ids)
                    ;;
            esac
            COMPREPLY=($(compgen -W "${vm_ids}" -- ${cur}))
            return 0
        fi
    fi

    return 0
}

# Register the completion function
complete -F _10mtctl_completion 10mtctl
