format_ps1() {
	local current_dir=$(realpath --relative-to=$ROOT/.. $PWD)
	if [ -z $CURRENT_FILE ]
	then
        echo "(No open file) $current_dir $ "
	else
		local open_file=$(realpath --relative-to=$PWD ${CURRENT_FILE:-$PWD})
        echo "(Open file: $open_file) $current_dir $ "
	fi
}

print() {
	local total_lines=$(wc -l < "$CURRENT_FILE")
	echo "[File: $(realpath --relative-to=$ROOT $CURRENT_FILE) ($total_lines lines total)]"
	cat $CURRENT_FILE | grep -n $ | head -n $(jq -n "[$CURRENT_LINE + $WINDOW, $WINDOW] | max | floor") | tail -n $(jq -n "2*$WINDOW")
}

print_search() {
	if [ -z "$SEARCH_RESULTS" ]
	then
		echo "No search results found."
		return
	fi
    export CURRENT_FILE=$(realpath ${SEARCH_FILES[$SEARCH_INDEX]})
	export CURRENT_LINE=${SEARCH_RESULTS[$SEARCH_INDEX]}
    echo "[Result "$(jq -n "($SEARCH_INDEX + 1)")"/${#SEARCH_RESULTS[@]} (Line $CURRENT_LINE)]"
	print
}

open() {
	if [ -z "$1" ]
	then
		echo "Usage: open <file>"
		return
	fi
    if [ -f "$1" ]
    then
        export CURRENT_FILE=$(realpath $1)
        export CURRENT_LINE=$WINDOW
        print
    else
        echo "File not found"
    fi
}

constrain_line() {
	if [ -z "$CURRENT_FILE" ]
	then
		echo "No file open. Use the open command first."
		return
	fi
	local max_line=$(wc -l < "$CURRENT_FILE")
	export CURRENT_LINE=$(jq -n "[$CURRENT_LINE, $max_line - $WINDOW + 1] | min") print
	export CURRENT_LINE=$(jq -n "[$CURRENT_LINE, $WINDOW] | max") print
}

scroll_down() {
	if [ -z "$CURRENT_FILE" ]
	then
		echo "No file open. Use the open command first."
		return
	fi
	export CURRENT_LINE=$(jq -n "$CURRENT_LINE + 2*$WINDOW - $OVERLAP")
    constrain_line
    print
}

scroll_up() {
	if [ -z "$CURRENT_FILE" ]
	then
		echo "No file open. Use the open command first."
		return
	fi
	export CURRENT_LINE=$(jq -n "$CURRENT_LINE - 2*$WINDOW + $OVERLAP")
    constrain_line
    print
}

