#!/usr/bin/env bash

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
    local task_description="${@: -1}"

    # Check if the last argument is a file path
    if [[ -f "$task_description" ]]; then
        # Read the prompt from the file
        task_description=$(cat "$task_description")
        # All preceding arguments are file paths
        local files=("${@:1:$#-1}")
    else
        # All preceding arguments are file paths
        local files=("${@:1:$#-1}")
    fi

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

    local file_contents=""
    for file in "${files[@]}"; do
        file_contents+="<<<$file>>>
$(cat "$file")
"
    done

    # Define prompt content
    mkdir -p tmp
    local prompt_file="tmp/llm_prompt"
    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.
Use the <<<...>>> marker to define the path to the file and the start of the content.

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."
    local 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
    local result_file="tmp/llm_prompt_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"
    local 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"
}
uwl() {
    update-with-llm "$@"
}
