#!/bin/bash
#
# This script packages your project into a solution distributable that can be
# used as an input to the solution builder validation pipeline.
#
# Important notes and prereq's:
#   1. This script should be run from the repo's /deployment folder.
#
# This script will perform the following tasks:
#   1. Remove any old dist files from previous runs.
#   2. Install dependencies for the cdk-solution-helper; responsible for
#      converting standard 'cdk synth' output into solution assets.
#   3. Build and synthesize your CDK project.
#   4. Run the cdk-solution-helper on template outputs and organize
#      those outputs into the /global-s3-assets folder.
#   5. Organize source code artifacts into the /regional-s3-assets folder.
#   6. Remove any temporary files used for staging.
#
# Parameters:
#  - source-bucket-base-name: Name for the S3 bucket location where the template will source the Lambda
#    code from. The template will append '-[region_name]' to this bucket name.
#    For example: ./build-s3-dist.sh solutions v1.0.0
#    The template will then expect the source code to be located in the solutions-[region_name] bucket
#  - solution-name: name of the solution for consistency
#  - version-code: version of the package

[ "$DEBUG" == 'true' ] && set -x
set -e

# Important: CDK global version number
cdk_version=1.126.0

# Check to see if the required parameters have been provided:
if [ -z "$1" ] || [ -z "$2" ] || [ -z "$3" ]; then
    echo "Please provide the base source bucket name, trademark approved solution name and version where the lambda code will eventually reside."
    echo "For example: ./build-s3-dist.sh solutions trademarked-solution-name v1.0.0"
    exit 1
fi

# Get reference for all important folders
template_dir="$PWD"
staging_dist_dir="$template_dir/staging"
template_dist_dir="$template_dir/global-s3-assets"
build_dist_dir="$template_dir/regional-s3-assets"
source_dir="$template_dir/../source"

echo "------------------------------------------------------------------------------"
echo "[Init] Remove any old dist files from previous runs"
echo "------------------------------------------------------------------------------"

echo "rm -rf $template_dist_dir"
rm -rf $template_dist_dir
echo "mkdir -p $template_dist_dir"
mkdir -p $template_dist_dir
echo "rm -rf $build_dist_dir"
rm -rf $build_dist_dir
echo "mkdir -p $build_dist_dir"
mkdir -p $build_dist_dir
echo "rm -rf $staging_dist_dir"
rm -rf $staging_dist_dir
echo "mkdir -p $staging_dist_dir"
mkdir -p $staging_dist_dir

echo "------------------------------------------------------------------------------"
echo "[Init] Setup the Python environment"
echo "------------------------------------------------------------------------------"

echo "cd $source_dir"
cd $source_dir

# setup lambda layers (building sagemaker layer using lambda build environment for python 3.8)
echo 'docker run --entrypoint /bin/bash -v "$source_dir"/lib/blueprints/byom/lambdas/sagemaker_layer:/var/task public.ecr.aws/lambda/python:3.9 -c "cat requirements.txt; pip3 install -r requirements.txt -t ./python; exit"'
docker run --entrypoint /bin/bash -v "$source_dir"/lib/blueprints/byom/lambdas/sagemaker_layer:/var/task public.ecr.aws/lambda/python:3.9 -c "cat requirements.txt; pip3 install -r requirements.txt -t ./python; exit"

# Remove tests and cache stuff (to reduce size)
find "$source_dir"/lib/blueprints/byom/lambdas/sagemaker_layer/python -type d -name "tests" -exec rm -rfv {} +
find "$source_dir"/lib/blueprints/byom/lambdas/sagemaker_layer/python -type d -name "__pycache__" -exec rm -rfv {} +

echo "python3 -m venv .venv-prod"
python3 -m venv .venv-prod
echo "source .venv-prod/bin/activate"
source .venv-prod/bin/activate
echo "pip install -r requirements.txt"
pip install -r requirements.txt

# setup crhelper for custom resource (copying s3 assets)
echo "pip install -r ./lambdas/custom_resource/requirements.txt -t ./lambdas/custom_resource/"
pip install -r ./lambdas/custom_resource/requirements.txt -t ./lambdas/custom_resource/

# setup crhelper for custom resource (solution helper)
echo "pip install -r ./lambdas/solution_helper/requirements.txt -t ./lambdas/solution_helper/"
pip install -r ./lambdas/solution_helper/requirements.txt -t ./lambdas/solution_helper/

