#!/bin/bash
# Update AWS credentials in secrets.toml for Snowflake Streamlit Container Runtime
# This script assumes awsume is configured and prompts for MFA token
# Preserves existing Spotify credentials and only updates AWS section

set -e

SECRETS_FILE=".streamlit/secrets.toml"

echo "=== AWS Secrets Updater for Snowflake ==="
echo ""

# Check if MFA token was passed as argument
if [ -n "$1" ]; then
    mfa_token="$1"
else
    # Prompt for MFA token
    read -p "Enter your MFA token: " mfa_token
fi

if [ -z "$mfa_token" ]; then
    echo "Error: MFA token cannot be empty"
    exit 1
fi

echo ""
echo "Running awsume with MFA token..."

# Run awsume with the MFA token
# Note: awsume may show warnings about not being sourced, but -o default still works
awsume songwhip -o default --mfa-token "$mfa_token" || true

# Wait a moment for credentials to be written
sleep 1

# Extract credentials
echo "Extracting AWS credentials..."
aws_access_key_id=$(aws configure get aws_access_key_id)
aws_secret_access_key=$(aws configure get aws_secret_access_key)
aws_session_token=$(aws configure get aws_session_token)

# Verify we got all credentials
if [ -z "$aws_access_key_id" ] || [ -z "$aws_secret_access_key" ] || [ -z "$aws_session_token" ]; then
    echo "Error: Failed to retrieve AWS credentials"
    exit 1
fi

echo "AWS credentials retrieved successfully."
echo ""

# Check if secrets.toml exists
if [ -f "$SECRETS_FILE" ]; then
    echo "Updating existing $SECRETS_FILE..."

    # Create a temp file
    tmp_file=$(mktemp)

    # Read the file and update AWS credentials while preserving other content
    awk -v access_key="$aws_access_key_id" \
        -v secret_key="$aws_secret_access_key" \
        -v session_token="$aws_session_token" \
        -v date="$(date)" '
    BEGIN { in_aws_section = 0; aws_updated = 0 }

    # Detect AWS section by looking for AWS_ keys
    /^AWS_ACCESS_KEY_ID[[:space:]]*=/ {
        print "# AWS credentials updated: " date
        print "AWS_ACCESS_KEY_ID = \"" access_key "\""
        in_aws_section = 1
        aws_updated = 1
        next
    }
    /^AWS_SECRET_ACCESS_KEY[[:space:]]*=/ { next }
    /^AWS_SESSION_TOKEN[[:space:]]*=/ { next }
    /^AWS_REGION[[:space:]]*=/ {
        print "AWS_SECRET_ACCESS_KEY = \"" secret_key "\""
        print "AWS_SESSION_TOKEN = \"" session_token "\""
        print "AWS_REGION = \"us-east-1\""
        in_aws_section = 0
        next
    }

    # Skip comment lines that mention AWS credentials timing
    /^# AWS credentials updated:/ { next }
    /^# Note: These are temporary session credentials/ { next }
    /^# Regenerate using:/ { next }

    # Print all other lines
    { print }

    END {
        # If AWS section was not found, append it
        if (aws_updated == 0) {
            print ""
            print "# AWS Credentials (for DynamoDB access)"
            print "# AWS credentials updated: " date
            print "# Note: These are temporary session credentials and will expire"
            print "# Regenerate using: ./generate-secrets.sh"
            print "AWS_ACCESS_KEY_ID = \"" access_key "\""
            print "AWS_SECRET_ACCESS_KEY = \"" secret_key "\""
            print "AWS_SESSION_TOKEN = \"" session_token "\""
            print "AWS_REGION = \"us-east-1\""
        }
    }
    ' "$SECRETS_FILE" > "$tmp_file"

    # Replace original file
    mv "$tmp_file" "$SECRETS_FILE"

    echo "Updated $SECRETS_FILE with new AWS credentials."
else
    echo "Creating new $SECRETS_FILE..."

    # Ensure .streamlit directory exists
    mkdir -p .streamlit

    # Create new secrets.toml with placeholders for Spotify
    cat > "$SECRETS_FILE" << EOF
# Heavy Rotation App - Secrets Configuration
# Generated: $(date)

# Spotify API Credentials
# Get these from https://developer.spotify.com/dashboard
SPOTIFY_CLIENT_ID = "your_client_id_here"
SPOTIFY_CLIENT_SECRET = "your_client_secret_here"

# AWS Credentials (for DynamoDB access)
# Note: These are temporary session credentials and will expire
# Regenerate using: ./generate-secrets.sh
AWS_ACCESS_KEY_ID = "$aws_access_key_id"
AWS_SECRET_ACCESS_KEY = "$aws_secret_access_key"
AWS_SESSION_TOKEN = "$aws_session_token"
AWS_REGION = "us-east-1"
EOF

    echo "Created $SECRETS_FILE"
    echo ""
    echo "NOTE: Remember to add your Spotify credentials to the file!"
fi

echo ""
echo "Done!"