#!/bin/bash

: "${PGPT_WORKER_MODE:=mixed}"  # mixed, worker, flower
: "${PGPT_FLOWER_PORT:=5555}"
: "${PGPT_FLOWER_PASSWORD:=flower}"
: "${PGPT_FLOWER_URL_PREFIX:=}"
: "${PGPT_FLOWER_PERSISTENT:=true}"
: "${PGPT_FLOWER_PERSISTENT_DB:=celery_flower.db}"
: "${PGPT_FLOWER_MAXIMUM_TASKS:=1000}"
: "${PGPT_CELERY_LOG_LEVEL:=info}"
: "${PGPT_CELERY_QUEUES:=}"
: "${PGPT_CELERY_HOSTNAME:=default}"
: "${PGPT_CELERY_POOL:=prefork}"
: "${PGPT_CELERY_CONCURRENCY:=}"
: "${PGPT_CELERY_TIME_LIMIT:=}"
: "${PGPT_WORKER_APP_MODULE:=private_gpt}"
: "${API_ENABLED:=true}"
: "${API_PORT:=8090}"

CELERY_APP_MODULE="${PGPT_WORKER_APP_MODULE}.celery"
HEALTHCHECK_APP_MODULE="${PGPT_WORKER_APP_MODULE}.celery.healthcheck:app"

# Create directory for PID files
mkdir -p /tmp/private_gpt_pids

# Add signal handling
cleanup() {
    echo "Shutting down services..."
    jobs -p | xargs -r kill
    exit 0
}
trap cleanup EXIT INT TERM

# Run flower on the specified port
if [[ "$PGPT_WORKER_MODE" == "flower" || "$PGPT_WORKER_MODE" == "mixed" ]]; then
    echo "Starting flower on port $PGPT_FLOWER_PORT"
    FLOWER_ARGS=()
    if [[ -n "$PGPT_FLOWER_URL_PREFIX" ]]; then
        FLOWER_ARGS+=("--url_prefix=$PGPT_FLOWER_URL_PREFIX")
    fi
    if [[ -n "$PGPT_FLOWER_PORT" ]]; then
        FLOWER_ARGS+=("--port=$PGPT_FLOWER_PORT")
    fi
    if [[ -n "$PGPT_FLOWER_PASSWORD" ]]; then
        FLOWER_ARGS+=("--basic_auth=flower:$PGPT_FLOWER_PASSWORD")
    fi
    if [[ "$PGPT_WORKER_MODE" == "mixed" ]]; then
        FLOWER_ARGS+=("--logging=none")
    fi
    if [[ "$PGPT_FLOWER_PERSISTENT" == "true" ]]; then
        FLOWER_ARGS+=("--persistent=True")
        if [[ -n "$PGPT_FLOWER_PERSISTENT_DB" ]]; then
            FLOWER_ARGS+=("--db=$PGPT_FLOWER_PERSISTENT_DB")
        fi
        if [[ -n "$PGPT_FLOWER_MAXIMUM_TASKS" ]]; then
            FLOWER_ARGS+=("--max-tasks=$PGPT_FLOWER_MAXIMUM_TASKS")
        fi
    fi

    # print all args except password
    args=("${FLOWER_ARGS[@]}")
    args=("${args[@]//:$PGPT_FLOWER_PASSWORD/}")

    echo "Starting flower with args: ${args[*]}"
    python -m celery --app "$CELERY_APP_MODULE" flower "${FLOWER_ARGS[@]}" &
fi

# Run celery worker using the config
if [[ "$PGPT_WORKER_MODE" == "worker" || "$PGPT_WORKER_MODE" == "mixed" ]]; then
  WORKER_ARGS=()
  if [[ -n "$PGPT_CELERY_LOG_LEVEL" ]]; then
      WORKER_ARGS+=("--loglevel=$PGPT_CELERY_LOG_LEVEL")
  fi
  if [[ -n "$PGPT_CELERY_QUEUES" ]]; then
      WORKER_ARGS+=("--queues=$PGPT_CELERY_QUEUES")
      if [[ -n "$PGPT_CELERY_HOSTNAME" ]]; then
        PGPT_CELERY_HOSTNAME="${PGPT_CELERY_QUEUES}"
      fi
  fi
  if [[ -n "$PGPT_CELERY_HOSTNAME" ]]; then
      WORKER_ARGS+=("--hostname=${PGPT_CELERY_HOSTNAME}%h")
  fi
  if [[ -n "$PGPT_CELERY_POOL" ]]; then
      WORKER_ARGS+=("--pool=$PGPT_CELERY_POOL")
  fi
  if [[ -n "$PGPT_CELERY_CONCURRENCY" ]]; then
      WORKER_ARGS+=("--concurrency=$PGPT_CELERY_CONCURRENCY")
  fi
  if [[ -n "$PGPT_CELERY_TIME_LIMIT" ]]; then
      WORKER_ARGS+=("--time-limit=$PGPT_CELERY_TIME_LIMIT")
  fi

  echo "Starting celery worker with args: ${WORKER_ARGS[*]}"
  python -m celery --app "$CELERY_APP_MODULE" worker "${WORKER_ARGS[@]}" &
fi

# Start the Healthcheck server
if [[ "$API_ENABLED" == "true" ]]; then
    echo "Starting healthcheck server on port $API_PORT"
    python -m uvicorn "$HEALTHCHECK_APP_MODULE" --host 0.0.0.0 --port $API_PORT --no-access-log --log-level critical &
fi

# Wait for all background processes
wait