search() {
	if [ $# -gt 1 ]; then
		echo "search allows only one search term at a time."
		return
	fi
	if [ -z "$1" ]
	then
		echo "Usage: search <term>"
		return
	fi
	if [ -z "$CURRENT_FILE" ]
	then
		echo "No file open. Use the open command first."
		return
	fi
	export SEARCH_FILES=( $(grep -nH "$@" $CURRENT_FILE | cut -d: -f1) )
	export SEARCH_RESULTS=( $(grep -nH "$@" $CURRENT_FILE | cut -d: -f2) )
	export SEARCH_INDEX=0
	print_search
}

search_all() {
	if [ $# -gt 1 ]; then
		echo "search_all allows only one search term at a time."
		return
	fi
	if [ -z "$1" ]
	then
		echo "Usage: search_all <term>"
		return
	fi
	export SEARCH_FILES=( $(grep -rnH "$@" | cut -d: -f1) )
	export SEARCH_RESULTS=( $(grep -rnH "$@" | cut -d: -f2) )
	export SEARCH_INDEX=0
	print_search
}

next() {
	if [ -z "$CURRENT_FILE" ]
	then
		echo "No file open. Use the open command first."
		return
	fi
	export SEARCH_INDEX=$(jq -n "($SEARCH_INDEX + 1) % ([${#SEARCH_RESULTS[@]}, 1] | max)")
	print_search
}

prev() {
	if [ -z "$CURRENT_FILE" ]
	then
		echo "No file open. Use the open command first."
		return
	fi
	export SEARCH_INDEX=$(jq -n "($SEARCH_INDEX - 1) % ([${#SEARCH_RESULTS[@]}, 1] | max)")
	print_search
}

# @yaml
# signature: "edit <n>:<m>\n<replacement_text>\nEND_OF_EDIT"
# docstring: replaces lines [n,m) with the given text in the open file. The replacement text is terminated by a line with only END_OF_EDIT on it. All of the <replacement text> will be entered, so make sure your indentation is formatted properly.
# end_name: END_OF_EDIT
edit() {
    if [ -z "$CURRENT_FILE" ]
	then
		echo "No file open. Use the open command first."
		return
	fi

    local start_line="$(echo $1: | cut -d: -f1)"
    local end_line="$(echo $1: | cut -d: -f2)"

    local replacement=()
    while IFS= read -r line
    do
		replacement+=("$line")
    done

	if [ -z "$start_line" ] || [ -z "$end_line" ]
	then
		echo "Usage: edit <start_line>:<end_line>"
		return
	fi

    local re='^[0-9]+$'
    if ! [[ $start_line =~ $re ]]; then
        echo "Usage: edit <start_line>:<end_line>"
        echo "Error: start_line must be a number"
        return
    fi
    if ! [[ $end_line =~ $re ]]; then
        echo "Usage: edit <start_line>:<end_line>"
        echo "Error: end_line must be a number"
        return
    fi

    # Bash array starts at 0, so let's adjust
    local start_line=$((start_line - 1))
    local end_line=$((end_line - 1))

    # Read the file line by line into an array
    mapfile -t lines < "$CURRENT_FILE"
    local new_lines=("${lines[@]:0:$start_line}" "${replacement[@]}" "${lines[@]:$((end_line))}")
    # Write the new stuff directly back into the original file
    printf "%s\n" "${new_lines[@]}" >| "$CURRENT_FILE"

    export CURRENT_LINE=$start_line
    constrain_line
    print
}

goto() {
	if [ $# -gt 1 ]; then
		echo "goto allows only one line number at a time."
		return
	fi
    if [ -z "$CURRENT_FILE" ]
	then
		echo "No file open. Use the open command first."
		return
	fi
	if [ -z "$1" ]
	then
		echo "Usage: goto <line>"
		return
	fi
    if ! [[ $1 =~ ^[0-9]+$ ]]
    then
		echo "Usage: goto <line>"
        echo "Error: <line> must be a number"
        return
    fi
	local max_line=$(wc -l < "$CURRENT_FILE")
	if [ $1 -gt $max_line ]
	then
		echo "Error: <line> must be less than or equal to $max_line"
		return
	fi
    export CURRENT_LINE=$1
    constrain_line
    print
}

submit() {
	cd $ROOT
	git add -A
	git diff --cached > model.patch
	echo "<<SUBMISSION||"
	cat model.patch
	echo "||SUBMISSION>>"
}

get_symbols() {
	if [ -z "$1" ]; then
		if [ -z "$CURRENT_FILE" ]; then
			echo "Usage: get_symbols <file>"
			return
		fi
		file_path="$CURRENT_FILE"
	else
		file_path="$1"
	fi

	if [ ! -f "$file_path" ]; then
		echo "File not found: $file_path"
		return
	fi

	/root/miniconda3/bin/python3 /root/tools/get_symbols.py "$file_path"
}

summarize() {
	if [ -z "$1" ]; then
		if [ -z "$CURRENT_FILE" ]; then
			echo "Usage: summarize <file>"  # bash command is summarize
			return
		fi
		file_path="$CURRENT_FILE"
	else
		file_path="$1"
	fi

	if [ ! -f "$file_path" ]; then
		echo "File not found: $file_path"
		return
	fi

	/root/miniconda3/bin/python3 /root/tools/get_summary.py "$file_path"
}

# this is a utility function for the get_assistant_data function in SWEEnv
get_assistant_data(){
	# check if /root/.get_assistant_data exists, if it does print the contents
	# otherwise print an empty json dict
	if [ -f "/root/.assistant_data" ]; then
		cat /root/.assistant_data
		rm -f /root/.assistant_data
	else
		echo "{}"
	fi
}