AWSTemplateFormatVersion: "2010-09-09"
Description: Private Amazon Braket environment - Braket VPC 
Parameters:
    VPCName:
      Description: "Braket VPC"
      Type: String
      Default: "Private VPC containing Braket and S3 private endpoints"

Mappings:
  BraketSubnetConfig:
    VPC:
      CIDR: "10.10.0.0/16"
    Private1:
      CIDR: "10.10.1.0/24"
    Private2:
      CIDR: "10.10.2.0/24"
        
  # This mapping accounts for the scenario when certain AZs
  # are not available to use (this differs on a per account
  # per customer basis). E.g., if the 'b' AZ is not available
  # in a specific region in one's account then updating the
  # list contained in the mapping below here will allow a
  # different AZ to be chosen. Since this template builds a
  # private endpoint for Amazon Braket, only regions where
  # the service exists are allowed.

  BraketAZRegions:
    us-east-1:            # IonQ
      AZs: ["a", "b"]
    us-west-1:            # Rigetti
      AZs: ["b", "c"]
    us-west-2:            # DWave
      AZs: ["a", "b", "c"]

Resources:

  VPC:
    Type: "AWS::EC2::VPC"
    Properties:
      EnableDnsSupport: "true"
      EnableDnsHostnames: "true"
      CidrBlock:
        Fn::FindInMap:
          - "BraketSubnetConfig"
          - "VPC"
          - "CIDR"
      Tags:
        - Key: Name
          Value: "Amazon Braket"
  
  UpdateDefaultSecurityGroup:
    Type: AWS::EC2::SecurityGroupIngress
    Properties: 
      GroupId:
        Fn::GetAtt: [VPC, DefaultSecurityGroup]
      CidrIp: 10.0.0.0/8
      IpProtocol: tcp
      FromPort: 443
      ToPort: 443
      Description: "Allows traffic between VPC and Braket notebook."

  PrivateSubnet1:
    Type: "AWS::EC2::Subnet"
    Properties:
      VpcId:
        Ref: "VPC"
      AvailabilityZone:
        Fn::Sub:
          - "${AWS::Region}${AZ}"
          - AZ: !Select [ 0, !FindInMap [ "BraketAZRegions", !Ref "AWS::Region", "AZs" ] ]
      CidrBlock:
        Fn::FindInMap:
          - "BraketSubnetConfig"
          - "Private1"
          - "CIDR"

  PrivateSubnet2:
    Type: "AWS::EC2::Subnet"
    Properties:
      VpcId:
        Ref: "VPC"
      AvailabilityZone:
        Fn::Sub:
          - "${AWS::Region}${AZ}"
          - AZ: !Select [ 1, !FindInMap [ "BraketAZRegions", !Ref "AWS::Region", "AZs" ] ]
      CidrBlock:
        Fn::FindInMap:
          - "BraketSubnetConfig"
          - "Private2"
          - "CIDR"

  RouteTable:
    Type: "AWS::EC2::RouteTable"
    Properties: 
      VpcId: 
        Ref: "VPC"

  PrivateSubnetRouteTableAssociation1:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      SubnetId:
        Ref: PrivateSubnet1
      RouteTableId:
        Ref: RouteTable

  PrivateSubnetRouteTableAssociation2:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      SubnetId:
        Ref: PrivateSubnet2
      RouteTableId:
        Ref: RouteTable

  BraketEndpoint:
    Type: 'AWS::EC2::VPCEndpoint'
    Properties:
      PrivateDnsEnabled: 'false'
      ServiceName: !Sub 'com.amazonaws.${AWS::Region}.braket'
      SubnetIds:
        - !Ref PrivateSubnet1
        - !Ref PrivateSubnet2
      VpcEndpointType: 'Interface'
      VpcId: !Ref VPC

  S3Endpoint:
    Type: 'AWS::EC2::VPCEndpoint'
    Properties:
      PrivateDnsEnabled: 'false'
      ServiceName: !Sub 'com.amazonaws.${AWS::Region}.s3'
      VpcEndpointType: 'Gateway'
      VpcId: !Ref VPC

  HostedZone: 
    Type: "AWS::Route53::HostedZone"
    Properties: 
      HostedZoneConfig: 
        Comment: 'Hosted zone for private inter-region DNS resolution'
      Name: !Sub 'braket.${AWS::Region}.amazonaws.com'
      VPCs: 
        - 
          VPCId: !Ref VPC
          VPCRegion: !Ref "AWS::Region"

  RecordSet: 
    Type: AWS::Route53::RecordSet
    Properties: 
      AliasTarget:
          DNSName: !Select [1,  !Split [ ":", !Select [0, !GetAtt BraketEndpoint.DnsEntries] ] ]
          HostedZoneId: !Select [0,  !Split [ ":", !Select [0, !GetAtt BraketEndpoint.DnsEntries] ] ]
      Comment: 'Record set pointing to Braket endpoint'
      HostedZoneId: !Ref HostedZone
      Name: !Sub 'braket.${AWS::Region}.amazonaws.com'
      Type: A

Outputs:
  VPCId:
    Description: "Amazon Braket VPC"
    Value: !Ref VPC

  RouteTable:
    Description: "Route table ID that needs to be updated with the VPC peering connection ID"
    Value: !Ref RouteTable

  S3Endpoint:
    Description: "S3 private endpoint to store experiment results."
    Value: !Ref S3Endpoint

  BraketEndpoint: 
    Description: "Amazon Braket private endpoint"
    Value: !Ref BraketEndpoint

  HostedZone:
    Description: "Hosted zone required to allow inter-regional DNS resolution."
    Value: !Ref HostedZone