AWSTemplateFormatVersion: '2010-09-09'
Description: 'CloudFormation template for SageMaker Unified Studio with trusted identity propagation - Complete version with all required permissions'

Resources:
  # S3 Bucket for customer data
  CustomerDataBucket:
    Type: AWS::S3::Bucket
    Properties:
      BucketName: !Sub 'smustipidcsqldemo-${AWS::AccountId}-${AWS::Region}'
      BucketEncryption:
        ServerSideEncryptionConfiguration:
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: AES256
      PublicAccessBlockConfiguration:
        BlockPublicAcls: true
        BlockPublicPolicy: true
        IgnorePublicAcls: true
        RestrictPublicBuckets: true
      VersioningConfiguration:
        Status: Enabled
      Tags:
        - Key: Purpose
          Value: SageMaker-Unified-Studio-CustomerData
        - Key: Environment
          Value: Demo

  # IAM Policy for Redshift with Glue/Lake Formation permissions
  RedshiftBasicPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      ManagedPolicyName: RedshiftBasicPolicy
      Description: 'Policy for Redshift operations with Glue and Lake Formation access'
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Sid: RedshiftDescribePermissions
            Effect: Allow
            Action:
              - 'redshift:DescribeQev2IdcApplications'
              - 'redshift-serverless:ListNamespaces'
              - 'redshift-serverless:ListWorkgroups'
              - 'redshift-serverless:GetWorkgroup'
            Resource: '*'
          - Sid: GlueDataCatalogAccess
            Effect: Allow
            Action:
              - 'lakeformation:GetDataAccess'
              - 'glue:GetTable'
              - 'glue:GetTables'
              - 'glue:SearchTables'
              - 'glue:GetDatabase'
              - 'glue:GetDatabases'
              - 'glue:GetPartitions'
              - 'lakeformation:GetResourceLFTags'
              - 'lakeformation:ListLFTags'
              - 'lakeformation:GetLFTag'
              - 'lakeformation:SearchTablesByLFTags'
              - 'lakeformation:SearchDatabasesByLFTags'
            Resource: '*'

  # IAM Role for Redshift IAM Identity Center Integration
  IAMIDCRedshiftRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: IAMIDCRedshiftRole
      Description: 'Allows Redshift clusters to call AWS services with Glue/Lake Formation permissions'
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: redshift.amazonaws.com
            Action:
              - 'sts:AssumeRole'
              - 'sts:SetContext'
      ManagedPolicyArns:
        - !Ref RedshiftBasicPolicy
      Policies:
        - PolicyName: LakeFormationGlueAccess
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - lakeformation:GetDataAccess
                  - glue:GetTable
                  - glue:GetTables
                  - glue:SearchTables
                  - glue:GetDatabase
                  - glue:GetDatabases
                  - glue:GetPartitions
                  - lakeformation:GetResourceLFTags
                  - lakeformation:ListLFTags
                  - lakeformation:GetLFTag
                  - lakeformation:SearchTablesByLFTags
                  - lakeformation:SearchDatabasesByLFTags
                Resource: '*'
      Tags:
        - Key: Purpose
          Value: SageMaker-Unified-Studio-RedshiftIDC
        - Key: Environment
          Value: Demo

  # IAM Role for Data Transfer between Glue and Redshift
  DataTransferRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: DataTransferRole
      Description: 'Allows Glue and Redshift to call AWS services for data transfer operations'
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - glue.amazonaws.com
                - redshift.amazonaws.com
                - redshift-serverless.amazonaws.com
            Action: 'sts:AssumeRole'
      ManagedPolicyArns:
        - 'arn:aws:iam::aws:policy/service-role/AWSGlueServiceRole'
      Policies:
        - PolicyName: S3DataAccess
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - 's3:GetObject'
                  - 's3:PutObject'
                  - 's3:DeleteObject'
                  - 's3:ListBucket'
                Resource:
                  - !Sub '${CustomerDataBucket.Arn}/*'
                  - !GetAtt CustomerDataBucket.Arn
              - Effect: Allow
                Action:
                  - 'glue:GetDatabase'
                  - 'glue:GetDatabases'
                  - 'glue:GetTable'
                  - 'glue:GetTables'
                  - 'glue:CreateTable'
                  - 'glue:UpdateTable'
                  - 'glue:DeleteTable'
                  - 'glue:GetPartitions'
                  - 'glue:CreatePartition'
                  - 'glue:UpdatePartition'
                  - 'glue:DeletePartition'
                  - 'glue:BatchCreatePartition'
                  - 'glue:BatchUpdatePartition'
                  - 'glue:BatchDeletePartition'
                  - 'redshift:DescribeNamespaces'
                  - 'redshift:DescribeClusters'
                  - 'redshift:GetClusterCredentials'
                  - 'redshift:CreateClusterUser'
                  - 'redshift:JoinGroup'
                  - 'redshift:DescribeQev2IdcApplications'
                  - 'redshift:ModifyCluster'
                  - 'redshift:ModifyClusterIamRoles'
                  - 'redshift-data:ExecuteStatement'
                  - 'redshift-data:GetStatementResult'
                  - 'redshift-data:DescribeStatement'
                  - 'redshift-data:ListStatements'
                  - 'redshift-data:CancelStatement'
                Resource: '*'
      Tags:
        - Key: Purpose
          Value: SageMaker-Unified-Studio-DataTransfer
        - Key: Environment
          Value: Demo

  # IAM Policy for Lake Formation Registration
  LakeFormationPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      ManagedPolicyName: LakeFormationRegistrationPolicy
      Description: 'Policy for Lake Formation data lake operations'
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Sid: LakeFormationPermissions
            Effect: Allow
            Action:
              - 'lakeformation:RegisterResource'
              - 'lakeformation:DeregisterResource'
              - 'lakeformation:DescribeResource'
              - 'lakeformation:ListResources'
              - 'lakeformation:GrantPermissions'
              - 'lakeformation:RevokePermissions'
              - 'lakeformation:ListPermissions'
              - 'lakeformation:GetEffectivePermissionsForPath'
              - 'lakeformation:GetDataAccess'
              - 'lakeformation:GetTemporaryGlueTableCredentials'
              - 'lakeformation:GetTemporaryGluePartitionCredentials'
              - 'lakeformation:CreateLakeFormationOptIn'
              - 'lakeformation:DeleteLakeFormationOptIn'
              - 'lakeformation:ListLakeFormationOptIns'
              - 'lakeformation:CreateLFTag'
              - 'lakeformation:DeleteLFTag'
              - 'lakeformation:DescribeLFTag'
              - 'lakeformation:ListLFTags'
              - 'lakeformation:GetResourceLFTags'
              - 'lakeformation:AddLFTagsToResource'
              - 'lakeformation:RemoveLFTagsFromResource'
              - 'glue:GetDatabase'
              - 'glue:GetDatabases'
              - 'glue:CreateDatabase'
              - 'glue:UpdateDatabase'
              - 'glue:DeleteDatabase'
              - 'glue:GetTable'
              - 'glue:GetTables'
              - 'glue:CreateTable'
              - 'glue:UpdateTable'
              - 'glue:DeleteTable'
              - 'glue:GetPartition'
              - 'glue:GetPartitions'
              - 'glue:CreatePartition'
              - 'glue:UpdatePartition'
              - 'glue:DeletePartition'
              - 'glue:BatchCreatePartition'
              - 'glue:BatchDeletePartition'
              - 'glue:BatchUpdatePartition'
            Resource: '*'
          - Sid: S3DataLakeAccess
            Effect: Allow
            Action:
              - 's3:GetObject'
              - 's3:PutObject'
              - 's3:DeleteObject'
              - 's3:ListBucket'
              - 's3:GetBucketLocation'
              - 's3:ListAllMyBuckets'
            Resource:
              - !Sub '${CustomerDataBucket.Arn}/*'
              - !GetAtt CustomerDataBucket.Arn
              - 'arn:aws:s3:::*'
          - Sid: IAMPassRole
            Effect: Allow
            Action:
              - 'iam:PassRole'
            Resource: !Sub 'arn:aws:iam::${AWS::AccountId}:role/LakeFormationRegistrationRole'

  # IAM Role for Lake Formation Registration
  LakeFormationRegistrationRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: LakeFormationRegistrationRole
      Description: 'Allows Lake Formation and Redshift to call AWS Services for data lake operations'
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lakeformation.amazonaws.com
                - redshift.amazonaws.com
            Action:
              - 'sts:AssumeRole'
              - 'sts:SetContext'
      ManagedPolicyArns:
        - !Ref LakeFormationPolicy
      Tags:
        - Key: Purpose
          Value: SageMaker-Unified-Studio-LakeFormation
        - Key: Environment
          Value: Demo

