import { Construct } from 'constructs';
import * as lambda from 'aws-cdk-lib/aws-lambda';
import * as firehose from 'aws-cdk-lib/aws-kinesisfirehose';
import * as s3 from 'aws-cdk-lib/aws-s3';
import * as iam from 'aws-cdk-lib/aws-iam';
import { Tags, Duration, RemovalPolicy } from 'aws-cdk-lib';
import * as logs from 'aws-cdk-lib/aws-logs';

export const createRealtimeEventLoggerFirehose = async (
  scope: Construct,
  streamArn: string,
  accountId: string,
  accountStage: string,
  region: string
) => {
  const cdkTagKey = 'CreatedBy';
  const cdkTagValue = 'CDK';

  const prodBucketName = 'suno-prod-modal-event-logger';
  const stagingBucketName = 'suno-staging-modal-events-logger';
  const s3Bucket = s3.Bucket.fromBucketName(
    scope,
    'ModalRealtimeEventLoggerTempBucket',
    accountStage === 'staging' ? stagingBucketName : prodBucketName
  );

  const logGroupID = 'ModalRealtimeEventLoggerFirehoseLogGroup';

  const lambdaID = 'ModalRealtimeEventLoggerParserLambda';
  const lambdaNamePrefix = 'modal-realtime-event-logger-parser';

  const iamID = 'ModalRealtimeFirehoseServiceRole';
  const iamNamePrefix = 'KinesisFirehose-modal-realtime-event-logger';

  const kinesisNamePrefix = 'modal-realtime-event-logger';

  const firehoseID = 'ModalRealtimeEventLoggerFirehose';
  const firehoseNamePrefix = 'modal-realtime-event-logger-firehose';

  const logGroup = new logs.LogGroup(scope, logGroupID, {
    logGroupName: `/aws/kinesisfirehose/${region}/${firehoseNamePrefix}`,
    removalPolicy: RemovalPolicy.DESTROY,
    retention: logs.RetentionDays.TWO_WEEKS,
  });

  const eventLoggerParserLambda = new lambda.Function(scope, lambdaID, {
    runtime: lambda.Runtime.PYTHON_3_10,
    handler: 'main.lambda_handler',
    code: lambda.Code.fromAsset(`lambda/${lambdaNamePrefix}`),
    timeout: Duration.seconds(60),
    functionName: `${lambdaNamePrefix}-lambda`,
  });

  const firehoseRole = new iam.Role(scope, iamID, {
    roleName: `${iamNamePrefix}-${region}`,
    assumedBy: new iam.ServicePrincipal('firehose.amazonaws.com'),
    inlinePolicies: {
      FirehosePolicy: new iam.PolicyDocument({
        statements: [
          new iam.PolicyStatement({
            effect: iam.Effect.ALLOW,
            actions: ['s3:AbortMultipartUpload', 's3:GetBucketLocation', 's3:GetObject', 's3:ListBucket', 's3:ListBucketMultipartUploads', 's3:PutObject'],
            resources: [`${s3Bucket.bucketArn}`, `${s3Bucket.bucketArn}/*`],
          }),
          new iam.PolicyStatement({
            effect: iam.Effect.ALLOW,
            actions: ['lambda:InvokeFunction', 'lambda:GetFunctionConfiguration'],
            resources: [`${eventLoggerParserLambda.functionArn}`, `${eventLoggerParserLambda.functionArn}:$LATEST`],
          }),
          new iam.PolicyStatement({
            effect: iam.Effect.ALLOW,
            actions: ['kms:GenerateDataKey', 'kms:Decrypt'],
            resources: [`arn:aws:kms:${region}:${accountId}:key/%FIREHOSE_POLICY_TEMPLATE_PLACEHOLDER%`],
            conditions: {
              StringEquals: {
                'kms:ViaService': `s3.${region}.amazonaws.com`,
              },
              StringLike: {
                'kms:EncryptionContext:aws:s3:arn': [`arn:aws:s3:::${accountId}/*`, `arn:aws:s3:::${accountId}`],
              },
            },
          }),
          new iam.PolicyStatement({
            effect: iam.Effect.ALLOW,
            actions: ['logs:PutLogEvents'],
            resources: [
              `arn:aws:logs:${region}:${accountId}:log-group:/aws/kinesisfirehose/${firehoseNamePrefix}-${region}:log-stream:*`,
              `arn:aws:logs:${region}:${accountId}:log-group:%FIREHOSE_POLICY_TEMPLATE_PLACEHOLDER%:log-stream:*`,
            ],
          }),
          new iam.PolicyStatement({
            effect: iam.Effect.ALLOW,
            actions: ['kinesis:DescribeStream', 'kinesis:GetShardIterator', 'kinesis:GetRecords', 'kinesis:ListShards'],
            resources: [streamArn],
          }),
          new iam.PolicyStatement({
            effect: iam.Effect.ALLOW,
            actions: ['kms:Decrypt'],
            resources: [`arn:aws:kms:${region}:${accountId}:key/%FIREHOSE_POLICY_TEMPLATE_PLACEHOLDER%`],
            conditions: {
              StringEquals: {
                'kms:ViaService': `kinesis.${region}.amazonaws.com`,
              },
              StringLike: {
                'kms:EncryptionContext:aws:kinesis:arn': `arn:aws:kinesis:${region}:${accountId}:stream/${kinesisNamePrefix}-${region}`,
              },
            },
          }),
        ],
      }),
    },
  });

  const firehoseStream = new firehose.CfnDeliveryStream(scope, firehoseID, {
    deliveryStreamName: `${firehoseNamePrefix}-${region}`,
    deliveryStreamType: 'KinesisStreamAsSource',
    kinesisStreamSourceConfiguration: {
      kinesisStreamArn: streamArn,
      roleArn: firehoseRole.roleArn,
    },
    extendedS3DestinationConfiguration: {
      bucketArn: s3Bucket.bucketArn,
      errorOutputPrefix: 'error/!{firehose:error-output-type}/!{timestamp:yyyy/MM/dd}/',
      customTimeZone: 'UTC',
      bufferingHints: {
        sizeInMBs: 64,
        intervalInSeconds: 60,
      },
      processingConfiguration: {
        enabled: true,
        processors: [
          {
            type: 'Lambda',
            parameters: [
              {
                parameterName: 'LambdaArn',
                parameterValue: eventLoggerParserLambda.functionArn,
              },
              {
                parameterName: 'BufferSizeInMBs',
                parameterValue: '1',
              },
              {
                parameterName: 'BufferIntervalInSeconds',
                parameterValue: '60',
              },
            ],
          },
          {
            type: 'AppendDelimiterToRecord',
            parameters: [
              {
                parameterName: 'Delimiter',
                parameterValue: '\\n',
              },
            ],
          },
        ],
      },
      prefix:
        'realtime/!{partitionKeyFromLambda:EventVersion}/!{partitionKeyFromLambda:year}/!{partitionKeyFromLambda:month}/!{partitionKeyFromLambda:date}/!{partitionKeyFromLambda:hour}/',
      roleArn: firehoseRole.roleArn,
      cloudWatchLoggingOptions: {
        enabled: true,
        logGroupName: logGroup.logGroupName,
        logStreamName: 'DestinationDelivery',
      },
      dynamicPartitioningConfiguration: {
        enabled: true,
        retryOptions: {
          durationInSeconds: 300,
        },
      },
    },
  });

  Tags.of(eventLoggerParserLambda).add(cdkTagKey, cdkTagValue);
  Tags.of(firehoseStream).add(cdkTagKey, cdkTagValue);
};