#!/bin/bash

function update-with-llm() {
    if [ $# -lt 2 ]; then
        echo "Usage: update-with-llm [glob] [prompt]"
        exit 1
    fi

    # Get the last argument as the task description
    TASK_DESCRIPTION="${@: -1}"

    # All preceding arguments are file paths
    FILES=("${@:1:$#-1}")

    echo -e "Load these files into prompt:\n\n${FILES[*]}\n"
    echo -e "Task: $TASK_DESCRIPTION"

    FILE_CONTENTS=""
    for FILE in "${FILES[@]}"; do
        FILE_CONTENTS+="<<<$FILE>>>
$(cat "$FILE")

"
    done

    # Define prompt content
    mkdir -p tmp
    PROMPT_FILE="tmp/llm_update"
    echo -e "\nWrite prompt to $PROMPT_FILE."
    cat << EOF > "$PROMPT_FILE"
Look at the code files below and do the following:

$TASK_DESCRIPTION

Output all files that you need to change in full again, including your changes.
In the same format as I provide the files below. Under no circumstances output any other text,
no additional infos, no code formatting chars. Only the code in the given format.

Here are the files:

$FILE_CONTENTS
EOF

    # Run the llm command
    echo -e "Send prompt and wait for the response of the $(llm models | sed -n 's/^Default: //p') LLM."
    RESULT=$(cat "$PROMPT_FILE" | llm)

    # Check if result is empty
    if [ -z "$RESULT" ]; then
        echo "No response from the model. Exiting."
        exit 1
    fi

    # Save the result to a file
    RESULT_FILE="tmp/llm_update_result"
    echo "$RESULT" > "$RESULT_FILE"
    echo -e "Saved response to $RESULT_FILE.\n"

    # Show a preview of result file
    less "$RESULT_FILE"

    # Ask for confirmation before updating files
    read -p "Do you want to apply these updates to the files? (y/n): " CONFIRM
    if [[ ! "$CONFIRM" =~ ^[Yy]$ ]]; then
        exit 0
    fi

    # Parse the response from the file and update the files directly
    echo -e "Parsing the response and updating files...\n"
    CURRENT_FILE=""
    while IFS= read -r LINE; do
        if [[ $LINE =~ ^"<<<"(.*)">>>" ]]; then
            CURRENT_FILE="${BASH_REMATCH[1]}"
            echo "Update file $CURRENT_FILE."
            > "$CURRENT_FILE"
        elif [[ -n $CURRENT_FILE ]]; then
            echo "$LINE" >> "$CURRENT_FILE"
        fi
    done < "$RESULT_FILE"
}