#!/usr/bin/env bash
set -euo pipefail

pids=()

generate_litellm_provider() {
	provider="$1"

	if [[ "$#" -lt 2 ]]; then
		>&2 echo "[+] Provider ${provider}: No key needed"
		cat "/assets/${provider}.yaml"
		return
	fi

	key_env="$2"

	if [[ -n "${!key_env:-}" ]]; then
		>&2 echo "[+] Provider ${provider}: Found key ${key_env}"
		cat "/assets/${provider}.yaml"
		return
	fi

	>&2 echo "[!] Provider ${provider}: No key found (set ${key_env})"
}

generate_litellm_config() {
	echo "model_list:"
	generate_litellm_provider "anthropic" "ANTHROPIC_API_KEY"
	generate_litellm_provider "github" "GITHUB_API_KEY"
	generate_litellm_provider "groq" "GROQ_API_KEY"
	generate_litellm_provider "huggingface"
	generate_litellm_provider "openai" "OPENAI_API_KEY"
	generate_litellm_provider "perplexity" "PERPLEXITY_API_KEY"

	if [[ -n "${LITELLM_MODELS_BASE64:-}" ]]; then
		echo "${LITELLM_MODELS_BASE64}" | base64 -d
	fi
}

start_litellm() {
	(
	# KISS: No persistence for LiteLLM
	unset DATABASE_URL

	litellm \
		--host "127.0.0.1" \
		--port "4000" \
		--config /tmp/litellm_config.yaml
	) &
	pids+=("$!")
}

start_open_webui() {
	(
	# By default, we expect it to be deployed to a private HF Space.
	# You can enable WEBUI_AUTH as needed
	if [[ -z "${WEBUI_AUTH:-}" ]]; then
		>&2 echo "[!] Enabling single user mode"
		export WEBUI_AUTH="False"
	fi

	# The less the user needs to configure, the better :)
	if [[ -n "${OPENAI_API_KEY:-}" ]] && [[ -z "${WEBUI_SECRET_KEY:-}${WEBUI_JWT_SECRET_KEY:-}" ]]; then
		>&2 echo "[!] Using OpenAI API key as Web UI secret key"
		export WEBUI_SECRET_KEY="${OPENAI_API_KEY}"
	fi

	if [[ -z "${ENABLE_RAG_WEB_SEARCH:-}${RAG_WEB_SEARCH_ENGINE:-}" ]]; then
		if [[ -n "${BRAVE_SEARCH_API_KEY:-}" ]]; then
			export RAG_WEB_SEARCH_ENGINE="brave"
			export ENABLE_RAG_WEB_SEARCH="True"
		fi
	fi

	if [[ -z "${DATABASE_URL:-}" ]]; then
		unset DATABASE_URL
	fi

	export ENABLE_OLLAMA_API="${ENABLE_OLLAMA_API:-False}"

	export OPENAI_API_BASE_URLS="http://localhost:4000"
	export OPENAI_API_KEY="sk-unused"

	export ENABLE_IMAGE_GENERATION="True"
	export IMAGES_OPENAI_API_BASE_URL="http://localhost:4000"
	export IMAGES_OPENAI_API_KEY="sk-unused"

	export PGSSLCERT=/tmp/postgresql.crt

	env
	/app/backend/start.sh
	) &
	pids+=("$!")
}

wait_litellm() {
	while ! curl -s http://localhost:4000 >/dev/null; do
		>&2 echo "[!] Waiting for LiteLLM..."
		sleep 1
	done
}

generate_litellm_config >/tmp/litellm_config.yaml
>&2 cat /tmp/litellm_config.yaml

start_litellm
wait_litellm
start_open_webui

for pid in "${pids[@]}"; do
    wait "$pid"
done