import * as cdk from 'aws-cdk-lib';
import { Construct } from 'constructs';
import * as iam from 'aws-cdk-lib/aws-iam';
import * as s3 from 'aws-cdk-lib/aws-s3';
import * as logs from 'aws-cdk-lib/aws-logs';
import * as firehose from 'aws-cdk-lib/aws-kinesisfirehose';
import * as secretsmanager from 'aws-cdk-lib/aws-secretsmanager';

export interface ConfigureStudioApiDatadogLogsOptions {
  accountStage: string;
  region: string;
  accountId: string;
  apiLogGroup?: logs.ILogGroup; // optional; only present when ECS service exists in this stack
}

export function configureStudioApiDatadogLogs(scope: Construct, options: ConfigureStudioApiDatadogLogsOptions) {
  const { accountStage, region, accountId, apiLogGroup } = options;


  const stage = accountStage;
  const idSuffix = stage === 'staging' ? 'Staging' : 'Prod';
  const nameSuffix = stage; // used in names

  // Known log group names per stage
  const apiLogGroupNameByStage: Record<string, string> = {
    staging:
      'BackendInfraStack-StudioApiFargateServiceStudioApiTaskDefinitionStudioApiContainer1LogGroup9ED9C304-WyjX9jYiaCWn',
    prod:
      'BackendInfraStack-StudioApiFargateServiceStudioApiTaskDefinitionStudioApiContainer1LogGroup9ED9C304-lEldMbUlJQaI',
  };
  const celeryLogGroupNameByStage: Record<string, string> = {
    staging:
      'BackendInfraStack-CeleryFargateServiceCeleryTaskDefinitionStudioApiCeleryContainer1LogGroupBA2EDC8B-ow6eEKJYArNG',
    prod:
      'BackendInfraStack-CeleryFargateServiceCeleryTaskDefinitionStudioApiCeleryContainer1LogGroupBA2EDC8B-MYK2EdzINxSe',
  };
  const migrationLogGroupNameByStage: Record<string, string> = {
    staging:
      'BackendInfraStack-StudioApiMigrationDefinitionStudioApiMigrationContainer1LogGroup055893A1-YoTElCwJAjG5',
    prod:
      'BackendInfraStack-StudioApiMigrationDefinitionStudioApiMigrationContainer1LogGroup055893A1-mxqWGpPi72y5',
  };

  // Resolve Datadog API key from Secrets Manager secret per stage
  const ddLogsSecret = secretsmanager.Secret.fromSecretNameV2(
    scope,
    `DatadogLogsFirehoseSecret${idSuffix}Shared`,
    `datadog-log-firehose-${nameSuffix}`
  );
  const ddApiKey = ddLogsSecret.secretValueFromJson('DD_API_KEY').unsafeUnwrap();

  // Shared Firehose role
  const firehoseRole = new iam.Role(scope, `StudioApiDatadogFirehoseRole${idSuffix}`, {
    roleName: `KinesisFirehoseServiceRole-studio-api-datadog-logs-${nameSuffix}`,
    assumedBy: new iam.ServicePrincipal('firehose.amazonaws.com'),
  });

  // Shared backup bucket
  const backupBucket = new s3.Bucket(scope, `StudioApiDatadogBackupBucket${idSuffix}`, {
    bucketName: `suno-studio-api-datadog-logs-${nameSuffix}`,
    blockPublicAccess: s3.BlockPublicAccess.BLOCK_ALL,
    encryption: s3.BucketEncryption.S3_MANAGED,
    enforceSSL: true,
    removalPolicy: cdk.RemovalPolicy.RETAIN,
  });

  firehoseRole.addToPolicy(
    new iam.PolicyStatement({
      sid: 's3PermissionsShared',
      effect: iam.Effect.ALLOW,
      actions: [
        's3:AbortMultipartUpload',
        's3:GetBucketLocation',
        's3:GetObject',
        's3:ListBucket',
        's3:ListBucketMultipartUploads',
        's3:PutObject',
      ],
      resources: [backupBucket.bucketArn, `${backupBucket.bucketArn}/*`],
    })
  );

  // Firehose service CloudWatch Logs group
  const fhLogGroup = new logs.LogGroup(scope, `StudioApiDatadogFirehoseLogGroup${idSuffix}`, {
    logGroupName: `/aws/kinesisfirehose/${region}/studio-api-datadog-logs-${nameSuffix}`,
    retention: logs.RetentionDays.ONE_WEEK,
  });

  // Allow Firehose to write its own CloudWatch Logs
  firehoseRole.addToPolicy(
    new iam.PolicyStatement({
      effect: iam.Effect.ALLOW,
      actions: ['logs:CreateLogStream', 'logs:PutLogEvents'],
      resources: [
        `${fhLogGroup.logGroupArn}:*`,
        `arn:aws:logs:${region}:${accountId}:log-group:${fhLogGroup.logGroupName}:log-stream:*`,
      ],
    })
  );

  // Shared Logs→Firehose subscription role
  const logsToFirehoseRole = new iam.Role(scope, `StudioApiLogsSubscriptionRole${idSuffix}`, {
    roleName: `StudioApiLogsToFirehoseRole-${nameSuffix}`,
    assumedBy: new iam.ServicePrincipal('logs.amazonaws.com'),
  });
  logsToFirehoseRole.addToPolicy(
    new iam.PolicyStatement({
      effect: iam.Effect.ALLOW,
      actions: ['firehose:PutRecord', 'firehose:PutRecordBatch', 'firehose:DescribeDeliveryStream'],
      resources: [`arn:aws:firehose:${region}:${accountId}:deliverystream/*`],
    })
  );

  // Helper to create a DirectPut HTTP delivery stream to Datadog with shared resources
  const createDdHttpStream = (id: string, name: string) =>
    new firehose.CfnDeliveryStream(scope, id, {
      deliveryStreamName: name,
      deliveryStreamType: 'DirectPut',
      httpEndpointDestinationConfiguration: {
        endpointConfiguration: {
          url: 'https://aws-kinesis-http-intake.logs.datadoghq.com/v1/input',
          accessKey: ddApiKey,
        },
        roleArn: firehoseRole.roleArn,
        s3BackupMode: 'FailedDataOnly',
        s3Configuration: {
          bucketArn: backupBucket.bucketArn,
          bufferingHints: { intervalInSeconds: 300, sizeInMBs: 5 },
          roleArn: firehoseRole.roleArn,
          cloudWatchLoggingOptions: { enabled: true, logGroupName: fhLogGroup.logGroupName, logStreamName: 'S3Backup' },
        },
        bufferingHints: { intervalInSeconds: 60, sizeInMBs: 4 },
        cloudWatchLoggingOptions: { enabled: true, logGroupName: fhLogGroup.logGroupName, logStreamName: 'DestinationDelivery' },
        requestConfiguration: { contentEncoding: 'GZIP' },
      },
    });

  // 1) studio-api (use provided group or resolve by stage)
  const apiLogGroupResolved =
    apiLogGroup ??
    logs.LogGroup.fromLogGroupName(
      scope,
      `StudioApiExistingLogGroup${idSuffix}`,
      apiLogGroupNameByStage[stage]
    );
  if (apiLogGroupResolved) {
    const apiStream = createDdHttpStream(
      `StudioApiDatadogLogsFirehose${idSuffix}`,
      `studio-api-datadog-logs-${nameSuffix}`
    );
    new logs.CfnSubscriptionFilter(scope, `StudioApiLogsSubscriptionFilter${idSuffix}`, {
      destinationArn: apiStream.attrArn,
      filterPattern: '',
      logGroupName: apiLogGroupResolved.logGroupName,
      roleArn: logsToFirehoseRole.roleArn,
      filterName: `studio-api-logs-subscription-filter-${nameSuffix}`,
    });
  }

  // 2) studio-api-celery (existing log group by name)
  const celeryExistingLogGroup = logs.LogGroup.fromLogGroupName(
    scope,
    `StudioApiCeleryExistingLogGroup${idSuffix}`,
    celeryLogGroupNameByStage[stage]
  );
  const celeryStream = createDdHttpStream(
    `StudioApiCeleryDatadogLogsFirehose${idSuffix}`,
    `studio-api-celery-datadog-logs-${nameSuffix}`
  );
  new logs.CfnSubscriptionFilter(scope, `StudioApiCeleryLogsSubscriptionFilter${idSuffix}`, {
    destinationArn: celeryStream.attrArn,
    filterPattern: '',
    logGroupName: celeryExistingLogGroup.logGroupName,
    roleArn: logsToFirehoseRole.roleArn,
    filterName: `studio-api-celery-logs-subscription-filter-${nameSuffix}`,
  });

  // 3) studio-api-migration (existing log group by name)
  const migrationExistingLogGroup = logs.LogGroup.fromLogGroupName(
    scope,
    `StudioApiMigrationExistingLogGroup${idSuffix}`,
    migrationLogGroupNameByStage[stage]
  );
  const migrationStream = createDdHttpStream(
    `StudioApiMigrationDatadogLogsFirehose${idSuffix}`,
    `studio-api-migration-datadog-logs-${nameSuffix}`
  );
  new logs.CfnSubscriptionFilter(scope, `StudioApiMigrationLogsSubscriptionFilter${idSuffix}`, {
    destinationArn: migrationStream.attrArn,
    filterPattern: '',
    logGroupName: migrationExistingLogGroup.logGroupName,
    roleArn: logsToFirehoseRole.roleArn,
    filterName: `studio-api-migration-logs-subscription-filter-${nameSuffix}`,
  });
}