#!/bin/bash

# ==============================================================================
#                 CONFIGURATION
# ==============================================================================

# Steps to run this script:
# 1. Set cloud_provider to the cloud provider you want to launch the instance to.
# 2. Fill out the relvant configuration variables for the cloud provider you are launching to.
# 3. Log in to your cloud provider via its CLI prior to running this script.
# 4. Run this script.
# 5. Connect to the instance with the provided SSH command at the end
# 6. Run the command "calypsoai-installer" to run the installer.

# Set Cloud Provider to launch instance to
cloud_provider="aws"                  # Set to 'aws', 'azure', or 'gcp'.

# --- AWS Configuration ---
AWS_REGION="us-east-2"
AWS_PROFILE="your-aws-profile"                   # AWS profile to use (leave empty for default)
AWS_KEY_PAIR_NAME="your-key-pair"   # Must exist in the specified AWS region.
AWS_SECURITY_GROUP_ID="your-security-group-id" # Must exist in your AWS account.
AWS_IMAGE_ID="resolve:ssm:/aws/service/canonical/ubuntu/server/24.04/stable/current/amd64/hvm/ebs-gp3/ami-id"   # Ubuntu 24.04 LTS for us-east-2 (must be updated if region is changed)
AWS_INSTANCE_TYPE="t3.medium"
AWS_INSTANCE_NAME="calypso-installer"
AWS_USER="ubuntu"
AWS_ROOT_VOLUME_SIZE="100"                       

# --- Azure Configuration ---
AZURE_RESOURCE_GROUP="your-resource-group" # Must exist in your Azure account.
AZURE_LOCATION="eastus"
AZURE_VM_NAME="calypso-installer"
AZURE_IMAGE="Ubuntu2404"              
AZURE_SIZE="Standard_B2s"
AZURE_ADMIN_USER="azureuser"
AZURE_SSH_PUBLIC_KEY_PATH="your-ssh-public-key.pub" # Path to your SSH public key.
AZURE_SSH_PRIVATE_KEY_PATH="your-ssh-private-key.pem" # Path to your SSH private key.
AZURE_DISK_SIZE="100"               

# --- GCP Configuration ---
GCP_PROJECT="your-gcp-project-id"    # Your GCP project ID where the VM will be created.
GCP_ZONE="us-central1-a"
GCP_MACHINE_TYPE="n2-standard-2"
GCP_IMAGE_FAMILY="ubuntu-2404-lts-amd64"
GCP_IMAGE_NAME="ubuntu-2404-noble-amd64-v20250725"
GCP_IMAGE_PROJECT="ubuntu-os-cloud"
GCP_INSTANCE_NAME="calypso-installer"
GCP_BOOT_DISK_SIZE="100"            

# --- Certificate File Configuration ---
CERT_FILEPATH=""                    # Path to certificate file (e.g., "/path/to/cert.crt")
CERT_KEY_FILEPATH=""                # Path to certificate key file (e.g., "/path/to/cert.key")
CERT_CHAIN_FILEPATH=""              # Path to certificate chain file (e.g., "/path/to/cert-chain.crt")

# Function to generate setup script based on cloud provider
generate_setup_script() {
    local provider=$1
    local username=""
    
    # Set username based on cloud provider
    case "$provider" in
        aws) username="ubuntu" ;;
        azure) username="azureuser" ;;
        gcp) username="ubuntu" ;;
    esac
    
    cat << 'SETUP_SCRIPT_EOF'
#!/bin/bash
echo "=========================================="
echo "F5 AI Security Installer Setup Starting..."
echo "=========================================="

# Update system packages
echo "Updating system packages..."
sudo apt-get update
sudo apt-get upgrade -y

# Install essential packages
echo "Installing essential packages..."
sudo apt-get install -y \
    curl \
    wget \
    git \
    unzip \
    software-properties-common \
    apt-transport-https \
    ca-certificates \
    gnupg \
    lsb-release \
    build-essential \
    libssl-dev \
    libffi-dev \
    python3-dev \
    python3-pip \
    python3-venv

# Install Python 3.12.8
echo "Setting up Python 3.12.8..."
sudo add-apt-repository ppa:deadsnakes/ppa -y
sudo apt-get update
sudo apt-get install -y python3.12 python3.12-venv python3.12-dev

