config.py:
import os

OPENAI_KEY = os.environ["OPENAI_KEY"]

HEADERS = {
    'Content-Type': 'application/json',
    'Authorization': f'Bearer {OPENAI_KEY}',
}

openai_api.py:
import requests
from . import config
from rich.markdown import Markdown


def raise_exception_if_no_key():
    if config.OPENAI_KEY is None:
        raise ValueError(Markdown("You need to set your OPENAI_KEY to use this script. \
        You can set it temporarily by running this on your terminal: `export OPENAI_KEY=YOUR_KEY_HERE`"))


def make_request_to_openai_api(data):
    response = requests.post('https://api.openai.com/v1/chat/completions', headers=config.HEADERS, json=data)
    if response.status_code != 200:
        raise ValueError("Your request to OpenAI API failed: ", response.json()['error'])
    return response.json()['choices'][0]['message']['content'].strip()


def request_to_completions(prompt, model="gpt-3.5-turbo", max_tokens=500, temperature=0):
    data = {
        "model": model,
        "messages": [
            {"role": "system", "content": "You are a shell command expert..."},
            {"role": "user", "content": prompt}
        ],
        "max_tokens": max_tokens,
        "temperature": temperature
    }
    return make_request_to_openai_api(data)

shell_utils.py:
import json
import os
import subprocess
from rich.console import Console
from rich.syntax import Syntax
from rich.table import Table
from rich import box

HISTORY_FILE = "history.json"
console = Console()


def load_history():
    if os.path.exists(HISTORY_FILE):
        with open(HISTORY_FILE, 'r') as f:
            return json.load(f)
    else:
        return []


def save_history(history):
    with open(HISTORY_FILE, 'w') as f:
        json.dump(history, f)


def print_history(history):
    table = Table(show_header=True, header_style="bold magenta", box=box.ROUNDED, show_lines=True)
    table.add_column("No.", width=10)
    table.add_column("Prompt", width=40, style="magenta")
    table.add_column("Response", width=50)

    for i, entry in enumerate(history):
        table.add_row(str(i + 1), entry['prompt'], entry['response'])
    console.print(table)


def execute_command(command):
    process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    stdout, stderr = process.communicate()

    if stdout:
        syntax = Syntax(stdout.decode(), "bash", theme="monokai", line_numbers=True)
        console.print("Output: ", syntax)
    if stderr:
        console.print("Error: ", stderr.decode(), style="bold red")

main.py:
import click

from .openai_api import raise_exception_if_no_key, request_to_completions
from .prompts import command_generation_prompt
from .shell_utils import load_history, save_history, print_history, execute_command
from rich.console import Console
from rich.syntax import Syntax
from threading import Thread, Event
from prompt_toolkit import prompt as prmt

raise_exception_if_no_key()
console = Console()

stop_spinner = Event()


def spin():
    with console.status("[bold green]Processing..."):
        while not stop_spinner.wait(timeout=0.1):
            pass


@click.command()
@click.option('--prompt', '-p', default=None, help='Prompt for ChatGPT')
@click.option('--cmd', '-c', default=None, help='Command to execute.')
@click.option('--interactive', '-i', is_flag=True, help='Run in interactive mode.')
def main(prompt, cmd, interactive):
    history = load_history()

    if prompt:
        process_other_prompt(prompt, history)
        save_history(history)

    if cmd:
        process_command_prompt(cmd, command_generation_prompt, history)
        save_history(history)
        return

    if interactive:
        running = True
        while running:
            prompt = get_prompt()

            if prompt == 'exit' or prompt == 'q':
                running = False
                continue

            if prompt == 'history':
                print_history(history)
            elif prompt.strip().startswith('command:'):
                process_command_prompt(prompt, command_generation_prompt, history)
            else:
                process_other_prompt(prompt, history)

            save_history(history)


def get_prompt():
    console.print("\nEnter a prompt: ", end="", style="bold cyan")
    return input()


def process_command_prompt(prompt, command_generation_prompt, history):
    request_prompt = command_generation_prompt + prompt.replace('command:', '', 1)

    spinner_thread = Thread(target=spin)
    spinner_thread.start()

    response = request_to_completions(request_prompt)

    stop_spinner.set()
    spinner_thread.join()
    stop_spinner.clear()

    console.print(f"\nThe command is: {response}\n", style="bold green")
    execute = input("Do you want to execute this command? (y/n) or edit (e): ")

    if execute.lower() == 'e':
        response = prmt('Enter your edited command: ', default=response)
        console.print(f"\nThe command is: {response}", style="bold green")
        execute = input("Do you want to execute this command? (y/n) or edit (e): ")

    if execute.lower() == 'y':
        execute_command(response)

    history.append({
        'prompt': prompt,
        'response': response
    })


def process_other_prompt(prompt, history):
    spinner_thread = Thread(target=spin)
    spinner_thread.start()

    response = request_to_completions(prompt)

    stop_spinner.set()
    spinner_thread.join()
    stop_spinner.clear()

    code_block_wrapper = '```'

    if code_block_wrapper in response:
        start = response.index(code_block_wrapper) + 3
        end = response.index(code_block_wrapper, start)
        code = response[start:end]
        syntax = Syntax(code, "python", theme="monokai", line_numbers=False)
        console.print(response[:start - 3])
        console.print(syntax)
        console.print(response[end + 3:])
    else:
        console.print(f"{response}", style="bold green")

    history.append({
        'prompt': prompt,
        'response': response
    })


if __name__ == "__main__":
    main()

prompts.py:
command_generation_prompt = """
    You are a shell command expert and your task is to provide functioning shell commands.
    I will ask you for a command I need and you will return a CLI command and nothing else. 
    Do not send it in a code block, quotes, or anything else, just provide the pure text CONTAINING ONLY THE COMMAND. DO NOT explain anything. 
    The command I need should do the following:"""

