import * as cdk from 'aws-cdk-lib';
import * as lambda from 'aws-cdk-lib/aws-lambda';
import * as iam from 'aws-cdk-lib/aws-iam';
import * as efs from 'aws-cdk-lib/aws-efs';
import * as ec2 from 'aws-cdk-lib/aws-ec2';
import { Construct } from 'constructs';
import { config } from '../../config'; // Import the configuration
import * as events from 'aws-cdk-lib/aws-events';
import * as targets from 'aws-cdk-lib/aws-events-targets';

export interface MLServingStackProps extends cdk.StackProps {
  accountStage: string;
}

export class MLServingStack extends cdk.Stack {
  constructor(scope: Construct, id: string, props: MLServingStackProps) {
    super(scope, id, props);

    // prod role
    if ('734185074900' === this.account) {
      // Define the VPC
      const vpc = ec2.Vpc.fromLookup(this, 'SunoMainVpc', {
        vpcId: config.prod.vpc,
      });

      const dbSubnets = config.prod.dbSubnets.map((subnetId) => ec2.Subnet.fromSubnetId(this, subnetId, subnetId));
      const lambdaSubnets = config.prod.lambdaSubnets.map((subnetId) => ec2.Subnet.fromSubnetId(this, subnetId, subnetId));

      // Create the EFS file system
      const fileSystem = new efs.FileSystem(this, 'SunoEfsFileSystem', {
        vpc,
        lifecyclePolicy: efs.LifecyclePolicy.AFTER_14_DAYS, // No lifecycle policy
        performanceMode: efs.PerformanceMode.GENERAL_PURPOSE,
        outOfInfrequentAccessPolicy: efs.OutOfInfrequentAccessPolicy.AFTER_1_ACCESS, // Ensure no transition to IA
        vpcSubnets: {
          subnets: dbSubnets,
        },
        fileSystemName: 'SunoEFSFileSystem',
      });

      // Create an access point
      const accessPoint = fileSystem.addAccessPoint('AccessPoint', {
        path: '/lambda',
        createAcl: {
          ownerUid: '1001',
          ownerGid: '1001',
          permissions: '755',
        },
        posixUser: {
          uid: '1001',
          gid: '1001',
        },
      });

      // Define the Lambda Function
      const sentenceTransformerLambda = new lambda.Function(this, 'sentenceTransformerLambda', {
        runtime: lambda.Runtime.PYTHON_3_9,
        handler: 'main.handler',
        code: lambda.Code.fromAsset('lambda/sentence-transformer'),
        vpc,
        filesystem: lambda.FileSystem.fromEfsAccessPoint(accessPoint, '/mnt/efs'),
        functionName: 'sentence-transformer-lambda',
        vpcSubnets: {
          subnets: lambdaSubnets,
        },
        timeout: cdk.Duration.minutes(15),
        memorySize: 4096,
        environment: {
          TRANSFORMERS_CACHE: '/tmp',
        },
        reservedConcurrentExecutions: 2,
      });
      // Define an alias for the Lambda function

      const sentenceTransformerLambdaAlias = new lambda.Alias(this, 'sentenceTransformerLambdaAlias', {
        aliasName: 'sentence-transformer-lambda-live',
        version: sentenceTransformerLambda.currentVersion,
        // provisionedConcurrentExecutions: 2, // Number of instances to keep warm
      });

      // Create a CloudWatch Events rule to trigger the Lambda function every 10 seconds
      const rule = new events.Rule(this, 'SentenceTransformerScheduleRule', {
        schedule: events.Schedule.rate(cdk.Duration.minutes(1)),
        ruleName: 'SentenceTransformerScheduleRule',
      });

      // Define the payload to send to the Lambda function
      const scheduledSentenceTransformerEventPayload = {
        body: JSON.stringify({ sentence: 'Scheduled invocation' }),
      };

      // Add the Lambda function alias as the target of the rule with the specified payload
      rule.addTarget(
        new targets.LambdaFunction(sentenceTransformerLambdaAlias, {
          event: events.RuleTargetInput.fromObject(scheduledSentenceTransformerEventPayload),
        })
      );

      // Grant CloudWatch Events permission to invoke the Lambda function
      sentenceTransformerLambdaAlias.addPermission('AllowCloudWatchInvoke', {
        principal: new iam.ServicePrincipal('events.amazonaws.com'),
        action: 'lambda:InvokeFunction',
        sourceArn: rule.ruleArn,
      });
    }
  }
}