#!/bin/bash

# Script to automatically update SSH keys in all main.tf files
# This script looks for terraform_key (or terraform_key.pub) in ~/.ssh directory

# Function to display usage
usage() {
    echo "Usage: $0 [ssh_key_name]"
    echo "  ssh_key_name: Name of the SSH key pair (default: terraform_key)"
    echo ""
    echo "Example:"
    echo "  $0                    # Uses default 'terraform_key'"
    echo "  $0 my_custom_key      # Uses 'my_custom_key' and 'my_custom_key.pub'"
    exit 1
}

# Set the SSH key name (default to terraform_key)
SSH_KEY_NAME="${1:-terraform_key}"

# Expand the home directory properly  
HOME_DIR="${HOME:-/home/$(whoami)}"
SSH_KEY_PATH="$HOME_DIR/.ssh/$SSH_KEY_NAME"
SSH_KEY_PUB_PATH="$HOME_DIR/.ssh/$SSH_KEY_NAME.pub"

# Check if SSH key exists
if [ ! -f "$SSH_KEY_PATH" ] && [ ! -f "$SSH_KEY_PUB_PATH" ]; then
    echo "Error: SSH key '$SSH_KEY_NAME' not found in $HOME_DIR/.ssh/"
    echo "Please generate your SSH key first:"
    echo "  ssh-keygen -t rsa -b 4096 -f $HOME_DIR/.ssh/$SSH_KEY_NAME"
    exit 1
fi

# Check if public key exists specifically (required for reading)
if [ ! -f "$SSH_KEY_PUB_PATH" ]; then
    echo "Error: SSH public key '$SSH_KEY_NAME.pub' not found in $HOME_DIR/.ssh/"
    exit 1
fi

# Get the public key content (remove any trailing whitespace)
PUBLIC_KEY=$(cat "$SSH_KEY_PUB_PATH" | tr -d '\n')

# Validate that we got a valid SSH key
if [[ ! "$PUBLIC_KEY" =~ ^ssh-[a-z]+[[:space:]]+[A-Za-z0-9+/]*[=]{0,3} ]]; then
    echo "Error: Invalid SSH public key format detected"
    exit 1
fi

echo "Found SSH public key:"
echo "$PUBLIC_KEY"
echo ""

# Find all main.tf files and update them
MAIN_TF_FILES=$(find . -name "main.tf" -type f)

if [ -z "$MAIN_TF_FILES" ]; then
    echo "No main.tf files found!"
    exit 1
fi

echo "Updating SSH key in the following files:"
echo "$MAIN_TF_FILES"
echo ""

# Replace the ssh_key line in all main.tf files using # as delimiter
echo "Replacing SSH key in all main.tf files..."
for file in $MAIN_TF_FILES; do
    sed -i "s#ssh_key = \".*\"#ssh_key = \"$PUBLIC_KEY\"#g" "$file"
done

# Verify the replacement worked
echo ""
echo "Verification:"
for file in $MAIN_TF_FILES; do
    echo "File: $file"
    grep "ssh_key =" "$file" | head -1
done

echo ""
echo "SSH key has been successfully updated in all main.tf files!"
echo "Backup files are saved with timestamp suffixes."