#!/bin/bash
set -e

# Save all environment variables for later SSH sessions
cat > /home/ubuntu/cloud-env.sh << EOF
export REGION_CODE=${REGION_CODE}
export REGION_NAME=${REGION_NAME}
export HOST_MONITOR1_PRIVATE_IP=${HOST_MONITOR1_PRIVATE_IP}
export HOST_MONITOR1_PUBLIC_IP=${HOST_MONITOR1_PUBLIC_IP}
EOF

# Configure AWS CLI for Bedrock access
echo "Configuring AWS credentials for Bedrock..."
CREDENTIALS=$(aws sts assume-role \
  --role-arn "arn:aws:iam::841162690310:role/AssumeBedrockRole" \
  --role-session-name "${USER_UNIQUE_TAG}-901-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')

# Write AWS credentials to cloud-env.sh
cat >> /home/ubuntu/cloud-env.sh << EOF
export AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID}
export AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY}
export AWS_SESSION_TOKEN=${AWS_SESSION_TOKEN}
EOF

# Tag subsequent ti commands emitted from this lab for telemetry.
echo 'export TI_TELEMETRY_TAG="tidb-labs"' >> /home/ubuntu/cloud-env.sh

# Clone lab materials
echo "Cloning lab materials..."
cd /home/ubuntu
git clone --depth 1 https://github.com/pingcap-inc/tidb-course-labs.git
cp /home/ubuntu/tidb-course-labs/labs/demo_901/template-config.yaml.solution /home/ubuntu/
rm -rf tidb-course-labs/
cp template-config.yaml.solution config.yaml.solution
sed -i'' \
  -e "s/<REGION_CODE>/${REGION_CODE}/g" \
  ./config.yaml.solution 2>/dev/null
rm -f template-config.yaml.solution

<< 'COMMENT'
# Install uv
echo "Installing uv..."
curl -LsSf https://astral.sh/uv/install.sh | sh
echo 'export PATH="$HOME/.local/bin:$PATH"' >> /home/ubuntu/cloud-env.sh

# Install Python dependencies
echo "Installing Python dependencies..."
export HOME=/home/ubuntu
export PATH="$HOME/.local/bin:$PATH"
uv venv --python=python3.12
source .venv/bin/activate
uv pip install -r ${LAB_DIR}/requirements.txt

# Install Hermes Agent
export HERMES_HOME=/home/ubuntu/.hermes
curl -fsSL https://raw.githubusercontent.com/NousResearch/hermes-agent/main/scripts/install.sh | bash -s -- --skip-setup
echo 'export PATH="$HOME/.hermes/bin:$PATH"' >> /home/ubuntu/cloud-env.sh

# Config Hermes
hermes config set model.provider bedrock
hermes config set model.default deepseek.v3.2
hermes config set model.base_url https://bedrock-runtime.us-west-2.amazonaws.com
hermes config set model.region us-west-2
COMMENT

# Add cloud-env.sh to .bashrc
echo "source /home/ubuntu/cloud-env.sh" >> /home/ubuntu/.bashrc