#!/bin/bash

# fixed by Dominik Weglarz (weglarz.it) 
# Configuration paths and environment settings
VENV_PATH="/opt/sgpt_env"
WRAPPER_PATH="/usr/local/bin/sgpt"
CONFIG_DIR="$HOME/.config/shell_gpt"
CONFIG_FILE="$CONFIG_DIR/.sgptrc"

ACTIVATION_KEY=""
DEFAULT_MODEL="azure/gpt-4o"
USE_LITELLM="true"
API_BASE_URL="https://apiservice5lh2czy2mmcta.azure-api.net/ecc-sqmp-v1/"
AZURE_API_VERSION="2024-09-01-preview"

check_prerequisites() {
    if ! python3 -m venv --help >/dev/null 2>&1; then
        printf "Installing required python3-venv package...\n"
        apt-get update -qq && apt-get install -y python3-venv || {
            printf "Error: Failed to install python3-venv.\n" >&2
            return 1
        }
    fi
}

install_dependencies() {
    printf "Setting up isolated virtual environment in %s...\n" "$VENV_PATH"
    rm -rf "$VENV_PATH"
    python3 -m venv "$VENV_PATH" || { printf "Error: Failed to create virtual environment.\n" >&2; return 1; }

    printf "Installing required packages inside virtual environment...\n"
    "$VENV_PATH/bin/pip" install --upgrade pip >/dev/null 2>&1
    "$VENV_PATH/bin/pip" install shell-gpt litellm || { printf "Error: Failed to install shell-gpt/litellm.\n" >&2; return 1; }
    "$VENV_PATH/bin/pip" install "click==8.1.7" || { printf "Error: Failed to pin click version.\n" >&2; return 1; }

    printf "Dependencies successfully installed without affecting system packages.\n"
}

prompt_for_key() {
    printf "Enter your AI Activation Key: "
    read -r ACTIVATION_KEY
    if [[ -z "$ACTIVATION_KEY" ]]; then
        printf "Error: Activation Key cannot be empty.\n" >&2
        return 1
    fi
}

configure_sgpt() {
    mkdir -p "$CONFIG_DIR" || { printf "Error: Failed to create config directory.\n" >&2; return 1; }

    cat <<EOF > "$CONFIG_FILE"
DEFAULT_MODEL=$DEFAULT_MODEL
USE_LITELLM=$USE_LITELLM
OPENAI_API_KEY=$ACTIVATION_KEY
API_BASE_URL=$API_BASE_URL
EOF

    printf "ShellGPT configuration saved to %s\n" "$CONFIG_FILE"
}

create_binary_wrapper() {
    printf "Creating global command wrapper at %s...\n" "$WRAPPER_PATH"
    cat <<EOF > "$WRAPPER_PATH"
#!/bin/bash
export AZURE_API_BASE="$API_BASE_URL"
export AZURE_API_VERSION="$AZURE_API_VERSION"
exec $VENV_PATH/bin/sgpt "\$@"
EOF
    chmod +x "$WRAPPER_PATH"
}

execute_sgpt() {
    printf "Verifying sgpt installation...\n"
    sgpt "Hi" || { printf "Error: sgpt verification failed.\n" >&2; return 1; }
    printf "\nSetup completed successfully. You can now use the 'sgpt' command.\n"
}

main() {
    check_prerequisites || exit 1
    install_dependencies || exit 1
    prompt_for_key || exit 1
    configure_sgpt || exit 1
    create_binary_wrapper
    execute_sgpt
}

main