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 { studioSecrets } from '../secrets/studio-secrets';
import { celerySecrets } from '../secrets/studio-api-celery-secrets';
import * as applicationautoscaling from 'aws-cdk-lib/aws-applicationautoscaling';
import * as cloudwatch from 'aws-cdk-lib/aws-cloudwatch';

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

export class CeleryService extends Construct {
  public service: ecs.FargateService;

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

    const taskRole = new iam.Role(this, 'CeleryTaskRole', {
      assumedBy: new iam.ServicePrincipal('ecs-tasks.amazonaws.com'),
      managedPolicies: [
        iam.ManagedPolicy.fromAwsManagedPolicyName('service-role/AmazonECSTaskExecutionRolePolicy'),
        iam.ManagedPolicy.fromAwsManagedPolicyName('AmazonSSMManagedInstanceCore'),
      ],
    });

    // Create a dedicated execution role with custom policy
    const executionRole = new iam.Role(this, 'CeleryExecutionRole', {
      assumedBy: new iam.ServicePrincipal('ecs-tasks.amazonaws.com'),
    });

    // Add the complete execution role policy based on your existing policy
    executionRole.addToPolicy(new iam.PolicyStatement({
      effect: iam.Effect.ALLOW,
      actions: [
        'ecr:BatchCheckLayerAvailability',
        'ecr:BatchGetImage',
        'ecr:GetDownloadUrlForLayer',
        'ecr:GetAuthorizationToken',
      ],
      resources: ['*'], 
    }));

    // Add CloudWatch Logs permissions
    executionRole.addToPolicy(new iam.PolicyStatement({
      effect: iam.Effect.ALLOW,
      actions: [
        'logs:CreateLogStream',
        'logs:PutLogEvents',
      ],
      resources: [`arn:aws:logs:${cdk.Aws.REGION}:${cdk.Aws.ACCOUNT_ID}:log-group:*`],
    }));

    // Add Secrets Manager permissions  
    executionRole.addToPolicy(new iam.PolicyStatement({
      effect: iam.Effect.ALLOW,
      actions: [
        'secretsmanager:DescribeSecret',
        'secretsmanager:GetSecretValue',
      ],
      resources: [
        `arn:aws:secretsmanager:${cdk.Aws.REGION}:${cdk.Aws.ACCOUNT_ID}:secret:*`,
      ],
    }));

    // Define a Fargate task definition
    const taskDefinition = new ecs.FargateTaskDefinition(this, 'CeleryTaskDefinition', {
      memoryLimitMiB: 16384,
      cpu: 4096,
      taskRole: taskRole,
      executionRole: executionRole,
      runtimePlatform: {
        operatingSystemFamily: ecs.OperatingSystemFamily.LINUX,
        cpuArchitecture: ecs.CpuArchitecture.ARM64,
      },
      family: 'BackendInfraStackCeleryFargateServiceCeleryTaskDefinitionC3748628',
    });

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

    const secretsStudio = studioSecrets(this, props.accountStage, true);
    const secretsCelery = celerySecrets(this, props.accountStage);

    // Add a celery worker container to the task definition
    const celeryContainer = taskDefinition.addContainer('StudioApiCeleryContainer1', {
      image: ecs.ContainerImage.fromEcrRepository(repository, 'latest'),
      logging: new ecs.AwsLogDriver({
        streamPrefix: 'CeleryContainer',
      }),
      secrets: { ...secretsStudio, ...secretsCelery },
      cpu: 4096,
      memoryLimitMiB: 16384,
      command: [
        'uv',
        'run',
        '--no-dev',
        'celery',
        '--app',
        'studio_api',
        'worker',
        '--queues',
        'celery,celery_tasks',
        '--concurrency',
        '2',
        '-l',
        'info',
      ],
      essential: true,
    });

    // Create a security group for the Fargate service
    const fargateServiceSecurityGroup = new ec2.SecurityGroup(this, 'FargateServiceSG', {
      vpc: props.vpc,
      allowAllOutbound: true,
      description: 'BackendInfraStack/CeleryFargateService/FargateServiceSG',
    });

    // Define the Fargate service
    this.service = new ecs.FargateService(this, 'StudioApiFargateService', {
      cluster: props.cluster,
      taskDefinition: taskDefinition,
      serviceName: 'CeleryService',
      deploymentController: {
        type: ecs.DeploymentControllerType.ECS,
      },
      securityGroups: [fargateServiceSecurityGroup],
      vpcSubnets: {
        subnets: props.subnets,
      },
      enableExecuteCommand: true,
      desiredCount: props.accountStage === 'prod' ? 100 : 1,
    });

