#!/usr/bin/env bash

SCRIPT=$(readlink -f "$0")
SCRIPTPATH=$(dirname "$SCRIPT")
cd "$SCRIPTPATH" || exit

if [[ ! -x $(command -v aichat) ]]; then
  echo "aichat is not installed."
  if [[ $(docker compose ps -q aichat) ]]; then
    docker compose cp "aichat:/usr/bin/aichat" "./tools/"
  else
    cargo install aichat@0.29.0
  fi
fi

source .env
mkdir -p ~/.config/aichat/
cat src/aichat/config.yaml | sed "s/__LLM_API_KEY__/${LLM_API_KEY}/" >~/.config/aichat/config.yaml
OLLAMA_HOST=$(cat .env | grep OLLAMA_HOST | cut -d'=' -f2 | cut -d'/' -f3)
if [[ ! -z $OLLAMA_HOST ]]; then
  cat ~/.config/aichat/config.yaml | sed "s/http/https/" >~/.config/aichat/config.yaml.tmp
  cat ~/.config/aichat/config.yaml.tmp | sed "s/localhost:11434/${OLLAMA_HOST}/" >~/.config/aichat/config.yaml
fi
DEFAULT_MODEL=$(cat .env | grep '^MODELS=' | cut -d'=' -f2 | cut -d',' -f2)
EMBEDDING_MODEL=$(cat .env | grep '^MODELS=' | cut -d'=' -f2 | cut -d',' -f1)
if [[ -z ${EMBEDDING_MODEL} ]] || [[ -z ${DEFAULT_MODEL} ]]; then
  echo "MODELS in .env are missing."
  exit 1
fi

cat ~/.config/aichat/config.yaml | sed "s/__EMBEDDING_MODEL__/ollama:${EMBEDDING_MODEL}/" >~/.config/aichat/config.yaml.tmp
cat ~/.config/aichat/config.yaml.tmp | sed "s/__DEFAULT_MODEL__/ollama:${DEFAULT_MODEL}/" >~/.config/aichat/config.yaml

MODELS=$(cat .env | grep '^MODELS=' | cut -d'=' -f2 | cut -d',' -f2- | tr ',' ' ')
for i in ${MODELS}; do
  echo "    - name: ${i}" >>~/.config/aichat/config.yaml
  echo
  read -p "Does the model ${i} support function calling? [y/N] " -n 1 -r
  if [[ $REPLY =~ ^[Yy]$ ]]; then
    echo "      supports_function_calling: true" >>~/.config/aichat/config.yaml
  else
    echo "      supports_function_calling: false" >>~/.config/aichat/config.yaml
  fi
  echo
  read -rp "What is the maximum number of input tokens ${i} supports? [128000] " max_input_tokens
  max_input_tokens=${max_input_tokens:-128000}
  echo "      max_input_tokens: ${max_input_tokens}" >>~/.config/aichat/config.yaml
done

mkdir -p ~/.config/aichat/roles
cp src/aichat/roles/* ~/.config/aichat/roles/