yt-audio() {
  nix-shell -p yt-dlp --run "yt-dlp -x $1 --audio-format mp3"
}

docker_mysql() {
  if [ -z "$1" ]; then
    echo "Usage: docker_mysql <mysql_container_name>"
    return 1
  fi
  docker exec -it "$1" mysql -u root -ppassword
}
# Tab-completion for docker_mysql function
_complete_docker_mysql() {
  local cur=${COMP_WORDS[COMP_CWORD]}
  COMPREPLY=($(compgen -W "$(docker ps --format '{{.Names}}' | grep mysql)" -- $cur))
}
# Register the completion for the docker_mysql function
complete -F _complete_docker_mysql docker_mysql

# Git restore file, like git restore but works for staged changes as well
gres() {
  git restore --source=HEAD --staged --worktree -- "$1"
}

# Nixos upgrade functions
upgrade() {
    local command="$1"
    local profile="$2"
    local flake_path="$HOME/nixos"
    
    # Define valid commands
    local valid_commands=("switch" "boot" "test" "build" "dry-build" "dry-activate" "edit" "repl" "build-vm" "build-vm-with-bootloader" "build-image")
    
    # Check if command is provided
    if [[ -z "$command" ]]; then
        echo "Error: No command specified"
        echo "Usage: upgrade <command> [profile]"
        echo "Valid commands: ${valid_commands[*]}"
        return 1
    fi

    # Validate command
    local valid_command=false
    for valid_cmd in "${valid_commands[@]}"; do
        if [[ "$command" == "$valid_cmd" ]]; then
            valid_command=true
            break
        fi
    done
    
    if [[ "$valid_command" != true ]]; then
        echo "Error: Invalid command '$command'"
        echo "Valid commands: ${valid_commands[*]}"
        return 1
    fi
    
    # Check if flake directory exists
    if [[ ! -d "$flake_path" ]]; then
        echo -e "Error: Flake directory '$flake_path' does not exist, check function definition:\n\n\t`function upgrade`"
        return 1
    fi
    
    # Check if flake.nix exists
    if [[ ! -f "$flake_path/flake.nix" ]]; then
        echo "Error: flake.nix not found in '$flake_path'"
        return 1
    fi
    
    # If profile is provided, validate it exists in flake.nix
    if [[ -n "$profile" ]]; then
        if ! awk '/nixosConfigurations = \{/,/^\s*\}/' "$flake_path/flake.nix" | \
             grep -E "^\s*$profile\s*=" > /dev/null; then
            echo "Error: Profile '$profile' not found in flake.nix"
            echo "Available profiles:"
            # Extract available profiles from nixosConfigurations section
            awk '/nixosConfigurations = \{/,/^\s*\}/' "$flake_path/flake.nix" | \
                grep -E "^\s*[a-zA-Z0-9_-]+\s*=.*nixpkgs\.lib\.nixosSystem" | \
                sed -E 's/^\s*([a-zA-Z0-9_-]+)\s*=.*/  - \1/'
            return 1
        fi
        # Build the full flake reference
        local flake_ref="$flake_path#$profile"
        echo "Running: sudo nixos-rebuild $command --flake $flake_ref"
        sudo nixos-rebuild "$command" --flake "$flake_ref"
    else
        echo "Error: No profile provided."
        echo "Usage: upgrade <command> [profile]"
        return 1
    fi
}

# Add tab completion for the function
_upgrade_completion() {
    local cur="${COMP_WORDS[COMP_CWORD]}"
    local prev="${COMP_WORDS[COMP_CWORD-1]}"
    
    case $COMP_CWORD in
        1)
            # Complete command names
            local commands="switch boot test build dry-build dry-activate edit repl build-vm build-vm-with-bootloader build-image"
            COMPREPLY=($(compgen -W "$commands" -- "$cur"))
            ;;
        2)
            # Complete profile names from flake.nix
            if [[ -f "$HOME/nixos/flake.nix" ]]; then
                local profiles=$(grep -E "^\s*[a-zA-Z0-9_-]+\s*=" "$HOME/nixos/flake.nix" | \
                    grep -A5 -B5 "nixosConfigurations" | \
                    sed -n 's/^\s*\([a-zA-Z0-9_-]*\)\s*=.*/\1/p' | \
                    sort -u)
                COMPREPLY=($(compgen -W "$profiles" -- "$cur"))
            fi
            ;;
    esac
}

# Register the completion function
complete -F _upgrade_completion upgrade