#!/bin/bash

[ ! -f .env ] || export $(grep -v '^#' .env | xargs)
[ ! -f .teagent ] || export $(grep -v '^#' .teagent | xargs)
if [ -z "$REPO" ]; then
  exit 1
fi
# Default 5 minute maximum
TIMEOUT_MAX="${TIMEOUT_MAX:-300}"
N_RETRIES="${RETRY_MAX:-5}"

teaissue() {
  tea issues create --title "$1 (TeAgent)" --body "$2" --login teagent --repo "${REPO}"
}

SERVER="${OLLAMA_HOST:-http://localhost:11434}"
MODEL="${OLLAMA_MODEL:-llama3.2}"

read -p "Source file listing cmd: " -e -i "find src -type f -iname \\\\*.py" source_file_fn
SOURCE_FILES=`eval "${source_file_fn}"`

# Select what files
FILE_SELECTIONS=()
INDEX=1
for source_file in ${SOURCE_FILES}; do
  if [ ! -s "${source_file}" ]; then
    continue
  fi
  FILE_SELECTIONS+=("${source_file}" "" on)
  INDEX=$((INDEX+1))
done

FILE_CHOICES=$(dialog --clear \
  --title "Files to process" \
  --checklist "Choose files:" \
  0 0 0 \
  "${FILE_SELECTIONS[@]}" \
  2>&1 >/dev/tty)
if [ "$?" -ne "0" ]; then
  exit 1
fi

TASK_SELECTIONS=(
  1 "Clean" off
  2 "UnitTest" off
  3 "Cancel" on
)

TASK_CHOICE=$(dialog --clear \
  --title "Agent Task" \
  --radiolist "Select Task:" \
  15 40 3 \
  "${TASK_SELECTIONS[@]}" \
  2>&1 >/dev/tty)
if [ "$?" -ne "0" ]; then
  exit 1
fi

case $TASK_CHOICE in
  1)
    TASK_NAME="Tests"
    PROMPT_BASE=$(cat << EOF
Please clean up the following code, leaving ample documentation.
If the code seems clean already, simply write: DONE
EOF
)
    ;;
  2)
    TASK_NAME="Cleanup"
    PROMPT_BASE=$(cat << EOF
Please write tests for the following source, leaving ample documentation.
The testing backend of choice is pytest.
If there don't seem to be any natural tests, simply write: DONE
EOF
)
      ;;
  3)
    echo Bye!
    exit 0
    ;;
  *)
    echo "Invalid choice!"
    exit 1
    ;;
esac

for source_file in $FILE_CHOICES; do
  if [ ! -s "${source_file}" ]; then
    continue
  fi
  echo $source_file;
  SOURCE=$(<"${source_file}")
  PROMPT="${PROMPT_BASE}
${SOURCE}"
  TRIES=0
  while :; do
    echo Sending over for "${source_file}"
    RESPONSE=$(timeout "${TIMEOUT_MAX}" curl -s "$SERVER/api/generate" \
      -H "Content-Type: application/json" \
      -d "$(jq -n --arg model "$MODEL" --arg prompt "$PROMPT" \
            '{model:$model, prompt:$prompt, stream:false}')" \
            | jq -r '.response')
    TRIES="$((TRIES+1))"
    if [ "$?" -eq "0" ]; then
      echo "Response success!"
      break
    fi
    if [ "$TRIES" -ge "$N_RETRIES" ]; then
      echo "Response timeout!"
      break
    fi
  done
  if [ ! -z "${RESPONSE}" ]; then
    teaissue "${TASK_NAME} ${source_file}" "${RESPONSE}"
  fi
done