#!/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}\`;"

mysql -e "CREATE TABLE IF NOT EXISTS \`${SERVERLESS_CLUSTER_DATABASE_NAME}\`.TEST2SQL(id int, name varchar(16));"

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

# 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 AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID} \
  -e AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY} \
  -e AWS_SESSION_TOKEN=${AWS_SESSION_TOKEN} \
  -v ./tidb-lab-ai:/root/tidb-lab-ai \
  cheesewang/tidb-lab-python-vscode:1.0