#!/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 [options] [ssh_key_name]"
    echo "  options:"
    echo "    -r, --remove   Remove SSH key from main.tf files"
    echo "    -h, --help     Display this help message"
    echo ""
    echo "  ssh_key_name: Name of the SSH key pair (default: terraform_key)"
    echo ""
    echo "Example:"
    echo "  $0                    # Updates with default 'terraform_key'"
    echo "  $0 my_custom_key      # Updates with 'my_custom_key'"
    echo "  $0 -r                 # Remove SSH key from files"
    echo "  $0 -r my_custom_key   # Remove SSH key from files"
    exit 1
}

# Parse command line arguments
REMOVE_KEY=false
SSH_KEY_NAME="terraform_key"

# Check if any arguments are provided
if [ $# -eq 0 ]; then
    # No arguments - use default behavior (update)
    :
elif [ "$1" = "-h" ] || [ "$1" = "--help" ]; then
    usage
elif [ "$1" = "-r" ] || [ "$1" = "--remove" ]; then
    # Remove mode enabled
    REMOVE_KEY=true
    if [ $# -gt 1 ]; then
        SSH_KEY_NAME="$2"
    fi
else
    # Normal update mode with key name provided as argument
    SSH_KEY_NAME="$1"
fi

# 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"

# If not removing keys, validate SSH key exists
if [ "$REMOVE_KEY" = false ]; then
    # 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 ""
fi

# 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 ""

# Process each file based on remove mode
for file in $MAIN_TF_FILES; do
    if [ "$REMOVE_KEY" = true ]; then
        echo "Removing SSH key from $file..."
        # Set ssh_key to empty string for idempotent removal
        sed -i "s/^[[:space:]]*ssh_key[[:space:]]*=[[:space:]]*\"[^\"]*\"/  ssh_key = \"\"/" "$file"
    else
        echo "Updating SSH key in $file..."
        # Update the ssh_key line with new value
        sed -i "s#ssh_key = \".*\"#ssh_key = \"$PUBLIC_KEY\"#g" "$file"
    fi
done

# Verify the replacement worked
echo ""
echo "Verification:"
for file in $MAIN_TF_FILES; do
    echo "File: $file"
    if [ "$REMOVE_KEY" = true ]; then
        # Show lines with empty ssh_key values
        grep "ssh_key = \"\"" "$file" | head -1
    else
        # Show updated ssh_key lines
        grep "ssh_key =" "$file" | head -1
    fi
done

echo ""
if [ "$REMOVE_KEY" = true ]; then
    echo "SSH key has been successfully removed (set to empty string) in all main.tf files!"
else
    echo "SSH key has been successfully updated in all main.tf files!"
fi