import * as cdk from 'aws-cdk-lib';
import { Construct } from 'constructs';
import { config } from '../../config'; // Import the configuration
import { createEventLoggerApiGateway } from './api-gateway/event-logger-api-gateway';
import * as certManager from 'aws-cdk-lib/aws-certificatemanager';
import { createEventLoggerKinesis } from './kinesis/event-logger-kinesis';
import { createEventLoggerFirehose } from './firehose/event-logger-firehose';
import { createHookEventsLambda } from './kinesis/hook-events-lambda';
import * as ec2 from 'aws-cdk-lib/aws-ec2';
import * as lambda from 'aws-cdk-lib/aws-lambda';
import * as sources from 'aws-cdk-lib/aws-lambda-event-sources';
import { multiRegionConfig } from '../../multi-region-config';
import { Duration } from 'aws-cdk-lib';
import * as iam from 'aws-cdk-lib/aws-iam';
import * as kinesis from 'aws-cdk-lib/aws-kinesis';
import * as logs from 'aws-cdk-lib/aws-logs';

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

export class EventStreamStack extends cdk.Stack {
  public readonly hooksLambda: lambda.Function;
  constructor(scope: Construct, id: string, props: EventStreamStackProps) {
    super(scope, id, props);

    // prod role
    const prodUSEast2CertificateArn = 'arn:aws:acm:us-east-2:734185074900:certificate/25359a18-026e-41b1-b30c-a8fa3df42819';
    const prodAPSoutheast1CertificateArn = 'arn:aws:acm:ap-southeast-1:734185074900:certificate/299c54bd-58d3-4239-aa64-852071c7be36';
    const prodEUcentral1CertificateArn = 'arn:aws:acm:eu-central-1:734185074900:certificate/de9b7438-ebbc-4a4d-ab0d-4364b8419d02';
    const stagingCertificateArn = 'arn:aws:acm:us-east-2:590183763515:certificate/39fc2086-d72c-430a-bf4d-f776d9ee3e93';
    const certificate = certManager.Certificate.fromCertificateArn(
      this,
      'EventLoggerApiGatewayCertificate',
      props.accountStage === 'staging'
        ? stagingCertificateArn
        : props.region === 'us-east-2'
        ? prodUSEast2CertificateArn
        : props.region === 'ap-southeast-1'
        ? prodAPSoutheast1CertificateArn
        : prodEUcentral1CertificateArn
    );

    // Create VPC for lambda
    const vpc = ec2.Vpc.fromVpcAttributes(this, 'EventLoggerKinesisVpc', {
      vpcId: multiRegionConfig[props.accountStage][props.region].vpc,
      availabilityZones: multiRegionConfig[props.accountStage][props.region].availabilityZones,
      privateSubnetIds: multiRegionConfig[props.accountStage][props.region].lambdaSubnets,
    });


    // Give lambda read-write access to all kinesis streams.
    const policyStatementReadWriteAllStreams = new iam.PolicyStatement({
      actions: [
        "kinesis:DescribeStream",
        "kinesis:DescribeStreamSummary",
        "kinesis:GetRecords",
        "kinesis:GetShardIterator",
        "kinesis:ListStreams",
        "kinesis:SubscribeToShard",
        "kinesis:PutRecord",
        "kinesis:PutRecords",
      ],
      resources: ["*"],
    });

    // Create hook events lambdas
    this.hooksLambda = createHookEventsLambda(this, vpc);
    this.hooksLambda.addToRolePolicy(policyStatementReadWriteAllStreams);
    
    const hooksMobileEventsLambda = new lambda.Function(this, 'HookEventsMobileToRecsStreamLoggerLambdaEventStreamStack', {
      runtime: lambda.Runtime.PYTHON_3_10,
      handler: 'main.lambda_handler',
      code: lambda.Code.fromAsset('lambda/hook-events-mobile-to-rec-stream-logger'),
      timeout: Duration.seconds(300),
      functionName: 'hook-events-mobile-to-rec-stream-logger-lambda',
      vpc: vpc,
      logRetention: logs.RetentionDays.TWO_WEEKS,
    });

    hooksMobileEventsLambda.addToRolePolicy(policyStatementReadWriteAllStreams);

    const hooksMobileNonPlaybackEventsLambda = new lambda.Function(this, 'HookEventsMobileNonPlaybackToRecsStreamLoggerLambdaEventStreamStack', {
      runtime: lambda.Runtime.PYTHON_3_10,
      handler: 'main.lambda_handler',
      code: lambda.Code.fromAsset('lambda/hooks-events-mobile-non-playback-to-rec-stream-logger'),
      timeout: Duration.seconds(300),
      functionName: 'hooks-events-mobile-non-playback-to-rec-stream-logger-lambda',
      vpc: vpc,
      logRetention: logs.RetentionDays.TWO_WEEKS,
    });

    hooksMobileNonPlaybackEventsLambda.addToRolePolicy(policyStatementReadWriteAllStreams);

    const hooksWebNonPlaybackEventsLambda = new lambda.Function(this, 'HookEventsWebNonPlaybackToRecsStreamLoggerLambdaEventStreamStack', {
      runtime: lambda.Runtime.PYTHON_3_10,
      handler: 'main.lambda_handler',
      code: lambda.Code.fromAsset('lambda/hook-events-web-non-playback-to-rec-stream-logger'),
      timeout: Duration.seconds(300),
      functionName: 'hook-events-web-non-playback-to-rec-stream-logger-lambda',
      vpc: vpc,
      logRetention: logs.RetentionDays.TWO_WEEKS,
    });

    hooksWebNonPlaybackEventsLambda.addToRolePolicy(policyStatementReadWriteAllStreams);

    const eventLoggerKinesis = createEventLoggerKinesis(this, props.accountStage, props.region);
    const eventLoggerFirehose = createEventLoggerFirehose(this, eventLoggerKinesis.streamArn, props.account, props.accountStage, props.region);
    const eventLoggerApiGateway = createEventLoggerApiGateway(this, certificate, props.accountStage, props.region);
    

    // Set EventLoggerStream readWrite permissions for the hook lambdas
    eventLoggerKinesis.grantReadWrite(this.hooksLambda);
    eventLoggerKinesis.grantReadWrite(hooksWebNonPlaybackEventsLambda);

    // Add EventLoggerStream as an event source for hookLambda
    this.hooksLambda.addEventSource(
      new sources.KinesisEventSource(eventLoggerKinesis, {
        batchSize: 100,
        startingPosition: lambda.StartingPosition.LATEST,
      })
    );

    // Add EventLoggerStream as an event source for web hook events lambda
    hooksWebNonPlaybackEventsLambda.addEventSource(
      new sources.KinesisEventSource(eventLoggerKinesis, {
        batchSize: 100,
        startingPosition: lambda.StartingPosition.LATEST,
      })
    );

    // Add EventLoggerTempBatchedStream as an event source for hooks processing mobile events.
    // event-logger-temp-batched only exists in us-east-2.
    if (props.region === 'us-east-2') {
      const eventLoggerTempBatchedKinesis = kinesis.Stream.fromStreamArn(
        this,
        'EventLoggerTempBatchStream',
        `arn:aws:kinesis:us-east-2:${props.account}:stream/event-logger-temp-batched`
      );
      hooksMobileEventsLambda.addEventSource(
        new sources.KinesisEventSource(eventLoggerTempBatchedKinesis, {
          batchSize: 100,
          startingPosition: lambda.StartingPosition.LATEST,
        })
      );
      hooksMobileNonPlaybackEventsLambda.addEventSource(
        new sources.KinesisEventSource(eventLoggerTempBatchedKinesis, {
          batchSize: 100,
          startingPosition: lambda.StartingPosition.LATEST,
        })
      );
    }

    // Backfill Lambda: reads failed raw events from S3 and writes decoded NDJSON grouped by event/date/hour
    const bucketName = props.accountStage === 'staging' ? 'event-logger-s3-temp' : 'event-logger-temp';
    const backfillLambda = new lambda.Function(this, 'EventLoggerBackfillLambdaEventStreamStack', {
      runtime: lambda.Runtime.PYTHON_3_10,
      handler: 'main.lambda_handler',
      code: lambda.Code.fromAsset('lambda/backfill-from-s3-event-logger-parser'),
      timeout: Duration.minutes(15),
      memorySize: 1024,
      functionName: 'event-logger-parser-backfill-lambda',
      logRetention: logs.RetentionDays.TWO_WEEKS,
      environment: {
        EVENT_LOGGER_BUCKET: bucketName,
      },
    });

    // Allow read/write on the configured bucket
    backfillLambda.addToRolePolicy(new iam.PolicyStatement({
      actions: ['s3:ListBucket'],
      resources: [`arn:aws:s3:::${bucketName}`],
    }));
    backfillLambda.addToRolePolicy(new iam.PolicyStatement({
      actions: ['s3:GetObject', 's3:PutObject'],
      resources: [`arn:aws:s3:::${bucketName}/*`],
    }));

    // Backend backfill Lambda (S3 -> S3) for backend event logger
    const backendBucketName = props.accountStage === 'staging' ? 'suno-staging-backend-event-logger' : 'suno-prod-backend-event-logger';
    const backendBackfill = new lambda.Function(this, 'BackendEventLoggerBackfillLambdaEventStreamStack', {
      runtime: lambda.Runtime.PYTHON_3_10,
      handler: 'main.lambda_handler',
      code: lambda.Code.fromAsset('lambda/backfill-from-s3-backend-event-logger-parser'),
      timeout: Duration.minutes(15),
      memorySize: 1024,
      functionName: 'backend-event-logger-parser-backfill-lambda',
      logRetention: logs.RetentionDays.TWO_WEEKS,
      environment: {
        BACKEND_EVENT_LOGGER_BUCKET: backendBucketName,
      },
    });

    // Permissions for backend backfill
    backendBackfill.addToRolePolicy(new iam.PolicyStatement({
      actions: ['s3:ListBucket'],
      resources: [
        `arn:aws:s3:::${backendBucketName}`
      ],
    }));
    backendBackfill.addToRolePolicy(new iam.PolicyStatement({
      actions: ['s3:GetObject', 's3:PutObject'],
      resources: [`arn:aws:s3:::${backendBucketName}/*`],
    }));
  }
}