import * as cdk from 'aws-cdk-lib';
import * as ec2 from 'aws-cdk-lib/aws-ec2';
import * as ecs from 'aws-cdk-lib/aws-ecs';
import * as ecr from 'aws-cdk-lib/aws-ecr';
import * as lambda from 'aws-cdk-lib/aws-lambda';
import * as iam from 'aws-cdk-lib/aws-iam';
import { Construct } from 'constructs';
import { certificate } from './certificate/studio-api-certificate';
import { createApiGateway } from './api-gateway/studio-api-gateway';
import { studioSecret, studioSecrets } from './secrets/studio-secrets';
import { setupLoadBalancer } from './load-balancer/studio-api-load-balancer';
import { FargateService, FargateServiceProps } from './ecs/studio-api-service';
import { CeleryService } from './ecs/celery-service';
import { CeleryBeatService } from './ecs/celery-beat-service';
import { createFlowerApiGateway } from './api-gateway/studio-api-flower-gateway';
import { setupFlowerLoadBalancer } from './load-balancer/studio-api-flower-load-balancer';
import { CeleryFlowerService } from './ecs/celery-flower-service';
import { Duration } from 'aws-cdk-lib';
import { config } from '../../config'; // Import the configuration
import { aws_apigatewayv2 } from 'aws-cdk-lib';
import { DatadogAgentService } from './ecs/datadog-agent-service';
import { DatadogPostgresService } from './ecs/datadog-postgres-service';
import { DatadogRedisService } from './ecs/datadog-redis-service';
import { FluteFargateService } from './ecs/studio-api-flute-service';
import * as servicediscovery from 'aws-cdk-lib/aws-servicediscovery';
import * as logs from 'aws-cdk-lib/aws-logs';
import { setupWebSocketLoadBalancer } from './load-balancer/studio-api-websocket-load-balancer';
import { TwitterListenerService } from './ecs/twitter-lister-service';
import { studioApiServiceOnlySecrets } from './secrets/studio-api-service-only-secret';
import { configureStudioApiDatadogLogs } from './firehose/studio-api-datadog-logs';
export interface BackendInfraStackProps extends cdk.StackProps {
  accountStage: string;
}