# setup crhelper for invoke lambda custom resource
echo "pip install -r ./lib/blueprints/byom/lambdas/invoke_lambda_custom_resource/requirements.txt -t ./lib/blueprints/byom/lambdas/invoke_lambda_custom_resource/"
pip install -r ./lib/blueprints/byom/lambdas/invoke_lambda_custom_resource/requirements.txt -t ./lib/blueprints/byom/lambdas/invoke_lambda_custom_resource/

echo "------------------------------------------------------------------------------"
echo "[Init] Install dependencies for the cdk-solution-helper"
echo "------------------------------------------------------------------------------"

echo "cd $template_dir/cdk-solution-helper"
cd $template_dir/cdk-solution-helper
echo "npm ci --only=prod"
npm ci --only=prod

echo "------------------------------------------------------------------------------"
echo "[Synth] CDK Project"
echo "------------------------------------------------------------------------------"

# Install the global aws-cdk package
echo "cd $source_dir"
cd $source_dir
echo "npm install -g aws-cdk@$cdk_version"
npm install -g aws-cdk@$cdk_version

#Run 'cdk synth for BYOM blueprints
echo "cdk synth DataQualityModelMonitorStack > lib/blueprints/byom/byom_data_quality_monitor.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false"
cdk synth DataQualityModelMonitorStack > lib/blueprints/byom/byom_data_quality_monitor.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false
echo "cdk synth ModelQualityModelMonitorStack > lib/blueprints/byom/byom_model_quality_monitor.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false"
cdk synth ModelQualityModelMonitorStack > lib/blueprints/byom/byom_model_quality_monitor.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false
echo "cdk synth ModelBiasModelMonitorStack > lib/blueprints/byom/byom_model_bias_monitor.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false"
cdk synth ModelBiasModelMonitorStack > lib/blueprints/byom/byom_model_bias_monitor.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false
echo "cdk synth ModelExplainabilityModelMonitorStack > lib/blueprints/byom/byom_model_explainability_monitor.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false"
cdk synth ModelExplainabilityModelMonitorStack > lib/blueprints/byom/byom_model_explainability_monitor.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false
echo "cdk synth SingleAccountCodePipelineStack > lib/blueprints/byom/single_account_codepipeline.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false"
cdk synth SingleAccountCodePipelineStack > lib/blueprints/byom/single_account_codepipeline.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false
echo "cdk synth MultiAccountCodePipelineStack > lib/blueprints/byom/multi_account_codepipeline.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false"
cdk synth MultiAccountCodePipelineStack > lib/blueprints/byom/multi_account_codepipeline.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false
echo "cdk synth BYOMRealtimePipelineStack > lib/blueprints/byom/byom_realtime_inference_pipeline.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false"
cdk synth BYOMRealtimePipelineStack > lib/blueprints/byom/byom_realtime_inference_pipeline.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false
echo "cdk synth BYOMCustomAlgorithmImageBuilderStack > lib/blueprints/byom/byom_custom_algorithm_image_builder.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false"
cdk synth BYOMCustomAlgorithmImageBuilderStack > lib/blueprints/byom/byom_custom_algorithm_image_builder.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false
echo "cdk synth BYOMBatchStack > lib/blueprints/byom/byom_batch_pipeline.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false"
cdk synth BYOMBatchStack > lib/blueprints/byom/byom_batch_pipeline.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false
echo "cdk synth AutopilotJobStack > lib/blueprints/byom/autopilot_training_pipeline.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false"
cdk synth AutopilotJobStack > lib/blueprints/byom/autopilot_training_pipeline.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false
echo "cdk synth TrainingJobStack > lib/blueprints/byom/model_training_pipeline.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false"
cdk synth TrainingJobStack > lib/blueprints/byom/model_training_pipeline.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false
echo "cdk synth HyperparamaterTunningJobStack > lib/blueprints/byom/model_hyperparameter_tunning_pipeline.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false"
cdk synth HyperparamaterTunningJobStack > lib/blueprints/byom/model_hyperparameter_tunning_pipeline.yaml --path-metadata false --version-reporting false --generate-bootstrap-version-rule false