# Create symlinks for python3.12
sudo ln -sf /usr/bin/python3.12 /usr/bin/python3
sudo ln -sf /usr/bin/python3.12 /usr/bin/python

# Upgrade pip
python3 -m pip install --upgrade pip

# Create application directory
echo "Creating application directory..."
sudo mkdir -p /opt/calypsoai-installer
sudo chmod 777 /opt/calypsoai-installer
cd /opt/calypsoai-installer
sleep 3

# Create certificate directory and copy certificate files
echo "Setting up certificate files..."
sudo mkdir -p /opt/cert_files
sudo chmod 755 /opt/cert_files

# Download the installer from S3
echo "Downloading F5 AI Security installer from S3..."
curl -L -o /opt/calypsoai-installer/cai-deploy-latest-linux-x86_64 https://calypsoai-installer.s3.us-east-1.amazonaws.com/installers/cai-deploy-latest-linux-x86_64

# Set permissions for the installer
echo "Setting installer permissions..."
sudo chmod +x *.exe 2>/dev/null || true
sudo chmod +x cai-deploy-* 2>/dev/null || true

# Create a symlink for easy access
if [ -f "/opt/calypsoai-installer/cai-deploy-latest-linux-x86_64" ]; then
    sudo ln -sf /opt/calypsoai-installer/cai-deploy-latest-linux-x86_64 /usr/local/bin/calypsoai-installer
fi

# Install uv
echo "Installing uv package manager..."
curl -LsSf https://astral.sh/uv/install.sh | sh
export PATH="$HOME/.cargo/bin:$PATH"
echo 'export PATH="$HOME/.cargo/bin:$PATH"' >> ~/.bashrc

# Install Pulumi
echo "Installing Pulumi..."
curl -fsSL https://get.pulumi.com | sh
export PATH="$HOME/.pulumi/bin:$PATH"
echo 'export PATH="$HOME/.pulumi/bin:$PATH"' >> ~/.bashrc

# Install kubectl
echo "Installing kubectl..."
curl -LO "https://dl.k8s.io/release/$(curl -L -s https://dl.k8s.io/release/stable.txt)/bin/linux/amd64/kubectl"
chmod +x kubectl
sudo mv kubectl /usr/local/bin/

# Install Helm
echo "Installing Helm..."
curl https://baltocdn.com/helm/signing.asc | gpg --dearmor | sudo tee /usr/share/keyrings/helm.gpg > /dev/null
echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/helm.gpg] https://baltocdn.com/helm/stable/debian/ all main" | sudo tee /etc/apt/sources.list.d/helm-stable-debian.list
sudo apt-get update
sudo apt-get install -y helm

# Install AWS CLI
echo "Installing AWS CLI..."
curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip"
unzip awscliv2.zip
sudo ./aws/install
rm -rf aws awscliv2.zip

# Install Azure CLI
echo "Installing Azure CLI..."
curl -sL https://aka.ms/InstallAzureCLIDeb | sudo bash

# Install Google Cloud CLI
echo "Installing Google Cloud CLI..."
curl https://sdk.cloud.google.com | bash
export PATH="$HOME/google-cloud-sdk/bin:$PATH"
echo 'export PATH="$HOME/google-cloud-sdk/bin:$PATH"' >> ~/.bashrc
source ~/.bashrc

# Install GKE gcloud auth plugin
echo "Installing GKE gcloud auth plugin..."
gcloud components install gke-gcloud-auth-plugin

# Add gcloud to PATH for current session
echo "Adding gcloud to PATH..."
export PATH="/usr/bin:$PATH"
source ~/.bashrc 2>/dev/null || true

# Install OCI CLI
echo "Installing OCI CLI..."
bash -c "$(curl -L https://raw.githubusercontent.com/oracle/oci-cli/master/scripts/install/install.sh)" -- --accept-all-defaults
export PATH="$HOME/bin:$PATH"
echo 'export PATH="$HOME/bin:$PATH"' >> ~/.bashrc

# Set up environment variables
echo "Setting up environment variables..."
cat >> ~/.bashrc << 'ENVEOF'

