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 codedeploy from 'aws-cdk-lib/aws-codedeploy';
import { Construct } from 'constructs';
import * as lambda from 'aws-cdk-lib/aws-lambda';
import * as iam from 'aws-cdk-lib/aws-iam';
import * as events from 'aws-cdk-lib/aws-events';
import * as targets from 'aws-cdk-lib/aws-events-targets';
import { ApplicationTargetGroup, ApplicationListener } from 'aws-cdk-lib/aws-elasticloadbalancingv2';
import * as elbv2 from 'aws-cdk-lib/aws-elasticloadbalancingv2';
import { studioSecret, studioSecrets } from '../secrets/studio-secrets';
import * as secretsmanager from 'aws-cdk-lib/aws-secretsmanager';
import { config } from '../../../config';
import { NamespaceType, DnsRecordType } from 'aws-cdk-lib/aws-servicediscovery';
import { datadogSecrets } from '../secrets/dd-secret';
import { Platform } from 'aws-cdk-lib/aws-ecr-assets';

export interface DatadogAgentServiceProps {
  accountStage: string;
  cluster: ecs.Cluster;
  vpc: ec2.IVpc;
  subnets: ec2.ISubnet[];
}

export class DatadogAgentService extends Construct {
  public service: ecs.FargateService;
  public nlb: elbv2.NetworkLoadBalancer;
  public nlbDnsName: string;