export class BackendInfraStack extends cdk.Stack {
  constructor(scope: Construct, id: string, props: BackendInfraStackProps) {
    super(scope, id, props);
    // Fetch the AWS account ID
    const accountId = cdk.Aws.ACCOUNT_ID;
    const region = cdk.Aws.REGION;

    const accountStage = props.accountStage;

    const vpc = ec2.Vpc.fromLookup(this, 'SunoMainVpc', {
      vpcName: 'suno-main-vpc',
    });

    const subnets = config[accountStage].appSubnets.map((subnetId) => ec2.Subnet.fromSubnetId(this, `AppSubnet-${subnetId}`, subnetId));
    const publicSubnets = config[accountStage].publicSubnets.map((subnetId) => ec2.Subnet.fromSubnetId(this, `PublicSubnet-${subnetId}`, subnetId));

    // Retrieve the secret from AWS Secrets Manager
    const secrets = studioSecrets(this, accountStage, true);
    const serviceOnlySecrets = studioApiServiceOnlySecrets(this, props.accountStage);

    // Create or reference an ACM certificate
    const cert = certificate(this, accountStage, false);
    const cloudFlareCert = certificate(this, accountStage, true);

    // Create an ECS cluster
    const cluster = new ecs.Cluster(this, 'StudioApiCluster', {
      vpc: vpc,
    });

    // Setup load balancer and target groups
    const { loadBalancer, blueTargetGroup, greenTargetGroup, httpListener, loadBalancerSecurityGroup } = setupLoadBalancer(this, vpc, subnets, cert);

    // Define a Fargate task definition for migration
    const migrationTaskDefinition = new ecs.FargateTaskDefinition(this, 'StudioApiMigrationDefinition', {
      memoryLimitMiB: 32768,
      cpu: 8192,
      runtimePlatform: {
        operatingSystemFamily: ecs.OperatingSystemFamily.LINUX,
        cpuArchitecture: ecs.CpuArchitecture.ARM64,
      },
      family: 'BackendInfraStackStudioApiMigrationDefinition6EB5C15F',
    });

    const repository = ecr.Repository.fromRepositoryAttributes(this, 'studioApiRepository', {
      repositoryArn: `arn:aws:ecr:${cdk.Aws.REGION}:${cdk.Aws.ACCOUNT_ID}:repository/studio-api`,
      repositoryName: 'studio-api',
    });

    console.log('Account Stage:', accountStage);

    // Add a celery worker container to the task definition
    const migrationContainer = migrationTaskDefinition.addContainer('StudioApiMigrationContainer1', {
      image: ecs.ContainerImage.fromEcrRepository(repository, 'latest'),
      logging: new ecs.AwsLogDriver({
        streamPrefix: 'MigrationContainer',
      }),
      secrets: { ...secrets, ...serviceOnlySecrets },
      cpu: 8192,
      memoryLimitMiB: 32768,
      command: ['./run_migrations_with_retry.sh'],
      essential: true,
      environment: {
        DJANGO_DB_OPTIONS: '{"lock_timeout": 5000}',
      },
    });

    // Create a Lambda function to handle the CodeDeploy deployment
    const migrationFunction = new lambda.Function(this, 'DbMigrationFunction', {
      runtime: lambda.Runtime.NODEJS_16_X,
      handler: 'index.handler',
      code: lambda.Code.fromAsset('lambda/backend-migration'),
      environment: {
        CLUSTER_NAME: cluster.clusterName,
        SUBNET_ID: subnets[0].subnetId,
        TASK_DEFINITION_FAMILY: migrationTaskDefinition.family,
      },
      functionName: 'DbMigrationFunctionBeforeInstall',
      timeout: Duration.seconds(300),
      retryAttempts: 0,
    });

    // Grant necessary permissions to the Lambda function
    migrationFunction.addToRolePolicy(
      new iam.PolicyStatement({
        actions: [
          'ecs:ListTaskDefinitions',
          'ecs:UpdateService',
          'ecs:RunTask',
          'ecs:DescribeTasks',
          'iam:PassRole',
          'codedeploy:PutLifecycleEventHookExecutionStatus',
          'codedeploy:CreateDeployment',
          'codedeploy:GetDeployment',
          'codedeploy:GetDeploymentConfig',
          'codedeploy:GetApplication',
          'codedeploy:GetApplicationRevision',
          'codedeploy:RegisterApplicationRevision',
          'codedeploy:GetDeploymentConfig',
          'codedeploy:GetDeploymentGroup',
          'codedeploy:GetDeploymentTarget',
          'codedeploy:ListApplications',
          'codedeploy:ListDeployments',
          'codedeploy:ListDeploymentConfigs',
          'codedeploy:ListDeploymentGroups',
          'codedeploy:ListDeploymentTargets',
        ],
        resources: ['*'],
      })
    );

    // Create a Lambda layer for the notifier function
    const notifierLayer = new lambda.LayerVersion(this, 'AfterAllowTrafficNotifierLayer', {
      code: lambda.Code.fromAsset('lambda/after-allow-traffic-notifier', {
        bundling: {
          image: lambda.Runtime.NODEJS_18_X.bundlingImage,
          command: [
            'bash', '-c',
            'mkdir -p /asset-output/nodejs && cp package.json /asset-output/nodejs/ && cd /asset-output/nodejs && npm install --omit=dev --cache=/tmp/.npm'
          ],
        },
      }),
      compatibleRuntimes: [lambda.Runtime.NODEJS_18_X],
      description: 'Contains the js-yaml dependency for the AfterAllowTraffic notifier.',
    });

    // Define the Lambda function for AfterAllowTraffic hook
    const afterAllowTrafficLambdaFunction = new lambda.Function(this, 'AfterAllowTrafficNotifier', {
      runtime: lambda.Runtime.NODEJS_18_X,
      handler: 'index.handler',
      code: lambda.Code.fromAsset('lambda/after-allow-traffic-notifier'),
      functionName: 'AfterAllowTrafficNotifier',
      layers: [notifierLayer],
      timeout: Duration.seconds(30),
      logRetention: logs.RetentionDays.TWO_WEEKS,
    });

    // Grant necessary permissions to the AfterAllowTraffic Lambda function
    afterAllowTrafficLambdaFunction.addToRolePolicy(
      new iam.PolicyStatement({
        actions: [
          'iam:PassRole',
          'sts:GetCallerIdentity',
          'ecr:DescribeImages',
          'ecs:ListTaskDefinitions',
          'ecs:describeTaskDefinition',
          'codedeploy:GetDeployment',
          'codedeploy:GetApplicationRevision',
          'codedeploy:PutLifecycleEventHookExecutionStatus',
        ],
        resources: ['*'],
      })
    );

    // CrossAccount_Codepipeline_Role for CodePipeline to deploy to this account
    const crossAccountCodepipelineRole = new iam.Role(this, 'CrossAccountCodepipelineRole', {
      roleName: 'CrossAccount_Codepipeline_Role',
      assumedBy: new iam.CompositePrincipal(
        new iam.AccountPrincipal('992382411129'),
        new iam.ServicePrincipal('codepipeline.amazonaws.com')
      ),
      description: 'Cross-account CodePipeline role for deployment operations',
    });

    // Add CodeDeploy permissions
    crossAccountCodepipelineRole.addToPolicy(
      new iam.PolicyStatement({
        effect: iam.Effect.ALLOW,
        actions: [
          'codedeploy:CreateDeployment',
          'codedeploy:GetDeployment',
          'codedeploy:GetDeploymentConfig',
          'codedeploy:GetApplication',
          'codedeploy:GetApplicationRevision',
          'codedeploy:RegisterApplicationRevision',
        ],
        resources: ['*'],
      })
    );

    // Add S3 permissions
    crossAccountCodepipelineRole.addToPolicy(
      new iam.PolicyStatement({
        effect: iam.Effect.ALLOW,
        actions: [
          's3:GetObject*',
          's3:PutObject',
          's3:PutObjectAcl',
        ],
        resources: [
          '*',
        ],
      })
    );

    crossAccountCodepipelineRole.addToPolicy(new iam.PolicyStatement({
      actions: [
          "kms:Decrypt",
          "kms:DescribeKey",
          "kms:Encrypt",
          "kms:GenerateDataKey",
          "kms:ReEncrypt*",
      ],
      resources: ["*"],
    }));

    // Add ECS permissions for task definition operations
    crossAccountCodepipelineRole.addToPolicy(new iam.PolicyStatement({
      actions: [
        "ecs:DescribeTaskDefinition",
        "ecs:RegisterTaskDefinition",
        "ecs:ListTaskDefinitions",
        "ecs:DescribeServices",
        "ecs:UpdateService",
      ],
      resources: ["*"],
    }));

    // Add IAM PassRole permission for ECS task execution
    crossAccountCodepipelineRole.addToPolicy(new iam.PolicyStatement({
      actions: [
        "iam:PassRole",
      ],
      resources: ["*"],
    }));

    const dataDogAgentService = new DatadogAgentService(this, 'dataDogAgentFargateService', {
      accountStage: accountStage,
      cluster: cluster,
      vpc: vpc,
      subnets: subnets,
    });

    const datadogPostgresService = new DatadogPostgresService(this, 'datadogPostgresFargateService', {
      accountStage: accountStage,
      cluster: cluster,
      vpc: vpc,
      subnets: subnets,
    });

    const datadogRedisService = new DatadogRedisService(this, 'datadogRedisFargateService', {
      accountStage: accountStage,
      cluster: cluster,
      vpc: vpc,
      subnets: subnets,
    });

    // Define the properties for the Fargate service
    const fargateService = new FargateService(this, 'StudioApiFargateService', {
      accountStage: accountStage,
      cluster: cluster,
      vpc: vpc,
      subnets: subnets,
      blueTargetGroup: blueTargetGroup,
      greenTargetGroup: greenTargetGroup,
      httpListener: httpListener,
      loadBalancerSecurityGroup: loadBalancerSecurityGroup,
    });

    // Configure Firehose + log subscriptions via helper (staging only)
    configureStudioApiDatadogLogs(this, {
      accountStage,
      region,
      accountId,
      apiLogGroup: fargateService.containerLogGroup,
    });

    if (!cluster.defaultCloudMapNamespace) {
      throw new Error('Cluster does not have a default Cloud Map namespace');
    }

    // Create a Cloud Map service
    const cloudMapService = new servicediscovery.Service(this, 'CloudMapService', {
      namespace: cluster.defaultCloudMapNamespace,
      name: 'studio-api',
      dnsRecordType: servicediscovery.DnsRecordType.A,
      dnsTtl: cdk.Duration.seconds(60),
      routingPolicy: servicediscovery.RoutingPolicy.WEIGHTED,
      loadBalancer: true,
    });

    // Ensure InstanceId is within the allowed length
    const cloudMapServiceInstanceId = accountStage === 'prod' ? 'BackendInfraStackCloudMapServicestudioapiprod8931D58B' : 'BackendInfraStackCloudMapServicestudioapistaging8B03086E';
    const cfnInstanceLogicalId = accountStage === 'prod' ? 'CloudMapServicestudioapiprod26E389A7' : 'CloudMapServicestudioapistaging2008D32C';
    new servicediscovery.CfnInstance(this, cfnInstanceLogicalId, {
      serviceId: cloudMapService.serviceId,
      instanceId: cloudMapServiceInstanceId,
      instanceAttributes: {
        AWS_ALIAS_DNS_NAME: loadBalancer.loadBalancerDnsName,
      },
    });

    // Define the properties for the Fargate service
    const celeryService = new CeleryService(this, 'CeleryFargateService', {
      accountStage: accountStage,
      cluster: cluster,
      vpc: vpc,
      subnets: subnets,
    });

    // Define the properties for the Fargate service
    const celeryBeatService = new CeleryBeatService(this, 'CeleryBeatFargateService', {
      accountStage: accountStage,
      cluster: cluster,
      vpc: vpc,
      subnets: subnets,
    });

    const {
      loadBalancer: flowerLoadBalancer,
      blueTargetGroup: flowerTargetGroup,
      httpListener: flowerHttpListener,
      loadBalancerSecurityGroup: flowerLoadBalancerSecurityGroup,
    } = setupFlowerLoadBalancer(this, vpc, subnets, cert);
    // Create a VPC Link for HTTP API (API Gateway v2)
    const flowerVpcLink = new aws_apigatewayv2.VpcLink(this, 'CeleryFlowerHttpApiVpcLink', {
      vpcLinkName: 'CeleryFlowerHttpApiVpcLink',
      vpc: vpc,
      subnets: {
        subnets: subnets,
      },
      securityGroups: [flowerLoadBalancerSecurityGroup],
    });
    const flowerApi = createFlowerApiGateway(this, cert, accountStage, flowerVpcLink, flowerHttpListener);
    // Define the properties for the Fargate service
    const flowerFargateService = new CeleryFlowerService(this, 'CeleryFlowerFargateService', {
      accountStage: accountStage,
      cluster: cluster,
      vpc: vpc,
      subnets: subnets,
      blueTargetGroup: flowerTargetGroup,
      loadBalancerSecurityGroup: flowerLoadBalancerSecurityGroup,
    });

    // Create a Cloud Map service
    const celeryFlowerCMService = new servicediscovery.Service(this, 'CFCMService', {
      namespace: cluster.defaultCloudMapNamespace,
      name: 'celery-flower',
      dnsRecordType: servicediscovery.DnsRecordType.A,
      dnsTtl: cdk.Duration.seconds(60),
      routingPolicy: servicediscovery.RoutingPolicy.WEIGHTED,
      loadBalancer: true,
    });

    // Ensure InstanceId is within the allowed length
    const celeryFlowerInstanceId = accountStage === 'prod' ? 'BackendInfraStackCFCMServiceceleryflowerprodC88A26E7' : 'BackendInfraStackCFCMServiceceleryflowerstagingA8CBEE11';
    const cfnCeleryFlowerInstanceLogicalId = accountStage === 'prod' ? 'CFCMServiceceleryflowerprod6DB2C992' : 'CFCMServiceceleryflowerstaging9C752650';
    new servicediscovery.CfnInstance(this, cfnCeleryFlowerInstanceLogicalId, {
      serviceId: celeryFlowerCMService.serviceId,
      instanceId: celeryFlowerInstanceId,
      instanceAttributes: {
        AWS_ALIAS_DNS_NAME: flowerLoadBalancer.loadBalancerDnsName,
      },
    });

    // Create a VPC Link for HTTP API (API Gateway v2)
    const vpcLink = new aws_apigatewayv2.VpcLink(this, 'HttpApiVpcLink', {
      vpcLinkName: 'HttpApiVpcLink',
      vpc: vpc,
      subnets: {
        subnets: subnets,
      },
      securityGroups: [loadBalancerSecurityGroup],
    });

    const api = createApiGateway(this, cert, accountStage, vpcLink, httpListener, false);
    const cloudFlareApi = createApiGateway(this, cloudFlareCert, accountStage, vpcLink, httpListener, true);

    // Setup load balancer and target groups
    // const {
    //   loadBalancer: grpcLoadBalancer,
    //   blueTargetGroup: grpcBlueTargetGroup,
    //   greenTargetGroup: grpcGreenTargetGroup,
    //   httpsListener: grpcHttpsListener,
    //   loadBalancerSecurityGroup: grpcLoadBalancerSecurityGroup,
    // } = setupGrpcLoadBalancer(this, vpc, publicSubnets, cert);

    // Setup load balancer and target groups
    // const {
    //   loadBalancer: webSocketLoadBalancer,
    //   blueTargetGroup: webSocketBlueTargetGroup,
    //   greenTargetGroup: webSocketGreenTargetGroup,
    //   httpsListener: webSocketHttpsListener,
    //   webSocketLoadBalancerSecurityGroup,
    // } = setupWebSocketLoadBalancer(this, vpc, publicSubnets, cloudFlareCert);

    // const fluteFargateService = new FluteFargateService(this, 'StudioApiFluteFargateService', {
    //   accountStage: accountStage,
    //   cluster: cluster,
    //   vpc: vpc,
    //   subnets: subnets,
    //   blueTargetGroup: webSocketBlueTargetGroup,
    //   greenTargetGroup: webSocketGreenTargetGroup,
    //   httpsListener: webSocketHttpsListener,
    //   loadBalancerSecurityGroup: webSocketLoadBalancerSecurityGroup,
    // });

    // // Create a Cloud Map service
    // const fluteCloudMapService = new servicediscovery.Service(this, 'FluteCloudMapService', {
    //   namespace: cluster.defaultCloudMapNamespace,
    //   name: 'studio-api-flute',
    //   dnsRecordType: servicediscovery.DnsRecordType.A,
    //   dnsTtl: cdk.Duration.seconds(60),
    //   routingPolicy: servicediscovery.RoutingPolicy.WEIGHTED,
    //   loadBalancer: true,
    // });

    // fluteCloudMapService.registerLoadBalancer('flute', webSocketLoadBalancer);

    const twitterListenerService = new TwitterListenerService(this, 'TwitterListenerFargateService', {
      accountStage: accountStage,
      cluster: cluster,
      vpc: vpc,
      subnets: subnets,
    });
  }
}