# F5 AI Security Installer Environment
export PATH="/opt/calypsoai-installer:$PATH"
export PULUMI_BACKEND_URL="file:///opt/calypsoai-installer/cai-pulumi-state"

# Cloud provider configurations
export AWS_DEFAULT_REGION=us-east-1
export AZURE_DEFAULT_LOCATION=eastus
export GOOGLE_CLOUD_PROJECT=your-project-id

# Harbor registry configuration
export HARBOR_REGISTRY=harbor.calypsoai.app
ENVEOF

# Set permissions for the application directory
sudo chmod -R 777 /opt/calypsoai-installer

# Clean up
echo "Cleaning up..."
sudo apt-get autoremove -y
sudo apt-get autoclean

echo "=========================================="
echo "F5 AI Security Installer Setup Complete!"
echo "=========================================="
echo ""
echo "Available tools:"
echo "- Pulumi: $(pulumi version)"
echo "- Helm: $(helm version --short)"
echo "- kubectl: $(kubectl version)"
echo "- AWS CLI: $(aws --version)"
echo "- Azure CLI: $(az version --output table)"
echo "- Google Cloud CLI: $(gcloud version --format='value(basic.version)' 2>/dev/null || echo 'Not available')"
echo "- OCI CLI: $(oci --version 2>/dev/null || echo 'Not available')"
echo ""
echo "To run the installer:"
echo "  calypsoai-installer"
echo ""
echo "Or navigate to the installer directory:"
echo "  cd /opt/calypsoai-installer"
echo ""
echo "Remember to configure your cloud provider credentials:"
echo "- AWS: aws configure"
echo "- Azure: az login"
echo "- GCP: gcloud auth login && gcloud auth application-default login"
echo ""
echo "For Harbor registry access:"
echo "  helm registry login harbor.calypsoai.app"
echo "=========================================="

echo "Setup script completed successfully!"
SETUP_SCRIPT_EOF
}

# Function to copy certificate files to remote VM
copy_certificate_files() {
    local ssh_command="$1"
    local provider="$2"
    local cert_dir="/opt/cert_files"
    
    echo "Copying certificate files to VM..."
    
    # Check if certificate files are specified and exist
    if [ -n "$CERT_FILEPATH" ] && [ -f "$CERT_FILEPATH" ]; then
        echo "Copying certificate file: $CERT_FILEPATH"
        if [ "$provider" = "gcp" ]; then
            gcloud compute scp --project="$GCP_PROJECT" --zone="$GCP_ZONE" "$CERT_FILEPATH" "$GCP_INSTANCE_NAME":$cert_dir/cert.crt
        else
            scp -o StrictHostKeyChecking=no $ssh_command "$CERT_FILEPATH" :$cert_dir/cert.crt
        fi
    fi
    
    if [ -n "$CERT_KEY_FILEPATH" ] && [ -f "$CERT_KEY_FILEPATH" ]; then
        echo "Copying certificate key file: $CERT_KEY_FILEPATH"
        if [ "$provider" = "gcp" ]; then
            gcloud compute scp --project="$GCP_PROJECT" --zone="$GCP_ZONE" "$CERT_KEY_FILEPATH" "$GCP_INSTANCE_NAME":$cert_dir/cert.key
        else
            scp -o StrictHostKeyChecking=no $ssh_command "$CERT_KEY_FILEPATH" :$cert_dir/cert.key
        fi
    fi
    
    if [ -n "$CERT_CHAIN_FILEPATH" ] && [ -f "$CERT_CHAIN_FILEPATH" ]; then
        echo "Copying certificate chain file: $CERT_CHAIN_FILEPATH"
        if [ "$provider" = "gcp" ]; then
            gcloud compute scp --project="$GCP_PROJECT" --zone="$GCP_ZONE" "$CERT_CHAIN_FILEPATH" "$GCP_INSTANCE_NAME":$cert_dir/cert-chain.crt
        else
            scp -o StrictHostKeyChecking=no $ssh_command "$CERT_CHAIN_FILEPATH" :$cert_dir/cert-chain.crt
        fi
    fi
    
    # Set proper permissions on the certificate files
    if [ "$provider" = "gcp" ]; then
        gcloud compute ssh --project="$GCP_PROJECT" --zone="$GCP_ZONE" "$GCP_INSTANCE_NAME" --command="sudo chmod 644 $cert_dir/*.crt 2>/dev/null || true; sudo chmod 600 $cert_dir/*.key 2>/dev/null || true"
    else
        local ssh_login=$(echo "$ssh_command" | sed 's/^-i [^ ]* //')
        ssh -o StrictHostKeyChecking=no $ssh_command "sudo chmod 644 $cert_dir/*.crt 2>/dev/null || true; sudo chmod 600 $cert_dir/*.key 2>/dev/null || true"
    fi
    
    echo "Certificate files copied successfully!"
}