  constructor(scope: Construct, id: string, props: DatadogAgentServiceProps) {
    super(scope, id);

    const datadogTaskDefinition = new ecs.FargateTaskDefinition(this, 'DatadogTaskDef', {
      memoryLimitMiB: 65536,
      cpu: 16384,
      ephemeralStorageGiB: 200,
      family: 'BackendInfraStackdataDogAgentFargateServiceDatadogTaskDef4A005502',
    });
    const dockerCredentials = secretsmanager.Secret.fromSecretNameV2(this, 'DockerCredentials', 'docker-credential');
    const secrets = studioSecrets(this, props.accountStage, true);

    const image = ecs.ContainerImage.fromAsset('./lib/backend-infra-stack/ecs/datadog-agent-container', {
      platform: Platform.LINUX_AMD64,
    });

    // Add the Datadog container to the task definition
    const datadogContainer = datadogTaskDefinition.addContainer('DatadogAgent', {
      image,
      // image: ecs.ContainerImage.fromRegistry('datadog/agent:latest', {
      //   credentials: dockerCredentials,
      // }),
      // image: ecs.ContainerImage.fromEcrRepository(datadogRepository, 'latest'),
      logging: new ecs.AwsLogDriver({
        streamPrefix: 'datadog-agent',
      }),
      environment: {
        DD_SITE: 'datadoghq.com',
        DD_ENV: props.accountStage,
        DD_DOGSTATSD_NON_LOCAL_TRAFFIC: 'true',
        DD_APM_ENABLED: 'true',
        DD_APM_NON_LOCAL_TRAFFIC: 'true',
        DD_INTEGRATIONS_ENABLED: 'true',
        DD_DATABASEMONITORING_ENABLED: 'true',
        DD_DBM_ENABLED: 'true',
        AWS_REGION: 'us-east-2',
        DD_REMOTE_CONFIGURATION_ENABLED: 'false',
        DD_SYSTEM_PROBE_ENABLED: 'false',
        DD_PROCESS_AGENT_ENABLED: 'false',
      },
      secrets: {
        DD_API_KEY: secrets['DD_API_KEY'],
      },
      cpu: 16384,
      memoryLimitMiB: 65536,
      essential: true,
      healthCheck: {
        command: ['CMD-SHELL', 'agent health'],
        interval: cdk.Duration.seconds(30),
        timeout: cdk.Duration.seconds(5),
        retries: 3,
        startPeriod: cdk.Duration.seconds(15),
      },
    });

    // Add port mappings for DogStatsD (UDP) and APM (TCP)
    datadogContainer.addPortMappings({
      containerPort: 8125,
      protocol: ecs.Protocol.UDP,
      name: 'dogstatsd-udp',
    });

    datadogContainer.addPortMappings({
      containerPort: 8126,
      protocol: ecs.Protocol.TCP,
      name: 'apm-tcp',
    });

    // Create a security group for the Fargate service
    const dataDogFargateServiceSecurityGroup = new ec2.SecurityGroup(this, 'DatadogAgentFargateServiceSG', {
      vpc: props.vpc,
      allowAllOutbound: true,
      securityGroupName: 'DatadogAgentFargateServiceSG',
      description: 'BackendInfraStack/dataDogAgentFargateService/DatadogAgentFargateServiceSG',
    });
    // Allow inbound traffic from other services in the same VPC
    dataDogFargateServiceSecurityGroup.addIngressRule(
      ec2.Peer.ipv4(props.vpc.vpcCidrBlock),
      ec2.Port.allTraffic(),
      'Allow inbound traffic from other services in the same VPC'
    );

    // Create Service Discovery Namespace
    const namespace = props.cluster.addDefaultCloudMapNamespace({
      name: 'suno-studio-api.local',
      type: NamespaceType.DNS_PRIVATE,
    });

    // Create Network Load Balancer for stable endpoint
    this.nlb = new elbv2.NetworkLoadBalancer(this, 'DatadogNLB', {
      vpc: props.vpc,
      internetFacing: false,
      vpcSubnets: {
        subnets: props.subnets,
      },
      loadBalancerName: `datadog-agent-nlb-${props.accountStage}`,
    });

    // Define the Fargate service first
    this.service = new ecs.FargateService(this, 'datadogAgentFargateService', {
      cluster: props.cluster,
      taskDefinition: datadogTaskDefinition,
      serviceName: 'dataDogAgentService',
      deploymentController: {
        type: ecs.DeploymentControllerType.ECS,
      },
      securityGroups: [dataDogFargateServiceSecurityGroup],
      vpcSubnets: {
        subnets: props.subnets,
      },
      enableExecuteCommand: true,
      cloudMapOptions: {
        name: 'datadog-agent',
        cloudMapNamespace: namespace,
        dnsRecordType: DnsRecordType.A, // Creates a DNS A record for the service
      },
      desiredCount: 2, // Start with 2 instances for HA
      minHealthyPercent: 50,
      maxHealthyPercent: 200,
    });

    // Add UDP listener and attach the service via addTargets
    const udpListener = this.nlb.addListener('DatadogUDPListener', {
      port: 8125,
      protocol: elbv2.Protocol.UDP,
    });
    udpListener.addTargets('DatadogUDPTargets', {
      port: 8125,
      targets: [this.service.loadBalancerTarget({
        containerName: 'DatadogAgent',
        containerPort: 8125,
        protocol: ecs.Protocol.UDP,
      })],
      targetGroupName: `dd-agent-udp-${props.accountStage}`,
      protocol: elbv2.Protocol.UDP,
      healthCheck: {
        enabled: true,
        protocol: elbv2.Protocol.TCP,
        port: '8126', // Health check on TCP port since UDP health checks are limited
        interval: cdk.Duration.seconds(30),
        healthyThresholdCount: 2,
        unhealthyThresholdCount: 2,
      },
      deregistrationDelay: cdk.Duration.seconds(30),
    });

    // Add TCP listener and attach the service via addTargets
    const tcpListener = this.nlb.addListener('DatadogTCPListener', {
      port: 8126,
      protocol: elbv2.Protocol.TCP,
    });
    tcpListener.addTargets('DatadogTCPTargets', {
      port: 8126,
      targets: [this.service.loadBalancerTarget({
        containerName: 'DatadogAgent',
        containerPort: 8126,
        protocol: ecs.Protocol.TCP,
      })],
      targetGroupName: `dd-agent-tcp-${props.accountStage}`,
      protocol: elbv2.Protocol.TCP,
      healthCheck: {
        enabled: true,
        protocol: elbv2.Protocol.TCP,
        port: '8126',
        interval: cdk.Duration.seconds(30),
        healthyThresholdCount: 2,
        unhealthyThresholdCount: 2,
      },
      deregistrationDelay: cdk.Duration.seconds(30),
    });
    // Configure autoscaling
    const scaling = this.service.autoScaleTaskCount({
      minCapacity: 2,
      maxCapacity: 10,
    });

    // Scale based on CPU utilization
    scaling.scaleOnCpuUtilization('CpuScaling', {
      targetUtilizationPercent: 45,
      scaleInCooldown: cdk.Duration.seconds(300),
      scaleOutCooldown: cdk.Duration.seconds(60),
    });

    // Scale based on memory utilization
    scaling.scaleOnMemoryUtilization('MemoryScaling', {
      targetUtilizationPercent: 70,
      scaleInCooldown: cdk.Duration.seconds(300),
      scaleOutCooldown: cdk.Duration.seconds(60),
    });

    // Store NLB DNS name for services to use
    this.nlbDnsName = this.nlb.loadBalancerDnsName;

    // Output the NLB DNS name
    new cdk.CfnOutput(this, `DatadogNLBDnsName${props.accountStage}`, {
      value: this.nlbDnsName,
      description: `Datadog Agent NLB DNS name for ${props.accountStage}`,
      exportName: `DatadogAgentNLBDnsName-${props.accountStage}`,
    });
  }
}