#!/bin/bash
# bootstrap.sh - Initialize SIA (Self-Improving Agent) environment for cloud deployment

set -eo pipefail # Exit on any error, pipe failures

# Hardcoded paths for cloud deployment
SIA_REPO_URL="ssh://git@git.nielsgeens.be:222/llm/SIA.git"
SIA_DIR="/root/sia"
DATA_DIR="/root/data"
MODELS_DIR="/root/models"
DESKTOP_DIR="/root/desktop"
STATIC_DIR="/root/static"
VENVS_DIR="/root/venvs"

# Print header
echo "==================================================="
echo "SIA Bootstrap Script - Cloud Deployment"
echo "==================================================="

# Create directory structure
echo "Creating directory structure..."
mkdir -p "$DATA_DIR/iterations"
mkdir -p "$DESKTOP_DIR"
mkdir -p "$VENVS_DIR"
cd "$DESKTOP_DIR"

# Set up SSH keys
echo "Setting up SSH keys for git access..."
mkdir -p ~/.ssh
chmod 700 ~/.ssh
ssh-keygen -t sia_git -N "" -f ~/.ssh/sia_git -C "sia-agent"
echo "New SSH key generated"

# Display public key for user to add to git server
echo "==================================================="
echo "Add this public key to your git server:"
cat ~/.ssh/sia_git.pub
echo "==================================================="

# Prompt user to confirm they've added the key
read -p "Press Enter once you've added the SSH key to the git server..."

# Clone SIA repository
echo "Cloning SIA repository..."
git clone "$SIA_REPO_URL" "$SIA_DIR"

# Create and setup virtual environments
echo "Setting up SIA virtual environments..."

# Setup ITB tool environment
echo "Creating ITB tool environment..."
python3 -m venv "$VENVS_DIR/itb"
"$VENVS_DIR/itb/bin/pip" install -e "$SIA_DIR/tools/itb"

# Setup Train tool environment
echo "Creating Train tool environment..."
python3 -m venv "$VENVS_DIR/train"
"$VENVS_DIR/train/bin/pip" install -e "$SIA_DIR/tools/train"

# Setup SIA core environment
echo "Creating SIA core environment..."
python3 -m venv "$VENVS_DIR/sia"
"$VENVS_DIR/sia/bin/pip" install -e "$SIA_DIR"

# Build web interface
echo "Building web interface"
cd "$SIA_DIR/web"

# Install Node.js if needed
if ! command -v node &> /dev/null; then
    echo "Installing Node.js..."
    curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.1/install.sh | bash
    export NVM_DIR="$HOME/.nvm"
    [ -s "$NVM_DIR/nvm.sh" ] && \. "$NVM_DIR/nvm.sh"
    nvm install node
fi

npm install
npm run build

mkdir -p "$STATIC_DIR"
cp -r "$SIA_DIR/web/dist/"* "$STATIC_DIR/"

echo "Web interface built successfully"

# Finetune model
echo "Starting model finetuning..."

COMMIT_ID=$(cd "$SIA_DIR" && git rev-parse HEAD)
echo "Current commit: $COMMIT_ID"

mkdir -p "$MODELS_DIR/$COMMIT_ID"
mkdir -p "$MODELS_DIR/current"

# Run finetuning using the train environment
"$VENVS_DIR/train/bin/train_deepseek" --output-dir "$MODELS_DIR/$COMMIT_ID"

ln -sf "$MODELS_DIR/$COMMIT_ID" "$MODELS_DIR/current"
echo "Finetuning complete, model linked to current"

# Initialize environment information
echo "Initializing environment information..."
mkdir -p "$DATA_DIR/environment"
cat > "$DATA_DIR/environment/sia_repo.md" << EOF
# SIA Repository Information

- Repository URL: $SIA_REPO_URL
- ssh key: ~/.ssh/sia_git.pub
EOF

# Create .env file for local model only
cat > "$SIA_DIR/.env" << EOF
SIA_DEEPSEEK_ENABLED=true
SIA_DEEPSEEK_MODEL=$MODELS_DIR/current
SIA_DEEPSEEK_TEMPERATURE=0.6
EOF

# Print header
echo "==================================================="
echo "SIA environment initialization complete!"
echo "==================================================="

# Start SIA using restart script
echo "Starting SIA..."
"$SIA_DIR/scripts/restart.sh"