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

# Config (override with env vars if needed)
VOLUME_NAME="${VOLUME_NAME:-dynamodb_data}"
DDB_CONTAINER_NAME="${DDB_CONTAINER_NAME:-dynamodb-local}"
ADMIN_CONTAINER_NAME="${ADMIN_CONTAINER_NAME:-dynamodb-admin}"
DDB_PORT="${DDB_PORT:-8123}"        # DynamoDB Local host port (set to 8123 if you prefer)
ADMIN_PORT="${ADMIN_PORT:-8124}"    # dynamodb-admin UI port
REGION="${AWS_DEFAULT_REGION:-us-east-2}"
ENDPOINT_URL="http://localhost:${DDB_PORT}"

require_cmd() {
  command -v "$1" >/dev/null 2>&1 || { echo "Missing required command: $1"; exit 1; }
}

require_cmd docker
require_cmd aws

echo "Exporting local AWS credentials..."
export AWS_ACCESS_KEY_ID="${AWS_ACCESS_KEY_ID:-local}"
export AWS_SECRET_ACCESS_KEY="${AWS_SECRET_ACCESS_KEY:-local}"
export AWS_DEFAULT_REGION="${REGION}"

ensure_volume() {
  if ! docker volume inspect "${VOLUME_NAME}" >/dev/null 2>&1; then
    echo "Creating volume: ${VOLUME_NAME}"
    docker volume create "${VOLUME_NAME}" >/dev/null
  else
    echo "Volume already exists: ${VOLUME_NAME}"
  fi
}

ensure_ddb_container() {
  if ! docker ps -a --format '{{.Names}}' | grep -q "^${DDB_CONTAINER_NAME}$"; then
    echo "Starting new DynamoDB Local container: ${DDB_CONTAINER_NAME} (host:${DDB_PORT} -> container:8000)"
    docker run -d --name "${DDB_CONTAINER_NAME}" \
      -p "${DDB_PORT}:8000" \
      -v "${VOLUME_NAME}:/home/dynamodblocal/data" \
      amazon/dynamodb-local \
      -jar DynamoDBLocal.jar >/dev/null
  else
    STATUS=$(docker inspect -f '{{.State.Status}}' "${DDB_CONTAINER_NAME}")
    if [ "${STATUS}" != "running" ]; then
      echo "Starting existing container: ${DDB_CONTAINER_NAME}"
      docker start "${DDB_CONTAINER_NAME}" >/dev/null
    else
      echo "Container already running: ${DDB_CONTAINER_NAME}"
    fi
  fi
}

wait_for_local() {
  echo -n "Waiting for DynamoDB Local on ${ENDPOINT_URL}"
  for i in {1..30}; do
    if aws dynamodb list-tables --endpoint-url "${ENDPOINT_URL}" >/dev/null 2>&1; then
      echo " ... ready"
      return
    fi
    echo -n "."
    sleep 1
  done
  echo
  echo "DynamoDB Local did not become ready in time."
  exit 1
}

ensure_table() {
  local TABLE_NAME="$1"
  local CREATE_CMD="$2"

  if aws dynamodb describe-table --table-name "${TABLE_NAME}" --endpoint-url "${ENDPOINT_URL}" >/dev/null 2>&1; then
    echo "Table exists: ${TABLE_NAME}"
  else
    echo "Creating table: ${TABLE_NAME}"
    eval "${CREATE_CMD}"
    echo "Waiting for table to become ACTIVE: ${TABLE_NAME}"
    aws dynamodb wait table-exists --table-name "${TABLE_NAME}" --endpoint-url "${ENDPOINT_URL}"
  fi
}

