stack

#!/usr/bin/env bash
set -euo pipefail
IFS=$'\n\t'

export DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
export TMUX_SESSION=${TMUX_SESSION:="helix"}
export WITH_RUNNER=${WITH_RUNNER:=""}
export WITH_DEMOS=${WITH_DEMOS:=""}

function mock-runner() {
   go run . runner \
    --mock-runner \
    --server-port 8090 \
    --api-host http://localhost:8080 \
    --api-token oh-hallo-insecure-token \
    --memory 24GB \
    --runner-id mock \
    --label gpu=4090 "$@"
}

function compose() {
  docker compose -f docker-compose.dev.yaml "$@"
}

function build() {
  compose build
}

function static-compile() {
  export CGO_ENABLED=0
  go build -ldflags '-extldflags "-static"' -o helix .
}

function start() {
  if tmux has-session -t "$TMUX_SESSION" 2>/dev/null; then
    echo "Session $TMUX_SESSION already exists. Attaching..."
    sleep 1
    tmux -2 attach -t $TMUX_SESSION
    exit 0;
  fi

  export MANUALRUN=1
  export LOG_LEVEL=debug

  echo "Starting docker compose"

  local COMPOSE_FLAGS=""
  local COMPOSE_PROFILES=""

  if [[ -n "$WITH_RUNNER" ]]; then
    COMPOSE_FLAGS="--build"
    COMPOSE_PROFILES="$COMPOSE_PROFILES --profile dev_gpu_runner"
  fi

  if [[ -n "$WITH_DEMOS" ]]; then
    COMPOSE_PROFILES="$COMPOSE_PROFILES --profile demos"
  fi

  eval "docker compose -f docker-compose.dev.yaml $COMPOSE_PROFILES up $COMPOSE_FLAGS -d"

  sleep 2

  echo "Creating tmux session $TMUX_SESSION..."

  # get the size of the window and create a session at that size
  local screensize=$(stty size)
  local width=$(echo -n "$screensize" | awk '{print $2}')
  local height=$(echo -n "$screensize" | awk '{print $1}')
  tmux -2 new-session -d -s $TMUX_SESSION -x "$width" -y "$(($height - 1))"

  tmux split-window -v -d
  tmux select-pane -t 1
  tmux split-window -v -d
  tmux select-pane -t 0
  tmux split-window -v -d
  tmux select-pane -t 1
  tmux split-window -v -d
  tmux select-pane -t 0
  tmux split-window -v -d

  tmux send-keys -t 0 './stack compose logs -f frontend' C-m
  tmux send-keys -t 1 './stack compose logs -f api' C-m
  tmux send-keys -t 2 './stack compose logs -f typesense' C-m

  if [[ -n "$WITH_RUNNER" ]]; then
    tmux send-keys -t 3 './stack compose --profile dev_gpu_runner -f docker-compose.dev.yaml exec dev_gpu_runner bash' C-m
    tmux send-keys -t 3 'go run . runner --api-host http://172.17.0.1:8080 --api-token oh-hallo-insecure-token --memory 24GB --runner-id dev-runner' C-m
  fi

  if [[ -n "$WITH_DEMOS" ]]; then
    tmux send-keys -t 4 'docker compose --profile demos -f docker-compose.dev.yaml exec demos bash' C-m
    tmux send-keys -t 4 'go run .' C-m
  fi

  tmux -2 attach-session -t $TMUX_SESSION
}

function stop() {
  echo "Removing docker containers"
  ./stack compose --profile dev_gpu_runner --profile demos down
  echo "Stopping tmux session $TMUX_SESSION..."
  tmux kill-session -t $TMUX_SESSION ||true
}

function up() {
  docker compose -f docker-compose.dev.yaml up -d $@
}

function rebuild() {
  docker compose -f docker-compose.dev.yaml up -d --build $@
}

function db() {
  local subcommand="${1-cli}"
  shift
  local containername="${1-postgres}"
  shift
  if [[ "$subcommand" == "cli" ]]; then
    ./stack compose exec $containername psql --user postgres "$@"
  elif [[ "$subcommand" == "pipe" ]]; then
    ./stack compose exec -T $containername psql --user postgres "$@"
  fi
}

# Regenerate test mocks
function generate() {
  go generate ./...
}

function psql() {
  db cli postgres "$@"
}

function psql_pipe() {
  db pipe postgres "$@"
}

function pgvector() {
  db cli pgvector "$@"
}

function pgvector_pipe() {
  db pipe pgvector "$@"
}

function install() {
  go install ./api/..
}

function update_openapi() {
	go install github.com/swaggo/swag/cmd/swag@latest && \
	swag init -g api/pkg/server/swagger.go \
	--parseDependency --parseInternal --parseDepth 3 \
	-o api/pkg/server
}

function test() {
  # Ingest env variables from .env file
  set -a
  source .env
  set +a

  # Check whether environment variables are set. If not, error
  if [[ -z "$TOGETHER_API_KEY" ]]; then
    echo "TOGETHER_API_KEY is not set"
    exit 1
  fi
  if [[ -z "$TOGETHER_BASE_URL" ]]; then
    echo "TOGETHER_BASE_URL is not set"
    exit 1
  fi

  # Ensure postgres, tika, typesense and chrome are running
  docker compose -f docker-compose.dev.yaml up -d postgres tika typesense chrome

  # Database config (running in a sidecar)
  export POSTGRES_HOST=localhost
  export TYPESENSE_URL=http://localhost:8108
  export TYPESENSE_API_KEY=typesense
  export TEXT_EXTRACTION_TIKA_URL=http://localhost:9998
  export RAG_CRAWLER_LAUNCHER_URL=http://localhost:7317

  go test -v ./...
}

eval "$@"