    // Redis endpoint configuration
    const redisEndpoint =
      props.accountStage === 'staging'
        ? 'valkey-cluster-default.ic34gd.ng.0001.use2.cache.amazonaws.com'
        : 'valkey-cluster-default.9rcjcr.ng.0001.use2.cache.amazonaws.com';

    // Create security group for Lambda to access Redis
    const lambdaSecurityGroup = new ec2.SecurityGroup(this, 'CeleryQueueMonitorLambdaSG', {
      vpc: props.vpc,
      allowAllOutbound: true,
      description: 'Security group for Celery queue monitor Lambda',
    });

    // Create Lambda function to monitor Redis queue length
    const queueMonitorLambda = new lambda.Function(this, 'CeleryQueueMonitor', {
      runtime: lambda.Runtime.PYTHON_3_11,
      handler: 'main.handler',
      timeout: cdk.Duration.minutes(5),
      vpc: props.vpc,
      vpcSubnets: {
        subnets: props.subnets,
      },
      securityGroups: [lambdaSecurityGroup],
      environment: {
        REDIS_ENDPOINT: redisEndpoint,
        REDIS_PORT: '6379',
        CELERY_QUEUES: 'celery,celery_tasks',
      },
      code: lambda.Code.fromAsset('lambda/celery-queue-monitor', {
        bundling: {
          image: lambda.Runtime.PYTHON_3_11.bundlingImage,
          command: ['bash', '-c', 'pip install -r requirements.txt -t /asset-output && cp -r . /asset-output'],
        },
      }),
    });

    // Add permissions for Lambda to publish CloudWatch metrics
    queueMonitorLambda.addToRolePolicy(
      new iam.PolicyStatement({
        actions: ['cloudwatch:PutMetricData'],
        resources: ['*'],
      })
    );

    // Add VPC configuration if Redis is in VPC
    if (props.vpc) {
      queueMonitorLambda.addToRolePolicy(
        new iam.PolicyStatement({
          actions: ['ec2:CreateNetworkInterface', 'ec2:DescribeNetworkInterfaces', 'ec2:DeleteNetworkInterface'],
          resources: ['*'],
        })
      );
    }

    // Schedule Lambda to run every minute
    const rule = new events.Rule(this, 'CeleryQueueMonitorRule', {
      schedule: events.Schedule.rate(cdk.Duration.minutes(1)),
    });
    rule.addTarget(new targets.LambdaFunction(queueMonitorLambda));

    // Configure Auto Scaling with Target Tracking
    const scalableTarget = new applicationautoscaling.ScalableTarget(this, 'CeleryScalableTarget', {
      serviceNamespace: applicationautoscaling.ServiceNamespace.ECS,
      scalableDimension: 'ecs:service:DesiredCount',
      resourceId: `service/${props.cluster.clusterName}/${this.service.serviceName}`,
      minCapacity: props.accountStage === 'prod' ? 50 : 1,
      maxCapacity: props.accountStage === 'prod' ? 200 : 10,
    });

    // Create CloudWatch metric
    const queueLengthMetric = new cloudwatch.Metric({
      namespace: 'Celery/Queue',
      metricName: 'TotalQueueLength',
      statistic: 'Average',
      period: cdk.Duration.minutes(1),
    });

    // Use Step Scaling for more precise control over empty queues
    const scaleUpPolicy = new applicationautoscaling.StepScalingPolicy(this, 'CeleryScaleUpPolicy', {
      scalingTarget: scalableTarget,
      metric: queueLengthMetric,
      scalingSteps: [
        { lower: 1, change: +5 }, // If any tasks in queue, add 5 workers quickly
        { lower: 10, change: +10 }, // If queue builds up, add 10 more
        { lower: 50, change: +20 }, // If queue is large, add 20 more
      ],
      adjustmentType: applicationautoscaling.AdjustmentType.CHANGE_IN_CAPACITY,
      cooldown: cdk.Duration.minutes(1),
      metricAggregationType: applicationautoscaling.MetricAggregationType.AVERAGE,
    });

    const scaleDownPolicy = new applicationautoscaling.StepScalingPolicy(this, 'CeleryScaleDownPolicy', {
      scalingTarget: scalableTarget,
      metric: queueLengthMetric,
      scalingSteps: [
        { upper: 0, change: -10 }, // If queue is empty, scale down aggressively
        { upper: 5, change: -2 }, // If queue is small, scale down slowly
      ],
      adjustmentType: applicationautoscaling.AdjustmentType.CHANGE_IN_CAPACITY,
      cooldown: cdk.Duration.minutes(3),
      metricAggregationType: applicationautoscaling.MetricAggregationType.AVERAGE,
    });
  }
}