AWSTemplateFormatVersion: "2010-09-09"
Transform: AWS::Serverless-2016-10-31
Description: Serverless Computer Vision Label Detection (uksb-1reh3a0p6)
Parameters:
  RekognitionModelProjectARN:
    Type: String
    Description: Amazon Rekognition Model Project ARN
  RekognitionModelProjectVersionARN:
    Type: String
    Description: Amazon Rekognition Model Project Version ARN

Mappings:
  Solution:
    Constants:
      Version: "v0.13"

Resources:
  SourceS3Bucket:
    Type: AWS::S3::Bucket
    DeletionPolicy: Retain
    UpdateReplacePolicy: Retain
    Properties:
      BucketEncryption:
        ServerSideEncryptionConfiguration:
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: AES256
      NotificationConfiguration:
        QueueConfigurations:
          - Event: s3:ObjectCreated:*
            Queue: !GetAtt SQSQueue.Arn
            Filter:
              S3Key:
                Rules:
                  - Name: suffix
                    Value: .png
          - Event: s3:ObjectCreated:*
            Queue: !GetAtt SQSQueue.Arn
            Filter:
              S3Key:
                Rules:
                  - Name: suffix
                    Value: .jpg

  SQSQueue:
    Type: AWS::SQS::Queue
    Properties:
      DelaySeconds: 0
      MaximumMessageSize: 262144
      MessageRetentionPeriod: 345600
      ReceiveMessageWaitTimeSeconds: 0
      VisibilityTimeout: 60

  SQSQueuePolicy:
    Type: AWS::SQS::QueuePolicy
    Properties:
      PolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Sid: "TestSID"
            Effect: "Allow"
            Principal:
              Service:
                - "lambda.amazonaws.com"
                - "s3.amazonaws.com"
            Action:
              - "sqs:DeleteMessage"
              - "sqs:ReceiveMessage"
              - "sqs:SendMessage"
            Resource: !Sub "arn:aws:sqs:${AWS::Region}:${AWS::AccountId}:${SQSQueue.QueueName}"
            Condition:
              StringEquals:
                "aws:SourceAccount": !Ref AWS::AccountId
      Queues:
        - !Sub "https://sqs.${AWS::Region}.amazonaws.com/${AWS::AccountId}/${SQSQueue.QueueName}"

  CustomCVStateMachine:
    Type: AWS::Serverless::StateMachine
    Properties:
      Definition:
        StartAt: Check SQS Queue
        States:
          Check SQS Queue:
            Type: Task
            Resource: !GetAtt SQSPollerFunction.Arn
            ResultPath: $.messageinqueue
            Next: Are there images to process?
          Are there images to process?:
            Type: Choice
            Choices:
              - Variable: $.messageinqueue
                StringEquals: incoming
                Next: Start Model
            Default: Finish
          Start Model:
            Type: Task
            Resource: !GetAtt StartModelFunction.Arn
            ResultPath: "$.runningstatus"
            Next: Start States
          Start States:
            Type: Choice
            Choices:
              - Variable: "$.runningstatus"
                StringEquals: RUNNING
                Next: Enable SQS Trigger
            Default: Wait for the model to start
          Wait for the model to start:
            Type: Wait
            Seconds: 900
            Next: Start Model
          Keep Model running for 1 hr:
            Type: Wait
            Seconds: 3540
            Next: Check Queue Again
          Check Queue Again:
            Type: Task
            Resource: !GetAtt SQSPollerFunction.Arn
            ResultPath: $.moremessagesinqueue
            Next: Are there more images?
          Are there more images?:
            Type: Choice
            Choices:
              - Variable: $.moremessagesinqueue
                StringEquals: stop
                Next: Disable SQS Trigger
            Default: Keep Model running for 1 hr
          Enable SQS Trigger:
            Type: Task
            ResultPath: $.alreadyrunning
            Parameters:
              - Action: enable
            Resource: !GetAtt ToggleTriggerFunction.Arn
            Next: Is another machine already running?
          Is another machine already running?:
            Type: Choice
            Choices:
                - Variable: $.alreadyrunning
                  StringEquals: Already_Running
                  Next: Finish
            Default: Keep Model running for 1 hr
          Disable SQS Trigger:
            Type: Task
            Parameters:
              - Action: disable
            Resource: !GetAtt ToggleTriggerFunction.Arn
            Next: Stop Model
          Stop Model:
            Type: Task
            Resource: !GetAtt StopModelFunction.Arn
            Next: Finish
          Finish:
            Type: Succeed
      Events:
        HourlyPollingSchedule:
          Type: Schedule # More info about Schedule Event Source: https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/sam-property-statemachine-schedule.html
          Properties:
            Description: Schedule to run the state machine every 1 hour
            Enabled: True # This schedule can be disabled based on the use case to avoid incurring charges.
            Schedule: "rate(1 hour)"
      Policies: # Find out more about SAM policy templates: https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/serverless-policy-templates.html
        - LambdaInvokePolicy:
            FunctionName: !Ref SQSPollerFunction
        - LambdaInvokePolicy:
            FunctionName: !Ref StartModelFunction
        - LambdaInvokePolicy:
            FunctionName: !Ref StopModelFunction
        - LambdaInvokePolicy:
            FunctionName: !Ref ToggleTriggerFunction

  SQSPollerFunction:
    Type: AWS::Serverless::Function # More info about Function Resource: https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/sam-resource-function.html
    Properties:
      Description: "Lambda function to Poll SQS queue for incoming messages from S3"
      CodeUri: functions/sqs_poller/
      Handler: app.lambda_handler
      Environment:
        Variables:
          SQS_Queue_URL: !Ref SQSQueue
      Runtime: python3.8
      MemorySize: 128
      Policies:
        - SQSPollerPolicy:
            QueueName: !GetAtt SQSQueue.QueueName
        - LambdaInvokePolicy:
            FunctionName: !Ref StartModelFunction
        - LambdaInvokePolicy:
            FunctionName: !Ref StopModelFunction

  StartModelFunction:
    Type: AWS::Serverless::Function # More info about Function Resource: https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/sam-resource-function.html
    Properties:
      Description: "Lambda function to Start the Rekognition model"
      CodeUri: functions/start_model/
      Handler: app.lambda_handler
      Environment:
        Variables:
          rekog_model_project_version_arn: !Ref RekognitionModelProjectVersionARN
          rekog_model_project_arn: !Ref RekognitionModelProjectARN
      Runtime: python3.8
      MemorySize: 128
      Timeout: 3
      Policies:
        - Statement:
            - Action:
                - "rekognition:DescribeProjectVersions"
                - "rekognition:StartProjectVersion"
              Effect: "Allow"
              Resource:
                - !Ref RekognitionModelProjectARN
                - !Ref RekognitionModelProjectVersionARN
        - LambdaInvokePolicy:
            FunctionName: !Ref ToggleTriggerFunction

  StopModelFunction:
    Type: AWS::Serverless::Function # More info about Function Resource: https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/sam-resource-function.html
    Properties:
      Description: "Lambda function to Stop the Rekognition model"
      CodeUri: functions/stop_model/
      Handler: app.lambda_handler
      Environment:
        Variables:
          rekog_model_project_version_arn: !Ref RekognitionModelProjectVersionARN
          rekog_model_project_arn: !Ref RekognitionModelProjectARN
      Runtime: python3.8
      MemorySize: 128
      Timeout: 3
      Policies:
        - Statement:
            - Action:
                - "rekognition:DescribeProjectVersions"
                - "rekognition:StopProjectVersion"
              Effect: "Allow"
              Resource:
                - !Ref RekognitionModelProjectARN
                - !Ref RekognitionModelProjectVersionARN
        - LambdaInvokePolicy:
            FunctionName: !Ref ToggleTriggerFunction

  ToggleTriggerFunction:
    Type: AWS::Serverless::Function # More info about Function Resource: https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/sam-resource-function.html
    Properties:
      Description: "Lambda function to toggle the SQS trigger for analysis"
      CodeUri: functions/toggle_trigger/
      Handler: app.lambda_handler
      Environment:
        Variables:
          analyze_lambda_arn: !GetAtt AnalyseImageFunction.Arn
      Runtime: python3.8
      MemorySize: 128
      Timeout: 3
      Policies:
        - Statement:
            - Action:
                - "lambda:ListEventSourceMappings"
                - "lambda:GetEventSourceMapping"
                - "lambda:UpdateEventSourceMapping"
              Effect: "Allow"
              Resource: "*"

  FinalS3Bucket:
    Type: "AWS::S3::Bucket"
    DeletionPolicy: Retain
    UpdateReplacePolicy: Retain
    Properties:
      BucketEncryption:
        ServerSideEncryptionConfiguration:
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: AES256

  AnalyseImageFunction:
    Type: AWS::Serverless::Function # More info about Function Resource: https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/sam-resource-function.html
    Properties:
      Description: "Lambda function to analyze Images using the Rekognition Model"
      CodeUri: functions/analyse_image/
      Handler: app.lambda_handler
      Environment:
        Variables:
          rekognition_model_project_version_arn: !Ref RekognitionModelProjectVersionARN
          Final_S3_Bucket_Name: !Ref FinalS3Bucket
      Events:
        SQSEvent:
          Type: SQS
          Properties:
            Queue: !GetAtt SQSQueue.Arn
            BatchSize: 1
            Enabled: false
      MemorySize: 128
      Timeout: 60
      Runtime: python3.8
      Policies:
        - Statement:
            - Action:
                - "rekognition:DetectCustomLabels"
              Effect: "Allow"
              Resource:
                - !Ref RekognitionModelProjectARN
                - !Ref RekognitionModelProjectVersionARN
        - S3CrudPolicy:
            BucketName: !Ref FinalS3Bucket
        - S3CrudPolicy:
            BucketName: !Ref SourceS3Bucket
        - SQSPollerPolicy:
            QueueName: !GetAtt SQSQueue.QueueName

Outputs:
  # CustomCVStateMachineHourlySchedule is an implicit Schedule event rule created out of Events key under Serverless::StateMachine
  # Find out more about other implicit resources you can reference within SAM
  # https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/sam-specification-generated-resources.html
  SourceS3BucketName:
    Description: "Name of the S3 bucket to hold the incoming images"
    Value: !Ref SourceS3Bucket
  FinalS3BucketName:
    Description: "Name of the final S3 bucket to hold the image and the inference json"
    Value: !Ref FinalS3Bucket
  CustomCVStateMachineARN:
    Description: "ARN of the Step Function"
    Value: !Ref CustomCVStateMachine
  CustomCVStateMachineHourlyPollingScheduleARN:
    Description: "ARN of the AWS EventBridge Rule"
    Value: !Ref CustomCVStateMachineHourlyPollingSchedule