#!/bin/bash

# Record start time
START_TIME=$(date +%s)

# Function to check list of required environment variables
check_env_vars() {
    local missing_vars=()
    for var in "$@"; do
        if [ -z "${!var}" ]; then
            missing_vars+=("$var")
        fi
    done
    if [ ${#missing_vars[@]} -ne 0 ]; then
        echo "Error: The following environment variables are not set: ${missing_vars[*]}"
        exit 1
    fi
}

# Function to check if doctl is available
check_doctl() {
    if ! command -v doctl &> /dev/null; then
        echo "Error: doctl is not installed or not in PATH"
        echo "Please install doctl: https://docs.digitalocean.com/reference/doctl/how-to/install/"
        exit 1
    fi
}

# Function to setup SSH private key
setup_ssh_key() {
    if [ -z "$SSH_PRIVATE_KEY_DECODED" ]; then
        echo "Error: SSH_PRIVATE_KEY_DECODED environment variable is not set"
        exit 1
    fi

    # Create SSH directory if it doesn't exist
    mkdir -p ~/.ssh
    chmod 700 ~/.ssh

    # Decode base64 encoded SSH private key
    echo "$SSH_PRIVATE_KEY_DECODED" | base64 -d > ~/.ssh/id_rsa

    # Ensure proper line endings for OpenSSH key format
    sed -i 's/\r$//' ~/.ssh/id_rsa

    # Add final newline if missing
    if [ -s ~/.ssh/id_rsa ] && [ "$(tail -c1 ~/.ssh/id_rsa | wc -l)" -eq 0 ]; then
        echo >> ~/.ssh/id_rsa
    fi

    chmod 600 ~/.ssh/id_rsa

    # Validate SSH key formats
    if ! ssh-keygen -y -f ~/.ssh/id_rsa > /dev/null 2>&1; then
        echo "Error: Invalid SSH private key format"
        echo "Key content preview (first 100 chars):"
        head -c 100 ~/.ssh/id_rsa
        echo
        echo "Key file size: $(wc -c < ~/.ssh/id_rsa) bytes"
        echo "Key file lines: $(wc -l < ~/.ssh/id_rsa) lines"
        exit 1
    fi

    echo "SSH private key has been set up and validated successfully"
}

# Function to destroy droplet
destroy_droplet() {
    local droplet_id="$1"

    if [ -z "$droplet_id" ]; then
        echo "Error: No droplet ID provided for destruction"
        return 1
    fi

    echo "Destroying droplet with ID: $droplet_id"

    # Authenticate doctl
    doctl auth init --access-token "$DO_TOKEN"

    # Destroy droplet
    if doctl compute droplet delete "$droplet_id" --force; then
        echo "Droplet $droplet_id destroyed successfully"
        return 0
    else
        echo "Error: Failed to destroy droplet $droplet_id"
        return 1
    fi
}

# Function to create droplet
create_droplet() {
    local name="${1:-runner-${RUNNER_ID}}"
    local size="${2:-s-1vcpu-1gb}"
    local image="${3:-200895606}"
    local region="${4:-sgp1}"

    echo "Creating droplet: $name"
    echo "Size: $size, Image: $image, Region: $region"

    # Authenticate doctl
    doctl auth init --access-token "$DO_TOKEN"

    # Create droplet
    local droplet_id=$(doctl compute droplet create "$name" \
        --size "$size" \
        --image "$image" \
        --region "$region" \
        --ssh-keys $(doctl compute ssh-key list --format ID --no-header | tr '\n' ',' | sed 's/,$//' || echo "") \
        --format ID \
        --no-header \
        --wait)

    if [ -n "$droplet_id" ]; then
        echo "Droplet created successfully with ID: $droplet_id"
        echo "Droplet name: $name"

        echo "droplet_id=$droplet_id" >> $GITHUB_OUTPUT

        # Get droplet IP
        local droplet_ip=$(doctl compute droplet get "$droplet_id" --format PublicIPv4 --no-header)
        echo "Droplet IP: $droplet_ip"
        echo "droplet_ip=$droplet_ip" >> $GITHUB_OUTPUT

        # Test SSH connection
        if test_ssh_connection "$droplet_ip"; then
            echo "Droplet is ready and accessible via SSH"

            # Setup VM with required services
            if setup_vm "$droplet_ip"; then
                echo "VM setup completed successfully"
            else
                echo "Error: VM setup failed"
                echo "Destroying failed droplet..."
                destroy_droplet "$droplet_id"
                return 1
            fi
        else
            echo "Warning: Droplet created but SSH connection failed"
            echo "Destroying failed droplet..."
            destroy_droplet "$droplet_id"
            return 1
        fi

        return 0
    else
        echo "Error: Failed to create droplet"
        return 1
    fi
}

# Function to test SSH connection
test_ssh_connection() {
    local droplet_ip="$1"
    local max_attempts=10
    local attempt=1
    local wait_time=10

    echo "Testing SSH connection to $droplet_ip..."

    while [ $attempt -le $max_attempts ]; do
        echo "Attempt $attempt/$max_attempts: Testing SSH connection..."

        # Test SSH connection with verbose error handling
        ssh_output=$(ssh -i ~/.ssh/id_rsa \
            -o StrictHostKeyChecking=no \
            -o UserKnownHostsFile=/dev/null \
            -o ConnectTimeout=10 \
            -o BatchMode=yes \
            -o PasswordAuthentication=no \
            root@"$droplet_ip" "echo 'SSH connection successful'" 2>&1)

        ssh_exit_code=$?

        if [ $ssh_exit_code -eq 0 ]; then
            echo "SSH connection to $droplet_ip established successfully"
            echo "Response: $ssh_output"
            return 0
        else
            echo "SSH connection failed (exit code: $ssh_exit_code)"
            echo "Error output: $ssh_output"

            if [ $attempt -lt $max_attempts ]; then
                echo "Waiting $wait_time seconds before retry..."
                sleep $wait_time
            fi
            attempt=$((attempt + 1))
        fi
    done

    echo "Error: Failed to establish SSH connection to $droplet_ip after $max_attempts attempts"
    return 1
}

# Function to setup VM after successful SSH connection
setup_vm() {
    local droplet_ip="$1"

    if [ -z "$droplet_ip" ]; then
        echo "Error: No droplet IP provided for VM setup"
        return 1
    fi

    echo "Setting up VM at $droplet_ip..."

    # Execute setup commands via SSH
    echo "Step 1: Running prepare.sh config.yml..."
    if ! ssh -i ~/.ssh/id_rsa \
        -o StrictHostKeyChecking=no \
        -o UserKnownHostsFile=/dev/null \
        -o ConnectTimeout=30 \
        root@"$droplet_ip" \
        "./prepare.sh config.yml RUNNER_ID $RUNNER_ID"; then
        echo "Error: Failed to run prepare.sh config.yml"
        return 1
    fi

    echo "Step 2: Running prepare.sh docker-compose.yml..."
    if ! ssh -i ~/.ssh/id_rsa \
        -o StrictHostKeyChecking=no \
        -o UserKnownHostsFile=/dev/null \
        -o ConnectTimeout=30 \
        root@"$droplet_ip" \
        "./prepare.sh docker-compose.yml RUNNER_ID $RUNNER_ID"; then
        echo "Error: Failed to run prepare.sh docker-compose.yml"
        return 1
    fi

    echo "Step 3: Starting docker compose services..."
    if ! ssh -i ~/.ssh/id_rsa \
        -o StrictHostKeyChecking=no \
        -o UserKnownHostsFile=/dev/null \
        -o ConnectTimeout=30 \
        root@"$droplet_ip" \
        "docker compose up -d"; then
        echo "Error: Failed to start docker compose services"
        return 1
    fi

    echo "Step 4: Checking docker compose status..."
    ssh -i ~/.ssh/id_rsa \
        -o StrictHostKeyChecking=no \
        -o UserKnownHostsFile=/dev/null \
        -o ConnectTimeout=30 \
        root@"$droplet_ip" \
        "docker compose ps"

    echo "VM setup completed successfully"
    return 0
}

# Check required environment variables
check_env_vars "SPEC" "IMAGE" "RUNNER_ID" "DO_TOKEN" "SSH_PRIVATE_KEY_DECODED"

# Check if doctl is available
check_doctl

# Setup SSH private key
setup_ssh_key

create_droplet "$RUNNER_ID" "$SPEC" "$IMAGE"

# Calculate and display execution time
END_TIME=$(date +%s)
EXECUTION_TIME=$((END_TIME - START_TIME))
echo "Script execution time: ${EXECUTION_TIME} seconds"
echo "ready=true" >> $GITHUB_OUTPUT