ensure_admin_container() {
  # Use host.docker.internal to reach DynamoDB Local from inside the admin container (macOS/Windows)
  local DYNAMO_ENDPOINT="http://host.docker.internal:${DDB_PORT}"

  if ! docker ps -a --format '{{.Names}}' | grep -q "^${ADMIN_CONTAINER_NAME}$"; then
    echo "Starting dynamodb-admin UI: ${ADMIN_CONTAINER_NAME} (http://localhost:${ADMIN_PORT}) -> ${DYNAMO_ENDPOINT}"
    docker run -d --name "${ADMIN_CONTAINER_NAME}" \
      -p "${ADMIN_PORT}:8001" \
      -e AWS_REGION="${REGION}" \
      -e AWS_ACCESS_KEY_ID="${AWS_ACCESS_KEY_ID}" \
      -e AWS_SECRET_ACCESS_KEY="${AWS_SECRET_ACCESS_KEY}" \
      -e DYNAMO_ENDPOINT="${DYNAMO_ENDPOINT}" \
      aaronshaf/dynamodb-admin >/dev/null
  else
    STATUS=$(docker inspect -f '{{.State.Status}}' "${ADMIN_CONTAINER_NAME}")
    if [ "${STATUS}" != "running" ]; then
      echo "Starting existing container: ${ADMIN_CONTAINER_NAME}"
      docker start "${ADMIN_CONTAINER_NAME}" >/dev/null
    else
      echo "Container already running: ${ADMIN_CONTAINER_NAME}"
    fi
  fi

  echo "dynamodb-admin UI: http://localhost:${ADMIN_PORT}"
}

main() {
  echo "Using config:"
  echo "  DynamoDB Local port: ${DDB_PORT}"
  echo "  dynamodb-admin port: ${ADMIN_PORT}"
  echo

  ensure_volume
  ensure_ddb_container
  wait_for_local

  # clip-meta-heavy (lyricsTable) - PK: clipId (S)
  ensure_table "clip-meta-heavy" "
aws dynamodb create-table \
  --table-name clip-meta-heavy \
  --attribute-definitions AttributeName=clipId,AttributeType=S \
  --key-schema AttributeName=clipId,KeyType=HASH \
  --billing-mode PAY_PER_REQUEST \
  --endpoint-url ${ENDPOINT_URL}
"

  # clip-gen-config (clipConfigTable) - PK: clipId (S), SK: type (S)
  ensure_table "clip-gen-config" "
aws dynamodb create-table \
  --table-name clip-gen-config \
  --attribute-definitions \
      AttributeName=clipId,AttributeType=S \
      AttributeName=type,AttributeType=S \
  --key-schema \
      AttributeName=clipId,KeyType=HASH \
      AttributeName=type,KeyType=RANGE \
  --billing-mode PAY_PER_REQUEST \
  --endpoint-url ${ENDPOINT_URL}
"

  # item-info (itemInfoTable) - PK: itemId (S), SK: type (S)
  ensure_table "item-info" "
aws dynamodb create-table \
  --table-name item-info \
  --attribute-definitions \
      AttributeName=itemId,AttributeType=S \
      AttributeName=type,AttributeType=S \
  --key-schema \
      AttributeName=itemId,KeyType=HASH \
      AttributeName=type,KeyType=RANGE \
  --billing-mode PAY_PER_REQUEST \
  --endpoint-url ${ENDPOINT_URL}
"

  # orpheus (Orpheus sessions/chat data) - PK: item_id (S), SK: type (S)
  ensure_table "orpheus" "
aws dynamodb create-table \
  --table-name orpheus \
  --attribute-definitions \
      AttributeName=item_id,AttributeType=S \
      AttributeName=type,AttributeType=S \
      AttributeName=user_id,AttributeType=S \
  --key-schema \
      AttributeName=item_id,KeyType=HASH \
      AttributeName=type,KeyType=RANGE \
  --global-secondary-indexes \
      'IndexName=UserTypeIndex,KeySchema=[{AttributeName=user_id,KeyType=HASH},{AttributeName=type,KeyType=RANGE}],Projection={ProjectionType=ALL}' \
  --billing-mode PAY_PER_REQUEST \
  --endpoint-url ${ENDPOINT_URL}
"

  # query-embedding-cache (QueryEmbeddingCacheHandler) - PK: query_hash (S), SK: embedding_type (S)
  ensure_table "query-embedding-cache" "
aws dynamodb create-table \
  --table-name query-embedding-cache \
  --attribute-definitions \
      AttributeName=query_hash,AttributeType=S \
      AttributeName=embedding_type,AttributeType=S \
  --key-schema \
      AttributeName=query_hash,KeyType=HASH \
      AttributeName=embedding_type,KeyType=RANGE \
  --billing-mode PAY_PER_REQUEST \
  --endpoint-url ${ENDPOINT_URL}
"

  echo
  echo "Tables present:"
  aws dynamodb list-tables --endpoint-url "${ENDPOINT_URL}"
  echo

  ensure_admin_container
  echo "Done."
}

main "$@"