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';

    // Reference the existing S3 bucket
    const prodBucketName = 'suno-prod-backend-event-logger';
    const stagingBucketName = 'suno-staging-backend-event-logger';
    const s3Bucket = s3.Bucket.fromBucketName(scope, 'BackendRealtimeEventLoggerTempBucket', accountStage === 'staging' ? stagingBucketName : prodBucketName);

    const logGroupID = 'BackendRealtimeEventLoggerFirehoseLogGroup';

    const lambdaID = 'BackendRealtimeEventLoggerParserLambda';
    const lambdaNamePrefix = 'backend-realtime-event-logger-parser';

    const iamID = 'BackendRealtimeFirehoseServiceRole';
    const iamNamePrefix = 'KinesisFirehose-backend-realtime-event-logger';
    
    const kinesisNamePrefix = 'backend-realtime-event-logger';

    const firehoseID = 'BackendRealtimeEventLoggerFirehose';
    const firehoseNamePrefix = 'backend-realtime-event-logger-firehose';

    // Create the CloudWatch log group
    const logGroup = new logs.LogGroup(scope, logGroupID, {
        logGroupName: `/aws/kinesisfirehose/${region}/${firehoseNamePrefix}`,
        removalPolicy: RemovalPolicy.DESTROY, // Adjust as needed
        retention: logs.RetentionDays.TWO_WEEKS,
    });

    // Create the Lambda function
    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`,
    });

    // Create the IAM Role
    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}`
                            }
                        }
                    })
                ]
            })
        },
    });

    // Create the Kinesis Firehose delivery stream
    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);
};