import * as cdk from 'aws-cdk-lib';
import * as ec2 from 'aws-cdk-lib/aws-ec2';
import * as elasticache from 'aws-cdk-lib/aws-elasticache';
import { Construct } from 'constructs';
import { config } from '../../config'; // Import the configuration
import * as s3 from 'aws-cdk-lib/aws-s3';
import * as lambda from 'aws-cdk-lib/aws-lambda';
import * as s3n from 'aws-cdk-lib/aws-s3-notifications';
import * as iam from 'aws-cdk-lib/aws-iam';
import * as mediaconvert from 'aws-cdk-lib/aws-mediaconvert';
import * as cloudfront from 'aws-cdk-lib/aws-cloudfront';
import * as origins from 'aws-cdk-lib/aws-cloudfront-origins';
import * as certManager from 'aws-cdk-lib/aws-certificatemanager';
import * as secretsmanager from 'aws-cdk-lib/aws-secretsmanager';
import { DatadogLambda } from 'datadog-cdk-constructs-v2';

export interface ProdS3StackProps extends cdk.StackProps {
  accountStage: string;
  account: string;
  region: string;
}

export class ProdS3Stack extends cdk.Stack {
  public readonly flinkAppBucket: s3.Bucket;
  public readonly recEventsLoggerBucket: s3.Bucket;

