#!/bin/bash

# Check if region parameter is provided
if [ $# -eq 0 ]; then
    echo "Usage: $0 <region>"
    echo "Example: $0 us-east-1"
    echo "         $0 us-west-2"
    exit 1
fi

REGION=$1

# SSO Configuration
SSO_START_URL="https://sunoai.awsapps.com/start/#"
SSO_REGION="us-east-1"
SSO_ACCOUNT_ID="734185074900"
SSO_ROLE_NAME="PowerUserAccess"

# Create profile name based on region (e.g., suno-east-1, suno-west-2)
REGION_SHORT=$(echo $REGION | sed 's/us-//' | sed 's/-[0-9]$//')
REGION_NUM=$(echo $REGION | grep -o '[0-9]$')
SSO_PROFILE_NAME="suno-${REGION_SHORT}-${REGION_NUM}"

echo "Using region: $REGION"
echo "Profile name: $SSO_PROFILE_NAME"

# Check if AWS CLI is configured properly
if ! command -v aws &> /dev/null; then
    echo "AWS CLI is not installed. Please install it first."
    exit 1
fi

# Configure AWS CLI for SSO
echo "Configuring AWS CLI for SSO..."
aws configure set sso_start_url $SSO_START_URL
aws configure set sso_region $SSO_REGION
aws configure set sso_account_id $SSO_ACCOUNT_ID
aws configure set sso_role_name $SSO_ROLE_NAME
aws configure set region $REGION
aws configure set output json

# Set up named profile
echo "Setting up named profile '$SSO_PROFILE_NAME'..."
aws configure set profile.$SSO_PROFILE_NAME.sso_start_url $SSO_START_URL
aws configure set profile.$SSO_PROFILE_NAME.sso_region $SSO_REGION
aws configure set profile.$SSO_PROFILE_NAME.sso_account_id $SSO_ACCOUNT_ID
aws configure set profile.$SSO_PROFILE_NAME.sso_role_name $SSO_ROLE_NAME
aws configure set profile.$SSO_PROFILE_NAME.region $REGION
aws configure set profile.$SSO_PROFILE_NAME.output json

# Log in to SSO
echo "Logging in to AWS SSO..."
aws sso login --profile $SSO_PROFILE_NAME

# Check if login was successful
if [ $? -eq 0 ]; then
    echo "SSO login successful!"
    echo ""
    echo "Setting $SSO_PROFILE_NAME as the default profile..."
    export AWS_PROFILE=$SSO_PROFILE_NAME
    echo "export AWS_PROFILE=$SSO_PROFILE_NAME" >> ~/.bashrc
    echo "export AWS_PROFILE=$SSO_PROFILE_NAME" >> ~/.zshrc
    echo ""
    echo "Profile $SSO_PROFILE_NAME is now set as the default profile."
    echo "This profile will be used for all subsequent AWS CLI commands."
    echo ""
    echo "Example: aws lambda list-functions"
    echo ""
    echo "Credentials will typically be valid for 8 hours."
    echo "To use this profile in new terminal sessions, restart your terminal or run:"
    echo "source ~/.bashrc (or ~/.zshrc)"
else
    echo "SSO login failed. Please check your SSO configuration and try again."
fi