Outputs:
  S3BucketName:
    Description: 'Name of the created S3 bucket for customer data'
    Value: !Ref CustomerDataBucket
    Export:
      Name: !Sub '${AWS::StackName}-S3Bucket'

  S3BucketArn:
    Description: 'ARN of the created S3 bucket'
    Value: !GetAtt CustomerDataBucket.Arn
    Export:
      Name: !Sub '${AWS::StackName}-S3BucketArn'

  IAMIDCRedshiftRoleArn:
    Description: 'ARN of the IAM Identity Center Redshift Role with Glue/Lake Formation permissions'
    Value: !GetAtt IAMIDCRedshiftRole.Arn
    Export:
      Name: !Sub '${AWS::StackName}-IAMIDCRedshiftRole'

  DataTransferRoleArn:
    Description: 'ARN of the Data Transfer Role'
    Value: !GetAtt DataTransferRole.Arn
    Export:
      Name: !Sub '${AWS::StackName}-DataTransferRole'

  LakeFormationRegistrationRoleArn:
    Description: 'ARN of the Lake Formation Registration Role'
    Value: !GetAtt LakeFormationRegistrationRole.Arn
    Export:
      Name: !Sub '${AWS::StackName}-LakeFormationRole'

  BucketLocation:
    Description: 'S3 location for Lake Formation registration'
    Value: !Sub 's3://${CustomerDataBucket}/'
    Export:
      Name: !Sub '${AWS::StackName}-S3Location'

  PolicyNote:
    Description: 'This template includes all required Glue and Lake Formation permissions'
    Value: 'All IAM roles include necessary permissions for federated queries and trusted identity propagation'