  constructor(scope: Construct, id: string, props: ProdS3StackProps) {
    super(scope, id, props);

    const bucketName = `suno-${props.accountStage}-${props.account}`;

    // Create Flink App bucket for all environments
    this.flinkAppBucket = new s3.Bucket(this, 'FlinkAppBucket', {
      bucketName: `flink-app-bucket-${props.accountStage}-${props.account}`,
      versioned: true,
      publicReadAccess: false,
      blockPublicAccess: s3.BlockPublicAccess.BLOCK_ALL,
      removalPolicy: cdk.RemovalPolicy.RETAIN,
    });

    // Create Rec Events Logger bucket for recommendation events
    this.recEventsLoggerBucket = new s3.Bucket(this, 'RecEventsLoggerBucket', {
      bucketName: `rec-events-logger-${props.accountStage}`,
      versioned: false,
      publicReadAccess: false,
      blockPublicAccess: s3.BlockPublicAccess.BLOCK_ALL,
      removalPolicy: cdk.RemovalPolicy.RETAIN,
    });

    // Export the bucket name for cross-stack reference
    new cdk.CfnOutput(this, 'FlinkAppBucketName', {
      value: this.flinkAppBucket.bucketName,
      description: 'S3 bucket name for Flink application code and checkpoints',
      exportName: `FlinkAppBucket-${props.accountStage}`,
    });

    new cdk.CfnOutput(this, 'FlinkAppBucketArn', {
      value: this.flinkAppBucket.bucketArn,
      description: 'S3 bucket ARN for Flink application code and checkpoints',
      exportName: `FlinkAppBucketArn-${props.accountStage}`,
    });

    // Export the Rec Events Logger bucket info for cross-stack reference
    new cdk.CfnOutput(this, 'RecEventsLoggerBucketName', {
      value: this.recEventsLoggerBucket.bucketName,
      description: 'S3 bucket name for recommendation events logging',
      exportName: `RecEventsLoggerBucket-${props.accountStage}`,
    });

    new cdk.CfnOutput(this, 'RecEventsLoggerBucketArn', {
      value: this.recEventsLoggerBucket.bucketArn,
      description: 'S3 bucket ARN for recommendation events logging',
      exportName: `RecEventsLoggerBucketArn-${props.accountStage}`,
    });

    // Add bucket policy to allow all roles in the account to access Flink app code
    this.flinkAppBucket.addToResourcePolicy(new iam.PolicyStatement({
      sid: 'AllowKDAAccessToCode',
      effect: iam.Effect.ALLOW,
      principals: [new iam.AccountRootPrincipal()],
      actions: ['s3:GetObject'],
      resources: [
        `${this.flinkAppBucket.bucketArn}/*`,
      ],
    }));

    // Add minimal bucket policy for Kinesis Firehose access
    this.recEventsLoggerBucket.addToResourcePolicy(new iam.PolicyStatement({
      sid: 'AllowFirehoseServiceAccess',
      effect: iam.Effect.ALLOW,
      principals: [new iam.ServicePrincipal('firehose.amazonaws.com')],
      actions: [
        's3:GetBucketLocation',
        's3:ListBucket',
        's3:ListBucketMultipartUploads',
      ],
      resources: [this.recEventsLoggerBucket.bucketArn],
    }));

    this.recEventsLoggerBucket.addToResourcePolicy(new iam.PolicyStatement({
      sid: 'AllowFirehoseObjectAccess',
      effect: iam.Effect.ALLOW,
      principals: [new iam.ServicePrincipal('firehose.amazonaws.com')],
      actions: [
        's3:PutObject',
        's3:GetObject',
        's3:AbortMultipartUpload',
        's3:ListMultipartUploadParts',
      ],
      resources: [`${this.recEventsLoggerBucket.bucketArn}/*`],
    }));

    if (props.accountStage === 'prod' && props.region === 'us-east-1') {
      const mediaFormats = ['.wav', '.mp3', '.mp4', '.jpg', '.jpeg', '.png', '.webp', '.npz', '.opus', '.webm'];
      const bucket = s3.Bucket.fromBucketName(this, 'SunoDataUploadsBucket', `suno-data-uploads`);

      // Create a Lambda function to handle S3 events
      const s3EventHandler = new lambda.Function(this, 'S3EventHandler', {
        runtime: lambda.Runtime.NODEJS_16_X,
        handler: 'index.handler',
        code: lambda.Code.fromAsset('lambda/studio-upload-s3-tag'), // Path to your Lambda function code
        environment: {
          BUCKET_NAME: bucket.bucketName,
        },
        functionName: 'suno-s3-event-handler',
      });
      bucket.grantReadWrite(s3EventHandler);

      // for (const format of mediaFormats) {
      //   bucket.addEventNotification(s3.EventType.OBJECT_CREATED, new s3n.LambdaDestination(s3EventHandler), { suffix: format });
      // }

      // bucket creation
      const mediaSourceBucket = new s3.Bucket(this, 'MediaSourceBucket', {
        bucketName: 'suno-media-sour',
        removalPolicy: cdk.RemovalPolicy.RETAIN,
      });
      const mediaDestinationBucket = new s3.Bucket(this, 'MediaDestinationBucket', {
        bucketName: 'suno-media-dest',
        removalPolicy: cdk.RemovalPolicy.RETAIN,
        cors: [
          {
            allowedMethods: [s3.HttpMethods.GET, s3.HttpMethods.HEAD],
            allowedOrigins: ['*'],
            allowedHeaders: ['*'],
            maxAge: 3000,
          },
        ],
      });

      const mediaConvertRole = new iam.Role(this, 'MediaConvertRole', {
        assumedBy: new iam.ServicePrincipal('mediaconvert.amazonaws.com'),
        managedPolicies: [iam.ManagedPolicy.fromAwsManagedPolicyName('AWSElementalMediaConvertFullAccess')],
      });

      mediaSourceBucket.grantRead(mediaConvertRole);
      mediaDestinationBucket.grantReadWrite(mediaConvertRole);
      bucket.grantRead(mediaConvertRole);

      // Create the Lambda role
      const mediaConvertProcessorRole = new iam.Role(this, 'MediaConvertProcessorRole', {
        assumedBy: new iam.ServicePrincipal('lambda.amazonaws.com'),
      });

      // Add policies to Lambda role
      mediaConvertProcessorRole.addToPolicy(
        new iam.PolicyStatement({
          actions: ['mediaconvert:CreateJob', 'mediaconvert:DescribeEndpoints'],
          resources: ['*'],
        })
      );

      mediaConvertProcessorRole.addToPolicy(
        new iam.PolicyStatement({
          actions: ['s3:GetObject', 's3:PutObject'],
          resources: [mediaSourceBucket.arnForObjects('*'), mediaDestinationBucket.arnForObjects('*')],
        })
      );

      mediaConvertProcessorRole.addToPolicy(
        new iam.PolicyStatement({
          actions: ['logs:CreateLogGroup', 'logs:CreateLogStream', 'logs:PutLogEvents'],
          resources: ['arn:aws:logs:*:*:*'],
        })
      );

      mediaConvertProcessorRole.addToPolicy(
        new iam.PolicyStatement({
          actions: ['iam:PassRole'],
          resources: [mediaConvertRole.roleArn],
        })
      );

      // Create Datadog API key secret
      const datadogApiKeySecret = new secretsmanager.Secret(this, 'datadogApiKeySecret', {
        secretName: 'datadog-api-key',
        description: 'Secret for datadog api key'
      })

      // Create Datadog Lambda Construct
      const datadogLambda = new DatadogLambda(this, 'datadogLambda', {
        pythonLayerVersion: 111,
        extensionLayerVersion: 81,
        site:'datadoghq.com',
        apiKeySecretArn: datadogApiKeySecret.secretArn,
      });

      // Create the Lambda function that will process uploads
      const mediaConvertHandler = new lambda.Function(this, 'MediaConvertHandler', {
        runtime: lambda.Runtime.PYTHON_3_12,
        handler: 'media_convert_handler.handler',
        role: mediaConvertProcessorRole,
        code: lambda.Code.fromAsset('lambda/media-convert-handler'), // Points to a directory called 'lambda'
        timeout: cdk.Duration.minutes(15),
        environment: {
          SOURCE_BUCKET: mediaSourceBucket.bucketName,
          DESTINATION_BUCKET: mediaDestinationBucket.bucketName,
          MEDIACONVERT_ROLE: mediaConvertRole.roleArn,
          MEDIACONVERT_ENDPOINT: 'https://mediaconvert.us-east-1.amazonaws.com',
          DEPLOY_REV: `${Date.now()}`,
        },
        functionName: 'suno-media-convert-handler',
      });

      // Grant the Lambda function access to the Datadog API key secret
      datadogApiKeySecret.grantRead(mediaConvertHandler);

      // Create an alias pointing to the versioned Lambda function
      const mediaConvertHandlerAlias = new lambda.Alias(this, 'MediaConvertHandlerAlias', {
        aliasName: 'prod',
        version: mediaConvertHandler.currentVersion,
        description: 'Current version alias for media convert handler',
      });

      datadogLambda.addLambdaFunctions([mediaConvertHandler]);

      // Add permission for S3 to invoke Lambda
      new lambda.CfnPermission(this, 'MediaConvertHandlerAllowLambdaEventNotificationSource', {
        action: 'lambda:InvokeFunction',
        principal: 's3.amazonaws.com',
        sourceArn: mediaSourceBucket.bucketArn,
        functionName: mediaConvertHandler.functionName,
      });

      // Add permission for S3 to invoke Lambda
      new lambda.CfnPermission(this, 'MediaConvertHandlerAllowLambdaEventNotificationDestination', {
        action: 'lambda:InvokeFunction',
        principal: 's3.amazonaws.com',
        sourceArn: mediaDestinationBucket.bucketArn,
        functionName: mediaConvertHandler.functionName,
      });

      // Trigger the Lambda when a file is uploaded to the source bucket
      mediaSourceBucket.addEventNotification(s3.EventType.OBJECT_CREATED, new s3n.LambdaDestination(mediaConvertHandlerAlias));

      // Create the logging bucket with ACLs enabled
      const logBucket = new s3.Bucket(this, 'LogBucket', {
        bucketName: 'suno-logs-cloudfront',
        removalPolicy: cdk.RemovalPolicy.DESTROY,
        objectOwnership: s3.ObjectOwnership.OBJECT_WRITER,
        accessControl: s3.BucketAccessControl.LOG_DELIVERY_WRITE,
        blockPublicAccess: new s3.BlockPublicAccess({
          blockPublicAcls: false,
          ignorePublicAcls: false,
          blockPublicPolicy: true,
          restrictPublicBuckets: true,
        }),
        enforceSSL: true,
        // Add lifecycle rules for 7-day TTL
        lifecycleRules: [
          {
            id: '7-day-ttl',
            enabled: true,
            expiration: cdk.Duration.days(7),
            abortIncompleteMultipartUploadAfter: cdk.Duration.days(1),
            noncurrentVersionExpiration: cdk.Duration.days(1),
          },
        ],
      });

      // Create a custom cache policy for HLS content
      const hlsCachePolicy = new cloudfront.CachePolicy(this, 'HLSCachePolicy', {
        cachePolicyName: 'suno-hls-cache-policy',
        minTtl: cdk.Duration.seconds(0),
        defaultTtl: cdk.Duration.hours(24),
        maxTtl: cdk.Duration.days(7),
        headerBehavior: cloudfront.CacheHeaderBehavior.allowList('Origin', 'Access-Control-Request-Method', 'Access-Control-Request-Headers'),
        cookieBehavior: cloudfront.CacheCookieBehavior.none(),
        queryStringBehavior: cloudfront.CacheQueryStringBehavior.none(),
        enableAcceptEncodingGzip: true,
        enableAcceptEncodingBrotli: true,
      });

      // Create a response headers policy for CORS and content types
      const hlsResponseHeadersPolicy = new cloudfront.ResponseHeadersPolicy(this, 'HLSResponseHeadersPolicy', {
        responseHeadersPolicyName: 'suno-hls-response-headers-policy',
        corsBehavior: {
          accessControlAllowOrigins: ['*'], // Or restrict to specific domains
          accessControlAllowHeaders: ['*'],
          accessControlAllowCredentials: false,
          accessControlAllowMethods: ['GET', 'HEAD'],
          accessControlMaxAge: cdk.Duration.seconds(600),
          originOverride: true,
        },
      });

      // Create CloudFront distribution with specific behaviors for HLS content
      const distribution = new cloudfront.Distribution(this, 'HLSDistribution', {
        defaultBehavior: {
          origin: new origins.S3Origin(mediaDestinationBucket),
          allowedMethods: cloudfront.AllowedMethods.ALLOW_GET_HEAD,
          viewerProtocolPolicy: cloudfront.ViewerProtocolPolicy.REDIRECT_TO_HTTPS,
          cachePolicy: hlsCachePolicy,
          responseHeadersPolicy: hlsResponseHeadersPolicy,
          compress: true,
        },
        // Add specific behavior for m3u8 files with shorter TTL
        additionalBehaviors: {
          '*.m3u8': {
            origin: new origins.S3Origin(mediaDestinationBucket),
            allowedMethods: cloudfront.AllowedMethods.ALLOW_GET_HEAD,
            viewerProtocolPolicy: cloudfront.ViewerProtocolPolicy.REDIRECT_TO_HTTPS,
            cachePolicy: new cloudfront.CachePolicy(this, 'M3U8CachePolicy', {
              minTtl: cdk.Duration.seconds(0),
              defaultTtl: cdk.Duration.minutes(2), // Shorter TTL for playlist files
              maxTtl: cdk.Duration.hours(1),
              headerBehavior: cloudfront.CacheHeaderBehavior.allowList('Origin'),
              cookieBehavior: cloudfront.CacheCookieBehavior.none(),
              queryStringBehavior: cloudfront.CacheQueryStringBehavior.none(),
              enableAcceptEncodingGzip: true,
              enableAcceptEncodingBrotli: true,
            }),
            responseHeadersPolicy: hlsResponseHeadersPolicy,
            compress: true,
          },
        },
        domainNames: ['cdn3.suno.ai'],
        certificate: certManager.Certificate.fromCertificateArn(
          this,
          'Certificate',
          'arn:aws:acm:us-east-1:734185074900:certificate/7267093e-d4fb-4328-be8c-199e026a0e18'
        ),
        comment: 'HLS streaming content distribution (Production)',
        defaultRootObject: 'index.html',
        priceClass: cloudfront.PriceClass.PRICE_CLASS_ALL,
        httpVersion: cloudfront.HttpVersion.HTTP3,
        enableLogging: true,
        logBucket: logBucket,
        logFilePrefix: 'distribution-logs/',
        // Add error responses for better handling
        errorResponses: [
          {
            httpStatus: 404,
            responseHttpStatus: 404,
            responsePagePath: '/404.html',
            ttl: cdk.Duration.minutes(30),
          },
        ],
      });

      // Output values
      new cdk.CfnOutput(this, 'SourceBucketName', {
        value: mediaSourceBucket.bucketName,
        description: 'The name of the source bucket',
      });

      new cdk.CfnOutput(this, 'DestinationBucketName', {
        value: mediaDestinationBucket.bucketName,
        description: 'The name of the destination bucket',
      });

      new cdk.CfnOutput(this, 'CloudFrontDistributionDomainName', {
        value: distribution.distributionDomainName,
        description: 'The domain name of the CloudFront distribution',
      });

      new cdk.CfnOutput(this, 'MediaConvertRoleArn', {
        value: mediaConvertRole.roleArn,
        description: 'The ARN of the MediaConvert service role',
      });

      // Create a Lambda layer for the requests library
      const requestsLayer = new lambda.LayerVersion(this, 'RequestsLayer', {
        code: lambda.Code.fromAsset('lambda/media-convert-status-handler/layers', {
          bundling: {
            image: lambda.Runtime.PYTHON_3_12.bundlingImage,
            command: ['bash', '-c', 'pip install -r requirements.txt -t /asset-output/python && ' + 'cp requirements.txt /asset-output/']
          }
        }),
        compatibleRuntimes: [lambda.Runtime.PYTHON_3_12],
        description: 'Layer containing the requests library',
        layerVersionName: 'requests-layer',
      })

      // Create secrets for studio_api auth tokens in staging and prod
      const studioApiStagingAuthSecret = new secretsmanager.Secret(this, 'StudioApiStagingAuthSecret', {
        secretName: 'studio-api-staging-secret',
        description: 'Secret for lambda/modal auth token in staging'
      })

      const studioApiProdAuthSecret = new secretsmanager.Secret(this, 'studioApiProdAuthSecret', {
        secretName: 'studio-api-prod-secret',
        description: 'Secret for lambda/modal auth token in prod'
      })

      const studioApiDevAuthSecret = new secretsmanager.Secret(this, 'studioApiDevAuthSecret', {
        secretName: 'studio-api-dev-secret',
        description: 'Secret for lambda/modal auth token in dev'
      })

      // Create a Lambda function to handle MediaConvert job status changes
      const mediaConvertStatusHandler = new lambda.Function(this, 'MediaConvertStatusHandler', {
        runtime: lambda.Runtime.PYTHON_3_12,
        handler: 'media-convert-status-handler.handler',
        code: lambda.Code.fromAsset('lambda/media-convert-status-handler'),
        timeout: cdk.Duration.minutes(5),
        environment: {
          DEV_SECRET_ARN: studioApiDevAuthSecret.secretArn,
          STAGING_SECRET_ARN: studioApiStagingAuthSecret.secretArn,
          PROD_SECRET_ARN: studioApiProdAuthSecret.secretArn
        },
        functionName: 'suno-media-convert-status-handler',
        layers: [requestsLayer],
      });

      // Create an alias pointing to the function's current version
      const mediaConvertStatusHandlerAlias = new lambda.Alias(this, 'MediaConvertStatusHandlerAlias', {
        aliasName: 'prod',
        version: mediaConvertStatusHandler.currentVersion,
        description: 'Current version alias for media convert status handler',
      });

      studioApiStagingAuthSecret.grantRead(mediaConvertStatusHandler)
      studioApiProdAuthSecret.grantRead(mediaConvertStatusHandler)
      studioApiDevAuthSecret.grantRead(mediaConvertStatusHandler)

      // Grant the Lambda permissions to read/write to the buckets
      // mediaSourceBucket.grantReadWrite(mediaConvertStatusHandler);
      // mediaDestinationBucket.grantReadWrite(mediaConvertStatusHandler);

      // Import EventBridge module
      const events = require('aws-cdk-lib/aws-events');
      const targets = require('aws-cdk-lib/aws-events-targets');

      // Create EventBridge rule to listen for MediaConvert job status changes
      const mediaConvertJobStatusRule = new events.Rule(this, 'MediaConvertJobStatusRule', {
        eventPattern: {
          source: ['aws.mediaconvert'],
          detailType: ['MediaConvert Job State Change'],
          detail: {
            status: ['COMPLETE', 'ERROR']
          }
        },
        ruleName: 'suno-mediaconvert-job-status-rule',
        description: 'Rule to capture MediaConvert job completion or error events',
      });

      // Add the Lambda function as a target for the EventBridge rule
      mediaConvertJobStatusRule.addTarget(new targets.LambdaFunction(mediaConvertStatusHandlerAlias));

      // Add permission for EventBridge to invoke Lambda
      mediaConvertStatusHandler.addPermission('AllowEventBridgeInvocation', {
        principal: new iam.ServicePrincipal('events.amazonaws.com'),
        sourceArn: mediaConvertJobStatusRule.ruleArn,
      });

      const livingRadioChatModerationLambdaFileName = 'living-radio-chat-moderation';
      const livingRadioChatModerationLambdaFunctionName = `suno-${livingRadioChatModerationLambdaFileName}`;
      const ablyAwsAccountId = '203461409171';
      const sunoStagingAblyExternalId = '0bLJ1Q.dhvfXw';
      const sunoProdAblyExternalId = '0bLJ1Q.nv36Vw';

      // role for Ably to invokeliving radio chat moderation lambda
      const livingRadioChatModerationExecutionRole = new iam.Role(this, 'LivingRadioChatModerationRole', {
        assumedBy: new iam.AccountPrincipal(ablyAwsAccountId),
        externalIds: [sunoStagingAblyExternalId, sunoProdAblyExternalId],
      });

      livingRadioChatModerationExecutionRole.addToPolicy(
        new iam.PolicyStatement({
          actions: ['lambda:InvokeAsync', 'lambda:InvokeFunction'],
          resources: [`arn:aws:lambda:${props.region}:${props.account}:function:${livingRadioChatModerationLambdaFunctionName}`],
        })
      );

      // Create a Lambda function to handle living radio chat moderation
      const livingRadioChatModerationHandler = new lambda.Function(this, 'livingRadioChatModerationHandler', {
        runtime: lambda.Runtime.PYTHON_3_12,
        handler: 'main.handler',
        code: lambda.Code.fromAsset(`lambda/${livingRadioChatModerationLambdaFileName}`),
        timeout: cdk.Duration.minutes(5),
        environment: {
          DEV_SECRET_ARN: studioApiDevAuthSecret.secretArn,
          STAGING_SECRET_ARN: studioApiStagingAuthSecret.secretArn,
          PROD_SECRET_ARN: studioApiProdAuthSecret.secretArn
        },
        functionName: livingRadioChatModerationLambdaFunctionName,
        layers: [requestsLayer],
      });

      studioApiStagingAuthSecret.grantRead(livingRadioChatModerationHandler)
      studioApiProdAuthSecret.grantRead(livingRadioChatModerationHandler)
      studioApiDevAuthSecret.grantRead(livingRadioChatModerationHandler)
    }
  }
}