#!/bin/bash

MODEL="gpt-4o"      # gpt-4o, gpt-4, gpt-3.5-turbo, gpt-3.5, gpt-3, curie, babbage, davinci
TEMPERATURE=0.2     # 0.0 to 1.0 (higher is more creative)
MAX_TOKENS=3072     # 1 to 4096 (higher requires more API credits, but can be more accurate)

# The base prompt describing the rules the AI should follow
read -r -d '' BASE_PROMPT << EOM
Do not explain anything.
Only respond with a command that can be executed on the linux terminal.
Never use code tags.
Always put everything in a single command, e.g. using pipes, semicolons or double ampersands.
Break up the command in multiple lines when using redirections, semicolons or double ampersands, but make sure the command can still be executed (e.g. by using a backslash).
Do not use any formatting except bash escaped newlines.
The question is:
EOM

# INTERNAL VARIABLES - DO NOT CHANGE
ALL_YES=false       # If true, the AI will execute the command without asking for confirmation

# Check if the OPENAI_API_KEY environment variable is set
# Ignore if -a is passed as a parameter
if [[ "$1" != "-a" && -z "$OPENAI_API_KEY" ]]; then
    OPENAI_API_KEY=$(printenv OPENAI_API_KEY)
    
    if [ -z "$OPENAI_API_KEY" ]; then
        
        # Check if the OPENAI_API_KEY is set in the ~/.openai file
        if [ -f ~/.openai ]; then
            source ~/.openai
        fi
        
        if [ -z "$OPENAI_API_KEY" ]; then
            echo "Please set the OPENAI_API_KEY environment variable."
            exit 1
        fi
    fi
fi

# check if -y is passed as a parameter and set ALL_YES to true
if [ "$1" == "-y" ]; then
    ALL_YES=true
    shift
fi

function run_ai()
{
    if [ -z "$*" ]; then
        echo "Usage: runai <command>"
        return
    fi
    
    # if ALL_YES is set to true, print an info message in red bold font
    if [ $ALL_YES == true ]; then
        echo -e "\033[1;31mRunning the command without asking for confirmation!\033[0m"
    fi
    
    echo -e -n "\e[38;5;12m"
    
    command=$(OPENAI_API_KEY=$OPENAI_API_KEY openai api chat.completions.create -m gpt-4o -t 0.4 -M 3072 -g "user" "$BASE_PROMPT $*")
    
    echo -e -n "\e[0m"
    
    if [ $ALL_YES == false ]; then
        echo -e "The generated command is:\n\033[0;33m$command\n\033[1;31m"
        read -p "Do you want to execute the command? [Y/n]: " input
        input=$(echo "$input" | tr '"'"'[A-Z]'"'"' '"'"'[a-z]'"'"')
    else
        echo -e "Running the command:\n\033[0;33m$command\n\033[1;31m"
    fi
    
    if [[ $ALL_YES == true || $input == "y" || $input == "yes" ]]; then
        echo -e "\033[0m";
        history -s "runai $*"
        history -s "$command"
        eval "$command"
    fi
}

# if the parameter -a is passed, set the alias, otherwise run the AI

if [ "$1" == "-a" ]; then
    # check first, if an alias with the same name already exists
    if grep -q "alias runai=" ~/.bashrc; then
        echo "The alias runai already exists."
        exit 0
    else
        echo "alias runai='source $0'" >> ~/.bashrc
        echo "The alias runai has been set. You can now run the AI by typing runai <command>."
    fi
else
    run_ai "$*"
fi