AWSTemplateFormatVersion: '2010-09-09'
Description: Prefect 2 Cluster, Roles, and Service Infrastructure

Parameters:
  envUpper:
    Type: String
    Description: "The name of the environment to be used as a part of the prefix for all resource names"
    AllowedValues: ["DEV", "PROD"]
  environment:
    Type: String
    Description: "The name of the environment to be used as a part of the prefix for all resource names"
    AllowedValues: [ "dev", "prod" ]
  cpu:
    Type: String
    Description: Allow Dynamic CPU configuration
    Default: 512
    AllowedValues: [256, 512, 1024, 2048, 4096]
  memory:
    Type: String
    Description: Allow Increasing Memory - from 8192 on requires 4096 CPU and increases in 1024 increments default
    Default: 1024
    AllowedValues: [512, 1024, 2048, 4096, 5120, 6144, 7168, 8192 , 9216, 10240]
  awsaccountid:
    Type: String
    Description: AWS Account ID
    Default: xxxx
  region:
    Type: String
    Description: AWS region name
    Default: us-east-2
  image:
    Type: String
    Description: Docker image from prefect which the agent will use - NOT OUR PY310-base image
    Default: prefecthq/prefect:2.14.13-python3.10
  InstanceType:
    Type: String
    Description: Choosing t3 micro because it is free
    Default: t3.micro
  PrefectECSServiceSubnet:
    Type: String
    Description: Subnet ID
    Default: subnet-0056ad26401f64636
  ECSClusterName:
    Type: String
    Default: dev-orion-prefect-2
  SecurityGroupIDs:
    Type: CommaDelimitedList
    Default: sg-336ce346
  SubnetIDs:
    Type: CommaDelimitedList
    Description: Subnets in the main VPC to use
    Default: >-
      subnet-abc123,subnet-abc234,subnet-abc345
  VpcID:
    Type: String
    Description: Main VPC ID
    Default: vpc-123


#Dict to reference
Mappings:
  EnvConfig:
    dev:
        PrefectApiUrl: "arn:aws:secretsmanager:us-east-2:xxxx:secret:DEV_PREFECT:api_url::"
        PrefectApiKey: "arn:aws:secretsmanager:us-east-2:xxxx:secret:DEV_PREFECT:api_key::"
    prod:
        PrefectApiUrl: "arn:aws:secretsmanager:us-east-2:xxxx:secret:PROD_PREFECT:api_url::"
        PrefectApiKey: "arn:aws:secretsmanager:us-east-2:xxxx:secret:PROD_PREFECT:api_key::"


Resources:
  PrefectFargateCluster:
    Type: AWS::ECS::Cluster
    Properties:
      ClusterName: !Sub "prefect-2-cluster-${environment}"
  PrefectLogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Ref PrefectFargateCluster
      RetentionInDays: 7
  ExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Sub "${environment}_prefect_ecs_execution_role"
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service: ecs-tasks.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: SecretsFromParameterStoreCloudWatch
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - ssm:GetParameters
                  - logs:CreateLogStream
                  - logs:CreateLogGroup
                  - logs:PutLogEvents
                  - secretsmanager:GetSecretValue
                  - kms:Decrypt
                Resource: "*"
        - PolicyName: SecretsManagerDecryptPermission
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - ssm:GetParameters
                  - secretsmanager:GetSecretValue
                  - kms:Decrypt
                Resource:
                  - !Sub "arn:aws:secretsmanager:us-east-2:xxxx:secret:${envUpper}_PREFECT*"
                  - "arn:aws:secretsmanager:us-east-2:xxxx:secret:Docker"
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AmazonECSTaskExecutionRolePolicy
  TaskRole:
    Type: AWS::IAM::Role
    Properties:
      #Change to lower?
      RoleName: !Sub "${environment}_prefect_ecs_task_role"
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service: ecs-tasks.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: PrefectRequirements
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - s3:*
                  - ecr:BatchCheckLayerAvailability
                  - ecr:BatchGetImage
                  - ecr:GetDownloadUrlForLayer
                  - ecr:GetAuthorizationToken
                  - ecs:RegisterTaskDefinition
                  - ecs:DeregisterTaskDefinition
                  - ecs:DescribeTaskDefinition
                  - ecs:DescribeTasks
                  - ecs:RunTask
                  - logs:GetLogEvents
                  - ec2:DescribeSubnets
                  - ec2:DescribeVpcs
                  - ecs:StopTask
                  - iam:PassRole
                  - secretsmanager:GetSecretValue
                  - kms:Decrypt
                Resource: "*"

  PrefectTaskDefinition:
    Type: AWS::ECS::TaskDefinition
    Properties:
      Family: !Ref PrefectFargateCluster
      Cpu: !Ref cpu
      Memory: !Ref memory
      NetworkMode: awsvpc
      ExecutionRoleArn: !Ref ExecutionRole
      TaskRoleArn: !Ref TaskRole
      ContainerDefinitions:
        - Name: !Sub ${environment}-prefect-container
          Image: !Ref image
          EntryPoint:
            - "bash"
            - "-c"
          StopTimeout: 120
          Environment:
            - Name: PREFECT_LOGGING_LEVEL
              Value: DEBUG
            - Name: AWS_RETRY_MODE
              Value: adaptive
            - Name: AWS_MAX_ATTEMPTS
              Value: 99
          Command:
            - !Sub "pip install prefect-aws && prefect worker start -t ecs -p ecs-worker-pool"
          LogConfiguration:
            LogDriver: awslogs
            Options:
              awslogs-region: !Ref AWS::Region
              awslogs-group: !Ref PrefectLogGroup
              awslogs-stream-prefix: prefect
          Secrets:
            - Name: PREFECT_API_URL
              ValueFrom: !FindInMap [ EnvConfig, !Ref environment, PrefectApiUrl ]
            - Name: PREFECT_API_KEY
              ValueFrom: !FindInMap [ EnvConfig, !Ref environment, PrefectApiKey ]
      RequiresCompatibilities:
        - FARGATE


  ECSService:
    Type: AWS::ECS::Service
    Properties:
      Cluster: !Ref PrefectFargateCluster
      TaskDefinition: !Ref PrefectTaskDefinition
      LaunchType: FARGATE
      ServiceName: !Sub ${environment}-prefect-service
      SchedulingStrategy: REPLICA
      DesiredCount: 1
      NetworkConfiguration:
        AwsvpcConfiguration:
          AssignPublicIp: ENABLED
          SecurityGroups: !Ref SecurityGroupIDs
          Subnets: !Ref SubnetIDs
      EnableECSManagedTags: true