#!/bin/bash
set -e

# Restore the original password
ORIGINAL_PASSWORD=$(echo "${SERVERLESS_CLUSTER_PASSWORD}" | sed 's/\\\$/\$/g' | sed 's/\\\!/\!/g' | sed 's/\\\^/\^/g')

# Configure AWS CLI
CREDENTIALS=$(aws sts assume-role \
  --role-arn "arn:aws:iam::841162690310:role/AssumeBedrockRole" \
  --role-session-name "${USER_UNIQUE_TAG}-421-session" \
  --duration-seconds 7200)

AWS_ACCESS_KEY_ID=$(echo $CREDENTIALS | jq -r '.Credentials.AccessKeyId')
AWS_SECRET_ACCESS_KEY=$(echo $CREDENTIALS | jq -r '.Credentials.SecretAccessKey')
AWS_SESSION_TOKEN=$(echo $CREDENTIALS | jq -r '.Credentials.SessionToken')

# Creating the database if it doesn't exist
echo "Creating the database if it doesn't exist..."
cat > ~/.my.cnf <<EOF
[client]
host=${SERVERLESS_CLUSTER_HOST}
port=${SERVERLESS_CLUSTER_PORT}
user='${SERVERLESS_CLUSTER_USERNAME}'
password='${ORIGINAL_PASSWORD}'
EOF

chmod 600 ~/.my.cnf

mysql -e "CREATE DATABASE IF NOT EXISTS \`${SERVERLESS_CLUSTER_DATABASE_NAME}\`;"

# Clone the repository
echo "Cloning repository..."
git clone https://github.com/Icemap/tidb-lab-branches

# Run Docker container
echo "Starting Docker container..."
sudo docker run -d \
  -p 3000:8080 \
  -p 8501:8501 \
  -p 8502:8502 \
  -e SERVERLESS_CLUSTER_HOST=${SERVERLESS_CLUSTER_HOST} \
  -e SERVERLESS_CLUSTER_PORT=${SERVERLESS_CLUSTER_PORT} \
  -e SERVERLESS_CLUSTER_USERNAME=${SERVERLESS_CLUSTER_USERNAME} \
  -e SERVERLESS_CLUSTER_PASSWORD=${ORIGINAL_PASSWORD} \
  -e SERVERLESS_CLUSTER_DATABASE_NAME=${SERVERLESS_CLUSTER_DATABASE_NAME} \
  -e SERVERLESS_CLUSTER_ID=${SERVERLESS_CLUSTER_ID} \
  -e TIDBCLOUD_PROJECT_ID=${SERVERLESS_CLUSTER_PROJECT_ID} \
  -e TIDBCLOUD_ACCESS_TOKEN=${TIDBCLOUD_ACCESS_TOKEN} \
  -e AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID} \
  -e AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY} \
  -e AWS_SESSION_TOKEN=${AWS_SESSION_TOKEN} \
  -e BEDROCK_MODEL_ID='arn:aws:bedrock:us-west-2:841162690310:application-inference-profile/4i83xkynoouo' \
  -v ./tidb-lab-branches:/root/tidb-lab-branches \
  cheesewang/tidb-lab-python-vscode:1.0