#!/bin/bash

# SSO profile name
SSO_PROFILE_NAME="suno"

# Lambda function name and region
DEPLOYMENT_TYPE="${1:-staging}"  # Use first argument, default to staging if not set
if [ "$DEPLOYMENT_TYPE" = "staging" ]; then
    FUNCTION_NAME="ask-suno-staging"
elif [ "$DEPLOYMENT_TYPE" = "prod" ]; then
    FUNCTION_NAME="ask-suno"
else
    echo "Invalid DEPLOYMENT_TYPE. Must be 'staging', or 'prod'"
    exit 1
fi
REGION="us-east-1"

# Lambda Target Environment
PYTHON_VERSION="3.13"
TARGET_PLATFORM="manylinux_2_28_aarch64" # Match Lambda AL2023 arm64
IMPLEMENTATION="cp"

# Ensure we have valid credentials
echo "Checking AWS credentials..."
if ! aws sts get-caller-identity --profile $SSO_PROFILE_NAME &> /dev/null; then
  echo "AWS credentials are invalid or expired. Logging in..."
  chmod +x ./config_sso.sh
  ./config_sso.sh
  
  # Check if login was successful
  if ! aws sts get-caller-identity --profile $SSO_PROFILE_NAME &> /dev/null; then
    echo "SSO login failed. Please check your SSO configuration and try again."
    exit 1
  fi
else
  echo "Using existing valid AWS credentials"
fi

# Export the profile for all AWS operations
export AWS_PROFILE=$SSO_PROFILE_NAME

# Create directories
echo "Creating directories..."
rm -rf lambda_package lambda_wheels
mkdir -p lambda_package lambda_wheels

# Check if constants.py exists
if [ ! -f "constants.py" ]; then
  echo "Error: constants.py is missing. This file is required for deployment."
  exit 1
fi

# Download wheels for the target Lambda platform
echo "Downloading wheels for Lambda (${TARGET_PLATFORM}, Python ${PYTHON_VERSION})..."
pip download \
  -r requirements.in \
  --dest lambda_wheels \
  --platform $TARGET_PLATFORM \
  --python-version $PYTHON_VERSION \
  --implementation $IMPLEMENTATION \
  --only-binary=:all:

# Install downloaded wheels into the package directory
echo "Installing downloaded wheels into package directory..."
pip install \
  --target lambda_package \
  --no-index \
  --find-links lambda_wheels \
  -r requirements.in

# Copy lambda function code and local modules
echo "Copying lambda function code..."
cp lambda_function.py error_models.py constants.py instrumental_determinator.py suno.py anodyne.py lambda_package/

# Apply .uvignore patterns (or similar cleanup logic if needed)
# Note: .uvignore is specific to uv. You might need different cleanup steps.
echo "Cleaning up package (basic)..."
# Add specific cleanup rules here if needed, similar to the previous script
find lambda_package -type d -name "__pycache__" -exec rm -rf {} \; 2>/dev/null || true
find lambda_package -type f -name "*.pyc" -delete
find lambda_package -type f -name "*.pyo" -delete
find lambda_package -type d -name "tests" -exec rm -rf {} \; 2>/dev/null || true

# Create deployment package
echo "Creating deployment package..."
cd lambda_package && zip -r9 ../lambda_deployment_package.zip . && cd ..

# Deploy to AWS Lambda
echo "Deploying to AWS Lambda..."
aws lambda update-function-code \
  --function-name $FUNCTION_NAME \
  --zip-file fileb://lambda_deployment_package.zip \
  --region $REGION | cat

echo "Deployment complete!"

# Clean up wheel directory
rm -rf lambda_wheels