case "$cloud_provider" in
  aws)
    echo "Initializing instance launch on AWS..."
    

    echo "Verifying AWS credentials..."
    AWS_CMD="aws"
    if [ -n "$AWS_PROFILE" ]; then
        AWS_CMD="aws --profile $AWS_PROFILE"
        echo "Using AWS profile: $AWS_PROFILE"
    fi
    
    if ! $AWS_CMD sts get-caller-identity > /dev/null 2>&1; then
        echo "Error: AWS credentials are not valid or have expired."
        if [ -n "$AWS_PROFILE" ]; then
            echo "If using AWS SSO, please run: aws sso login --profile $AWS_PROFILE"
        else
            echo "If using AWS SSO, please run: aws sso login"
        fi
        echo "If using regular credentials, please run: aws configure"
        exit 1
    fi
    

    echo "AWS credentials verified"
    echo "Current AWS identity:"
    $AWS_CMD sts get-caller-identity --query 'Account' --output text | xargs -I {} echo "  Account: {}"
    $AWS_CMD sts get-caller-identity --query 'Arn' --output text | xargs -I {} echo "  ARN: {}"
    
    if [ ! -f "$AWS_KEY_PAIR_NAME.pem" ]; then
        echo "Error: Private key file '$AWS_KEY_PAIR_NAME.pem' not found."
        echo "It must be in the same directory with permissions set to 400."
        exit 1
    fi

    INSTANCE_ID=$($AWS_CMD ec2 run-instances \
        --region "$AWS_REGION" \
        --image-id "$AWS_IMAGE_ID" \
        --instance-type "$AWS_INSTANCE_TYPE" \
        --key-name "$AWS_KEY_PAIR_NAME" \
        --security-group-ids "$AWS_SECURITY_GROUP_ID" \
        --block-device-mappings "[{\"DeviceName\":\"/dev/sda1\",\"Ebs\":{\"VolumeSize\":$AWS_ROOT_VOLUME_SIZE,\"VolumeType\":\"gp3\",\"DeleteOnTermination\":true}}]" \
        --tag-specifications "ResourceType=instance,Tags=[{Key=Name,Value=$AWS_INSTANCE_NAME}]" \
        --query "Instances[0].InstanceId" --output text)

    if [ -z "$INSTANCE_ID" ]; then
        echo "Failed to launch EC2 instance. Check AWS CLI configuration and parameters."
        exit 1
    fi
    echo "Instance launch initiated. Instance ID: $INSTANCE_ID"

    echo "Waiting for instance to enter 'running' state..."
    $AWS_CMD ec2 wait instance-running --region "$AWS_REGION" --instance-ids "$INSTANCE_ID"
    echo "Instance is now running."
    


    PUBLIC_IP=$($AWS_CMD ec2 describe-instances --region "$AWS_REGION" --instance-ids "$INSTANCE_ID" --query "Reservations[0].Instances[0].PublicIpAddress" --output text)

    if [ -z "$PUBLIC_IP" ]; then
        echo "Failed to retrieve public IP. Check instance networking and security group rules."
        exit 1
    fi
    echo "Public IP Address: $PUBLIC_IP"

    echo "Waiting for instance to be ready for SSH connection..."
    sleep 10
    
    # Retry SSH connection with timeout
    echo "Attempting to connect to instance..."
    for attempt in {1..5}; do
        if timeout 30 ssh -o StrictHostKeyChecking=no -o ConnectTimeout=10 -i "$AWS_KEY_PAIR_NAME.pem" "$AWS_USER@$PUBLIC_IP" "echo 'SSH connection successful'" 2>/dev/null; then
            echo "SSH connection established on attempt $attempt"
            break
        else
            echo "SSH connection failed on attempt $attempt, waiting 10 seconds..."
            sleep 10
        fi
    done

    # Copy certificate files to VM if specified
    copy_certificate_files "-i $AWS_KEY_PAIR_NAME.pem $AWS_USER@$PUBLIC_IP" "aws"

    echo "Running setup script on VM..."
    generate_setup_script "aws" | sed "s/SETUP_USERNAME/$AWS_USER/g" | ssh -o StrictHostKeyChecking=no -i "$AWS_KEY_PAIR_NAME.pem" "$AWS_USER@$PUBLIC_IP" bash

    echo ""
    echo "=========================================="
    echo "AWS EC2 Instance Setup Complete!"
    echo "=========================================="
    echo ""
    echo "Your instance is ready! Connect to it using:"
    echo ""
    echo "  ssh -i $AWS_KEY_PAIR_NAME.pem $AWS_USER@$PUBLIC_IP"
    echo ""
    echo "Once connected, you can run the F5 AI Security installer:"
    echo "  calypsoai-installer"
    echo ""
    echo "Or navigate to the installer directory:"
    echo "  cd /opt/calypsoai-installer"
    echo ""
    echo "Remember to configure your cloud provider credentials:"
    echo "- AWS: aws configure"
    echo "- Azure: az login"
    echo "- GCP: gcloud auth login && gcloud auth application-default login"
    echo ""
    echo "For Harbor registry access:"
    echo "  helm registry login harbor.calypsoai.app"
    echo "=========================================="
    ;;

  azure)
    echo "Initializing VM launch on Azure..."
    if [ ! -f "$AZURE_SSH_PUBLIC_KEY_PATH" ]; then
        echo "Error: SSH public key file '$AZURE_SSH_PUBLIC_KEY_PATH' not found."
        exit 1
    fi
    
    if [ ! -f "$AZURE_SSH_PRIVATE_KEY_PATH" ]; then
        echo "Error: SSH private key file '$AZURE_SSH_PRIVATE_KEY_PATH' not found."
        echo "It must be in the same directory with permissions set to 400."
        exit 1
    fi

    echo "Creating VM... this may take a few minutes."
    PUBLIC_IP=$(az vm create \
        --resource-group "$AZURE_RESOURCE_GROUP" \
        --name "$AZURE_VM_NAME" \
        --image "$AZURE_IMAGE" \
        --size "$AZURE_SIZE" \
        --location "$AZURE_LOCATION" \
        --admin-username "$AZURE_ADMIN_USER" \
        --ssh-key-values "$AZURE_SSH_PUBLIC_KEY_PATH" \
        --os-disk-size-gb "$AZURE_DISK_SIZE" \
        --query "publicIpAddress" --output tsv)

    if [ -z "$PUBLIC_IP" ]; then
        echo "Failed to launch Azure VM. Check Azure CLI configuration and parameters."
        exit 1
    fi

    echo "VM created successfully. Public IP Address: $PUBLIC_IP"

    echo "Waiting for VM to be ready for SSH connection..."
    sleep 10
    
    # Retry SSH connection with timeout
    echo "Attempting to connect to VM..."
    for attempt in {1..5}; do
        if timeout 30 ssh -o StrictHostKeyChecking=no -o ConnectTimeout=10 -i "$AZURE_SSH_PRIVATE_KEY_PATH" "$AZURE_ADMIN_USER@$PUBLIC_IP" "echo 'SSH connection successful'" 2>/dev/null; then
            echo "SSH connection established on attempt $attempt"
            break
        else
            echo "SSH connection failed on attempt $attempt, waiting 10 seconds..."
            sleep 10
        fi
    done

    # Copy certificate files to VM if specified
    copy_certificate_files "-i $AZURE_SSH_PRIVATE_KEY_PATH $AZURE_ADMIN_USER@$PUBLIC_IP" "azure"

    echo "Running setup script on VM..."
    generate_setup_script "azure" | sed "s/SETUP_USERNAME/$AZURE_ADMIN_USER/g" | ssh -o StrictHostKeyChecking=no -i "$AZURE_SSH_PRIVATE_KEY_PATH" "$AZURE_ADMIN_USER@$PUBLIC_IP" bash

    echo ""
    echo "=========================================="
    echo "Azure VM Setup Complete!"
    echo "=========================================="
    echo ""
    echo "Your VM is ready! Connect to it using:"
    echo ""
    echo "  ssh -i $AZURE_SSH_PRIVATE_KEY_PATH $AZURE_ADMIN_USER@$PUBLIC_IP"
    echo ""
    echo "Once connected, you can run the F5 AI Security installer:"
    echo "  calypsoai-installer"
    echo ""
    echo "Or navigate to the installer directory:"
    echo "  cd /opt/calypsoai-installer"
    echo ""
    echo "Remember to configure your cloud provider credentials:"
    echo "- AWS: aws configure"
    echo "- Azure: az login"
    echo "- GCP: gcloud auth login && gcloud auth application-default login"
    echo ""
    echo "For Harbor registry access:"
    echo "  helm registry login harbor.calypsoai.app"
    echo "=========================================="
    ;;

  gcp)
    echo "Initializing instance launch on GCP..."

    echo "Creating instance... this may take a few minutes."
    gcloud compute instances create "$GCP_INSTANCE_NAME" \
        --project="$GCP_PROJECT" \
        --zone="$GCP_ZONE" \
        --machine-type="$GCP_MACHINE_TYPE" \
        --image="$GCP_IMAGE_NAME" \
        --image-project="$GCP_IMAGE_PROJECT" \
        --boot-disk-size="$GCP_BOOT_DISK_SIZE"GB \
        --quiet

    if [ $? -ne 0 ]; then
        echo "Failed to launch GCP instance. Check gcloud CLI configuration and parameters."
        exit 1
    fi
    echo "Instance created successfully."

    PUBLIC_IP=$(gcloud compute instances describe "$GCP_INSTANCE_NAME" --project="$GCP_PROJECT" --zone="$GCP_ZONE" --format='value(networkInterfaces[0].accessConfigs[0].natIP)')
    echo "Public IP Address: $PUBLIC_IP"

    echo "Waiting for instance to be ready for SSH connection..."
    sleep 10
    
    # Retry gcloud SSH connection with timeout
    echo "Attempting to connect to instance..."
    for attempt in {1..5}; do
        if timeout 30 gcloud compute ssh --project="$GCP_PROJECT" --zone "$GCP_ZONE" "$GCP_INSTANCE_NAME" --command="echo 'SSH connection successful'" 2>/dev/null; then
            echo "SSH connection established on attempt $attempt"
            break
        else
            echo "SSH connection failed on attempt $attempt, waiting 10 seconds..."
            sleep 10
        fi
    done

    # Copy certificate files to VM if specified
    copy_certificate_files "" "gcp"

    echo "Running setup script on VM..."
    generate_setup_script "gcp" | sed "s/SETUP_USERNAME/ubuntu/g" | gcloud compute ssh --project="$GCP_PROJECT" --zone "$GCP_ZONE" "$GCP_INSTANCE_NAME" --command="bash -s"

    echo ""
    echo "=========================================="
    echo "GCP Compute Instance Setup Complete!"
    echo "=========================================="
    echo ""
    echo "Your instance is ready! Connect to it using:"
    echo ""
    echo "  gcloud compute ssh --project=$GCP_PROJECT --zone $GCP_ZONE $GCP_INSTANCE_NAME"
    echo ""
    echo "Once connected, you can run the F5 AI Security installer:"
    echo "  calypsoai-installer"
    echo ""
    echo "Or navigate to the installer directory:"
    echo "  cd /opt/calypsoai-installer"
    echo ""
    echo "Remember to configure your cloud provider credentials:"
    echo "- AWS: aws configure"
    echo "- Azure: az login"
    echo "- GCP: gcloud auth login && gcloud auth application-default login"
    echo ""
    echo "For Harbor registry access:"
    echo "  helm registry login harbor.calypsoai.app"
    echo "=========================================="
    ;;

  *)
    echo "Invalid cloud provider: '$cloud_provider'. Please choose 'aws', 'azure', or 'gcp'."
    exit 1
    ;;
esac

echo "Script finished."