# Replace %%VERSION%% in other templates
replace="s/%%VERSION%%/$3/g"
echo "sed -i -e $replace lib/blueprints/byom/byom_data_quality_monitor.yaml"
sed -i -e $replace lib/blueprints/byom/byom_data_quality_monitor.yaml
echo "sed -i -e $replace lib/blueprints/byom/byom_model_quality_monitor.yaml"
sed -i -e $replace lib/blueprints/byom/byom_model_quality_monitor.yaml
echo "sed -i -e $replace lib/blueprints/byom/byom_model_bias_monitor.yaml"
sed -i -e $replace lib/blueprints/byom/byom_model_bias_monitor.yaml
echo "sed -i -e $replace lib/blueprints/byom/byom_model_explainability_monitor.yaml"
sed -i -e $replace lib/blueprints/byom/byom_model_explainability_monitor.yaml
echo "sed -i -e $replace lib/blueprints/byom/byom_realtime_inference_pipeline.yaml"
sed -i -e $replace lib/blueprints/byom/byom_realtime_inference_pipeline.yaml
echo "sed -i -e $replace lib/blueprints/byom/single_account_codepipeline.yaml"
sed -i -e $replace lib/blueprints/byom/single_account_codepipeline.yaml
echo "sed -i -e $replace lib/blueprints/byom/multi_account_codepipeline.yaml"
sed -i -e $replace lib/blueprints/byom/multi_account_codepipeline.yaml
echo "sed -i -e $replace lib/blueprints/byom/byom_custom_algorithm_image_builder.yaml"
sed -i -e $replace lib/blueprints/byom/byom_custom_algorithm_image_builder.yaml
echo "sed -i -e $replace lib/blueprints/byom/byom_batch_pipeline.yaml"
sed -i -e $replace lib/blueprints/byom/byom_batch_pipeline.yaml
echo "sed -i -e $replace lib/blueprints/byom/autopilot_training_pipeline.yaml"
sed -i -e $replace lib/blueprints/byom/autopilot_training_pipeline.yaml
echo "sed -i -e $replace lib/blueprints/byom/model_training_pipeline.yaml"
sed -i -e $replace lib/blueprints/byom/model_training_pipeline.yaml
echo "sed -i -e $replace lib/blueprints/byom/model_hyperparameter_tunning_pipeline.yaml"
sed -i -e $replace lib/blueprints/byom/model_hyperparameter_tunning_pipeline.yaml

# Run 'cdk synth' for main templates to generate raw solution outputs
echo "cdk synth mlops-workload-orchestrator-single-account --path-metadata false --version-reporting false --generate-bootstrap-version-rule false --output=$staging_dist_dir"
cdk synth mlops-workload-orchestrator-single-account --path-metadata false --version-reporting false --generate-bootstrap-version-rule false --output=$staging_dist_dir
echo "cdk synth mlops-workload-orchestrator-multi-account --path-metadata false --version-reporting false --generate-bootstrap-version-rule false --output=$staging_dist_dir"
cdk synth mlops-workload-orchestrator-multi-account --path-metadata false --version-reporting false --generate-bootstrap-version-rule false --output=$staging_dist_dir

# Remove unnecessary output files
echo "cd $staging_dist_dir"
cd $staging_dist_dir
echo "rm tree.json manifest.json cdk.out"
rm tree.json manifest.json cdk.out

echo "------------------------------------------------------------------------------"
echo "[Packing] Template artifacts"
echo "------------------------------------------------------------------------------"

# Move outputs from staging to template_dist_dir
echo "Move outputs from staging to template_dist_dir"
echo "cp $template_dir/*.template $template_dist_dir/"
cp $staging_dist_dir/*.template.json $template_dist_dir/
rm *.template.json

# Rename all *.template.json files to *.template
echo "Rename all *.template.json to *.template"
echo "copy templates and rename"
for f in $template_dist_dir/*.template.json; do
    mv -- "$f" "${f%.template.json}.template"
done

# Run the helper to clean-up the templates and remove unnecessary CDK elements
echo "Run the helper to clean-up the templates and remove unnecessary CDK elements"
echo "node $template_dir/cdk-solution-helper/index"
node $template_dir/cdk-solution-helper/index
if [ "$?" = "1" ]; then
	echo "(cdk-solution-helper) ERROR: there is likely output above." 1>&2
	exit 1
fi

# Find and replace bucket_name, solution_name, and version
echo "Find and replace bucket_name, solution_name, and version"
cd $template_dist_dir
echo "Updating code source bucket in template with $1"
replace="s/%%BUCKET_NAME%%/$1/g"

echo "sed -i -e $replace $template_dist_dir/mlops-workload-orchestrator-single-account.template"
sed -i -e $replace $template_dist_dir/mlops-workload-orchestrator-single-account.template
echo "sed -i -e $replace $template_dist_dir/mlops-workload-orchestrator-multi-account.template"
sed -i -e $replace $template_dist_dir/mlops-workload-orchestrator-multi-account.template
replace="s/%%SOLUTION_NAME%%/$2/g"
echo "sed -i -e $replace $template_dist_dir/mlops-workload-orchestrator-single-account"
sed -i -e $replace $template_dist_dir/mlops-workload-orchestrator-single-account.template
echo "sed -i -e $replace $template_dist_dir/mlops-workload-orchestrator-multi-account.template"
sed -i -e $replace $template_dist_dir/mlops-workload-orchestrator-multi-account.template
replace="s/%%VERSION%%/$3/g"
echo "sed -i -e $replace $template_dist_dir/mlops-workload-orchestrator-single-account.template"
sed -i -e $replace $template_dist_dir/mlops-workload-orchestrator-single-account.template
echo "sed -i -e $replace $template_dist_dir/mlops-workload-orchestrator-multi-account.template"
sed -i -e $replace $template_dist_dir/mlops-workload-orchestrator-multi-account.template


echo "------------------------------------------------------------------------------"
echo "[Packing] Source code artifacts"
echo "------------------------------------------------------------------------------"
cd $staging_dist_dir
# ... For each asset.* source code artifact in the temporary /staging folder...
for d in `find . -mindepth 1 -maxdepth 1 -type d`; do

    # Rename the artifact, removing the period for handler compatibility
    pfname="$(basename -- $d)"
    fname="$(echo $pfname | sed -e 's/\.//g')"
    mv $d $fname

    # Zip the artifact
    cd $fname
    echo "zip -qr ../$fname.zip *"
    zip -qr ../$fname.zip *
    cd ..

    # Copy the zipped artifact from /staging to /regional-s3-assets
    echo "cp $fname.zip $build_dist_dir"
    cp $fname.zip $build_dist_dir

    # Remove the old, unzipped artifact from /staging
    echo "rm -rf $fname"
    rm -rf $fname

    # Remove the old, zipped artifact from /staging
    echo "rm $fname.zip"
    rm $fname.zip

    # ... repeat until all source code artifacts are zipped and placed in the
    # ... /regional-s3-assets folder

done

echo "Creating zip files for the blueprint stacks"
echo "mkdir -p $template_dir/bp_staging"
mkdir -p $template_dir/bp_staging
echo "cp -r $source_dir/lib/blueprints $template_dir/bp_staging/"
cp -r $source_dir/lib/blueprints $template_dir/bp_staging/

echo "cp -r $source_dir/lambdas/pipeline_orchestration/shared $template_dir/bp_staging/"
cp -r $source_dir/lambdas/pipeline_orchestration/shared $template_dir/bp_staging/

# Cleaning up blueprint staging folder
cd $template_dir/bp_staging/
rm -rf **/__pycache__/
rm -rf **/*.egg-info/
cd $template_dir/bp_staging/blueprints

# Loop through all blueprint directories in blueprints
for bp in `find . -mindepth 1 -maxdepth 1 -type d`; do

    cd $bp
    # Loop through all subdirectories in blueprints/<blueprint_type> e.g., byom
    for d in `find . -mindepth 1 -maxdepth 1 -type d`; do
        if [ $d != "./lambdas" ]; then
            # Remove any directory that is not 'lambdas'
            rm -rf $d
        fi
    done

    cd lambdas
    # Loop through all lambda directories of the <blueprint_type>
    for lambda in `find . -mindepth 1 -maxdepth 1 -type d`; do

        # Copying shared source codes to each lambda function
        echo "cp -r $template_dir/bp_staging/shared $lambda"
        cp -r $template_dir/bp_staging/shared $lambda

        # Removing './' from the directory name to use for zip file
        echo "lambda_dir_name=`echo $lambda | cut -d '/' -f 2`"
        lambda_dir_name=`echo $lambda | cut -d '/' -f 2`

        cd $lambda_dir_name

        # Creating the zip file for each lambda
        echo "zip -r9 ../$lambda_dir_name.zip *"
        zip -r9 ../$lambda_dir_name.zip *
        cd ..

        # Removing the lambda directories after creating zip files of them
        echo "rm -rf $lambda"
        rm -rf $lambda
    done
done

cd $template_dir/bp_staging/blueprints
# Remove all .py files in subdirectories since they are not necessary anymore
rm -f *.py
rm -f */*.py
rm -f */*/*.py

cd $template_dir/bp_staging

# Creating a zip file of blueprints directory and putting it in regional-s3-assets
echo "zip -r9 blueprints.zip blueprints"
zip -r9 blueprints.zip blueprints

echo "mv blueprints.zip $build_dist_dir"
mv blueprints.zip $build_dist_dir

# Remove temporary directory
rm -rf $template_dir/bp_staging

echo "------------------------------------------------------------------------------"
echo "[Cleanup] Remove temporary files"
echo "------------------------------------------------------------------------------"

# Delete the temporary /staging folder
echo "rm -rf $staging_dist_dir"
rm -rf $staging_dist_dir