import * as cdk from 'aws-cdk-lib';
import * as codepipeline from 'aws-cdk-lib/aws-codepipeline';
import * as codepipeline_actions from 'aws-cdk-lib/aws-codepipeline-actions';
import * as codebuild from 'aws-cdk-lib/aws-codebuild';
import * as ecr from 'aws-cdk-lib/aws-ecr';
import * as ecs from 'aws-cdk-lib/aws-ecs';
import * as iam from 'aws-cdk-lib/aws-iam';
import * as lambda from 'aws-cdk-lib/aws-lambda';
import * as sns from 'aws-cdk-lib/aws-sns';
import * as sns_subscriptions from 'aws-cdk-lib/aws-sns-subscriptions';
import { Construct } from 'constructs';
import { stagingConfig, stagingResources, productionConfig, productionResources } from './celery-flower-config';

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

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

    const accountStage = props.accountStage;

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

    const stagingBuildProject = this.createBuildProject(accountStage);
    const prodBuildProject = this.createProductionBuildProject(accountStage);

    const { autoApprovalTopic } = this.createAutoApprovalResources();

    const pipeline = this.createPipeline(
      sourceRepository,
      stagingBuildProject,
      prodBuildProject,
      autoApprovalTopic,
    );

    new cdk.CfnOutput(this, 'CeleryFlowerPipelineName', {
      value: pipeline.pipelineName,
      description: 'Name of the Celery Flower CodePipeline',
    });
  }

  private createAutoApprovalResources(): { autoApprovalTopic: sns.Topic } {
    const autoApprovalTopic = new sns.Topic(this, 'CeleryFlowerAutoApprovalTopic', {
      topicName: 'celery-flower-pipeline-auto-approval',
      displayName: 'Celery Flower Pipeline Auto Approval',
    });

    const autoApprovalFunction = new lambda.Function(this, 'CeleryFlowerAutoApprovalFunction', {
      runtime: lambda.Runtime.NODEJS_18_X,
      handler: 'index.handler',
      functionName: 'celery-flower-auto-approval',
      code: lambda.Code.fromInline(`
        const { CodePipelineClient, PutApprovalResultCommand } = require("@aws-sdk/client-codepipeline");
        exports.handler = async (event) => {
            const snsMessage = JSON.parse(event.Records[0].Sns.Message);
            const token = snsMessage.approval.token;
            const pipelineName = snsMessage.approval.pipelineName;
            const stageName = snsMessage.approval.stageName;
            const actionName = snsMessage.approval.actionName;
            if (pipelineName !== 'celery-flower-pipeline') return;
            const client = new CodePipelineClient({});
            const command = new PutApprovalResultCommand({
                pipelineName, stageName, actionName,
                result: { status: 'Approved', summary: 'Auto approved' },
                token
            });
            await client.send(command);
        };
      `),
      timeout: cdk.Duration.seconds(30),
    });

    autoApprovalTopic.addSubscription(new sns_subscriptions.LambdaSubscription(autoApprovalFunction));
    autoApprovalFunction.addToRolePolicy(new iam.PolicyStatement({
      actions: ['codepipeline:PutApprovalResult'],
      resources: [`arn:aws:codepipeline:${cdk.Aws.REGION}:${cdk.Aws.ACCOUNT_ID}:celery-flower-pipeline/*`],
    }));

    return { autoApprovalTopic };
  }

  private createBuildProject(accountStage: string): codebuild.PipelineProject {
    // Create IAM role for CodeBuild project
    const buildRole = new iam.Role(this, 'CeleryFlowerBuildRole', {
      assumedBy: new iam.ServicePrincipal('codebuild.amazonaws.com'),
    });

    // Logs
    buildRole.addToPolicy(new iam.PolicyStatement({
      actions: ['logs:CreateLogGroup', 'logs:CreateLogStream', 'logs:PutLogEvents'],
      resources: [
        `arn:aws:logs:${cdk.Aws.REGION}:${cdk.Aws.ACCOUNT_ID}:log-group:/aws/codebuild/*`,
        `arn:aws:logs:${cdk.Aws.REGION}:${cdk.Aws.ACCOUNT_ID}:log-group:/aws/codebuild/*:*`
      ],
    }));
    // Reports
    buildRole.addToPolicy(new iam.PolicyStatement({
      actions: ['codebuild:BatchPutCodeCoverages','codebuild:BatchPutTestCases','codebuild:CreateReport','codebuild:CreateReportGroup','codebuild:UpdateReport'],
      resources: [`arn:aws:codebuild:${cdk.Aws.REGION}:${cdk.Aws.ACCOUNT_ID}:report-group/*`],
    }));
    // S3
    buildRole.addToPolicy(new iam.PolicyStatement({
      actions: ['s3:GetObject','s3:GetObjectVersion','s3:PutObject','s3:GetBucketLocation','s3:ListBucket'],
      resources: ['arn:aws:s3:::codepipeline-*','arn:aws:s3:::codepipeline-*/*','arn:aws:s3:::aws-codebuild-*','arn:aws:s3:::aws-codebuild-*/*'],
    }));
    // KMS
    buildRole.addToPolicy(new iam.PolicyStatement({
      actions: ['kms:Decrypt','kms:DescribeKey','kms:Encrypt','kms:GenerateDataKey','kms:ReEncrypt*'],
      resources: ['*'],
    }));
    // STS
    buildRole.addToPolicy(new iam.PolicyStatement({ actions: ['sts:AssumeRole'], resources: ['*'] }));
    // ECR
    buildRole.addToPolicy(new iam.PolicyStatement({ actions: ['ecr:DescribeImages','ecr:GetAuthorizationToken'], resources: ['*'] }));

    const buildProject = new codebuild.PipelineProject(this, 'BuildProject', {
      projectName: `celery-flower-build`,
      role: buildRole,
      environment: {
        buildImage: codebuild.LinuxBuildImage.AMAZON_LINUX_2_5,
        privileged: false,
        computeType: codebuild.ComputeType.MEDIUM,
        environmentVariables: {
          STAGING_ACCOUNT_ID: { value: stagingConfig.accountId },
          TASK_DEFINITION_FAMILY: { value: stagingConfig.taskDefinitionFamily },
            CONTAINER_NAME: { value: stagingConfig.containerName },
        },
      },
      timeout: cdk.Duration.hours(1),
      buildSpec: codebuild.BuildSpec.fromObjectToYaml({
        version: '0.2',
        env: { 'exported-variables': ['COMMIT_ID'] },
        phases: {
          build: {
            commands: [
              'IMAGE_URI=$(jq -r .ImageURI imageDetail.json)',
              'echo "Image URI: $IMAGE_URI"',
              'IMAGE_DIGEST=$(echo $IMAGE_URI | sed "s/.*@//")',
              'IMAGE_TAGS=$(aws ecr describe-images --repository-name studio-api --image-ids imageDigest=$IMAGE_DIGEST --region $AWS_DEFAULT_REGION --query "imageDetails[0].imageTags" --output text)',
              'COMMIT_ID=""; for tag in $IMAGE_TAGS; do if [ "$tag" != "latest" ]; then COMMIT_ID="$tag"; break; fi; done',
              'echo "Commit ID: $COMMIT_ID"',
              `echo '[{"name":"'"$CONTAINER_NAME"'","imageUri":"'"$IMAGE_URI"'"}]' > imagedefinitions.json`,
              `STAGING_CREDS=$(aws sts assume-role --role-arn arn:aws:iam::$STAGING_ACCOUNT_ID:role/CrossAccount_Codepipeline_Role --role-session-name "get-staging-task-def")`,
              `export AWS_ACCESS_KEY_ID=$(echo $STAGING_CREDS | jq -r .Credentials.AccessKeyId)`,
              `export AWS_SECRET_ACCESS_KEY=$(echo $STAGING_CREDS | jq -r .Credentials.SecretAccessKey)`,
              `export AWS_SESSION_TOKEN=$(echo $STAGING_CREDS | jq -r .Credentials.SessionToken)`,
              `aws ecs describe-task-definition --task-definition $TASK_DEFINITION_FAMILY --region $AWS_DEFAULT_REGION | jq --arg image_uri "$IMAGE_URI" '.taskDefinition | .containerDefinitions[0].image = $image_uri | del(.taskDefinitionArn) | del(.revision) | del(.status) | del(.compatibilities) | del(.registeredAt) | del(.registeredBy)' > taskdef.json`,
              'unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN',
            ],
          },
        },
        artifacts: { 'exported-variables': ['COMMIT_ID'], files: ['taskdef.json', 'imagedefinitions.json'] },
      }),
    });

    return buildProject;
  }

  private createProductionBuildProject(accountStage: string): codebuild.PipelineProject {
    const productionBuildRole = new iam.Role(this, 'CeleryFlowerProductionBuildRole', { assumedBy: new iam.ServicePrincipal('codebuild.amazonaws.com') });
    productionBuildRole.addToPolicy(new iam.PolicyStatement({ actions: ['logs:CreateLogGroup','logs:CreateLogStream','logs:PutLogEvents'], resources: [`arn:aws:logs:${cdk.Aws.REGION}:${cdk.Aws.ACCOUNT_ID}:log-group:/aws/codebuild/*`, `arn:aws:logs:${cdk.Aws.REGION}:${cdk.Aws.ACCOUNT_ID}:log-group:/aws/codebuild/*:*`] }));
    productionBuildRole.addToPolicy(new iam.PolicyStatement({ actions: ['codebuild:BatchPutCodeCoverages','codebuild:BatchPutTestCases','codebuild:CreateReport','codebuild:CreateReportGroup','codebuild:UpdateReport'], resources: [`arn:aws:codebuild:${cdk.Aws.REGION}:${cdk.Aws.ACCOUNT_ID}:report-group/*`] }));
    productionBuildRole.addToPolicy(new iam.PolicyStatement({ actions: ['s3:GetObject','s3:GetObjectVersion','s3:PutObject','s3:GetBucketLocation','s3:ListBucket'], resources: ['arn:aws:s3:::codepipeline-*','arn:aws:s3:::codepipeline-*/*','arn:aws:s3:::aws-codebuild-*','arn:aws:s3:::aws-codebuild-*/*'] }));
    productionBuildRole.addToPolicy(new iam.PolicyStatement({ actions: ['kms:Decrypt','kms:DescribeKey','kms:Encrypt','kms:GenerateDataKey','kms:ReEncrypt*'], resources: ['*'] }));
    productionBuildRole.addToPolicy(new iam.PolicyStatement({ actions: ['sts:AssumeRole'], resources: ['*'] }));
    productionBuildRole.addToPolicy(new iam.PolicyStatement({ actions: ['ecr:DescribeImages','ecr:GetAuthorizationToken'], resources: ['*'] }));

    const productionBuildProject = new codebuild.PipelineProject(this, 'ProductionBuildProject', {
      projectName: `celery-flower-prod-build`,
      role: productionBuildRole,
      environment: {
        buildImage: codebuild.LinuxBuildImage.AMAZON_LINUX_2_5,
        privileged: false,
        computeType: codebuild.ComputeType.MEDIUM,
        environmentVariables: {
          PRODUCTION_ACCOUNT_ID: { value: productionConfig.accountId },
          TASK_DEFINITION_FAMILY: { value: productionConfig.taskDefinitionFamily },
            CONTAINER_NAME: { value: productionConfig.containerName },
        },
      },
      timeout: cdk.Duration.hours(1),
      buildSpec: codebuild.BuildSpec.fromObjectToYaml({
        version: '0.2',
        env: { 'exported-variables': ['COMMIT_ID'] },
        phases: {
          build: {
            commands: [
              'IMAGE_URI=$(jq -r .ImageURI imageDetail.json)',
              'echo "Image URI: $IMAGE_URI"',
              'IMAGE_DIGEST=$(echo $IMAGE_URI | sed "s/.*@//")',
              'IMAGE_TAGS=$(aws ecr describe-images --repository-name studio-api --image-ids imageDigest=$IMAGE_DIGEST --region $AWS_DEFAULT_REGION --query "imageDetails[0].imageTags" --output text)',
              'COMMIT_ID=""; for tag in $IMAGE_TAGS; do if [ "$tag" != "latest" ]; then COMMIT_ID="$tag"; break; fi; done',
              'echo "Commit ID: $COMMIT_ID"',
              `PROD_CREDS=$(aws sts assume-role --role-arn arn:aws:iam::$PRODUCTION_ACCOUNT_ID:role/CrossAccount_Codepipeline_Role --role-session-name "get-prod-task-def")`,
              `export AWS_ACCESS_KEY_ID=$(echo $PROD_CREDS | jq -r .Credentials.AccessKeyId)`,
              `export AWS_SECRET_ACCESS_KEY=$(echo $PROD_CREDS | jq -r .Credentials.SecretAccessKey)`,
              `export AWS_SESSION_TOKEN=$(echo $PROD_CREDS | jq -r .Credentials.SessionToken)`,
              `aws ecs describe-task-definition --task-definition $TASK_DEFINITION_FAMILY --region $AWS_DEFAULT_REGION | jq --arg image_uri "$IMAGE_URI" '.taskDefinition | .containerDefinitions[0].image = $image_uri | del(.taskDefinitionArn) | del(.revision) | del(.status) | del(.compatibilities) | del(.registeredAt) | del(.registeredBy)' > taskdef.json`,
              `echo '[{"name":"'"$CONTAINER_NAME"'","imageUri":"'"$IMAGE_URI"'"}]' > imagedefinitions.json`,
              'unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN',
            ],
          },
        },
        artifacts: { files: ['taskdef.json', 'imagedefinitions.json'] },
      }),
    });

    return productionBuildProject;
  }

  private createPipeline(
    sourceRepository: ecr.IRepository,
    stagingBuildProject: codebuild.PipelineProject,
    prodBuildProject: codebuild.PipelineProject,
    autoApprovalTopic: sns.Topic,
  ): codepipeline.Pipeline {
    const pipelineRole = new iam.Role(this, 'CeleryFlowerPipelineRole', { assumedBy: new iam.ServicePrincipal('codepipeline.amazonaws.com') });
    pipelineRole.addToPolicy(new iam.PolicyStatement({ actions: ['s3:GetObject','s3:GetObjectVersion','s3:PutObject','s3:GetBucketACL','s3:GetBucketLocation','s3:ListBucket'], resources: ['arn:aws:s3:::codepipeline-*/*','arn:aws:s3:::codepipeline-*'] }));
    pipelineRole.addToPolicy(new iam.PolicyStatement({ actions: ['kms:Decrypt','kms:Encrypt','kms:ReEncrypt*','kms:GenerateDataKey*','kms:DescribeKey'], resources: ['*'] }));

    const ecrSourceRole = new iam.Role(this, 'CeleryFlowerEcrSourceRole', { assumedBy: pipelineRole });
    ecrSourceRole.addToPolicy(new iam.PolicyStatement({ actions: ['ecr:GetDownloadUrlForLayer','ecr:BatchGetImage','ecr:DescribeImages'], resources: [sourceRepository.repositoryArn] }));
    ecrSourceRole.addToPolicy(new iam.PolicyStatement({ actions: ['s3:PutObject','s3:PutObjectLegalHold','s3:PutObjectRetention','s3:PutObjectTagging','s3:PutObjectVersionTagging'], resources: ['arn:aws:s3:::codepipeline-*/*'] }));
    ecrSourceRole.addToPolicy(new iam.PolicyStatement({ actions: ['kms:Decrypt','kms:Encrypt','kms:ReEncrypt*','kms:GenerateDataKey*','kms:DescribeKey'], resources: ['*'] }));

    const buildActionRole = new iam.Role(this, 'CeleryFlowerBuildActionRole', { assumedBy: pipelineRole });
    buildActionRole.addToPolicy(new iam.PolicyStatement({ actions: ['codebuild:StartBuild','codebuild:BatchGetBuilds','codebuild:StopBuild'], resources: ['*'] }));

    // Create dedicated role for Production CodeBuild action
    const prodBuildActionRole = new iam.Role(this, 'CeleryProdCodeBuildActionRole', {
        assumedBy: pipelineRole,
      });

      // Add policies for Production CodeBuild action
      prodBuildActionRole.addToPolicy(
        new iam.PolicyStatement({
          actions: [
          'codebuild:StartBuild',
          'codebuild:BatchGetBuilds',
          "codebuild:StopBuild"
          ],
          resources: ['*'],
        })
    );

    const deployActionRole = new iam.Role(this, 'CeleryFlowerDeployActionRole', {
      assumedBy: pipelineRole,
    });
    deployActionRole.addToPolicy(new iam.PolicyStatement({ actions: ['codebuild:StartBuild','codebuild:BatchGetBuilds','codebuild:StopBuild'], resources: ['*'] }));

    pipelineRole.addToPolicy(new iam.PolicyStatement({
      actions: ['sts:AssumeRole'],
      resources: [
        ecrSourceRole.roleArn,
        buildActionRole.roleArn,
        prodBuildActionRole.roleArn,
        deployActionRole.roleArn,
        'arn:aws:iam::590183763515:role/CrossAccount_Codepipeline_Role',
        'arn:aws:iam::734185074900:role/CrossAccount_Codepipeline_Role'
      ],
    }));

    const pipeline = new codepipeline.Pipeline(this, 'CeleryFlowerPipeline', {
      pipelineName: 'celery-flower-pipeline',
      role: pipelineRole,
      crossAccountKeys: true,
      pipelineType: codepipeline.PipelineType.V2,
      executionMode: codepipeline.ExecutionMode.SUPERSEDED,
      stages: [
        {
          stageName: 'Source',
          actions: [
            new codepipeline_actions.EcrSourceAction({
              actionName: 'Source',
              repository: sourceRepository,
              imageTag: 'latest',
              output: new codepipeline.Artifact('SourceArtifact'),
              role: ecrSourceRole,
            }),
          ],
        },
        {
          stageName: 'Build',
          actions: [
            new codepipeline_actions.CodeBuildAction({
              actionName: 'Build',
              project: stagingBuildProject,
              input: new codepipeline.Artifact('SourceArtifact'),
              outputs: [new codepipeline.Artifact('StagingBuildArtifact')],
              role: buildActionRole,
              variablesNamespace: 'BuildVariables',
              runOrder: 1,
            }),
          ],
        },
        {
          stageName: 'Staging-Deploy',
          actions: [
            new codepipeline_actions.EcsDeployAction({
              actionName: 'Staging-Deploy',
              input: new codepipeline.Artifact('StagingBuildArtifact'),
              service: ecs.BaseService.fromServiceArnWithCluster(
                this,
                'StagingCeleryFlowerServiceRef',
                `arn:aws:ecs:${cdk.Aws.REGION}:${stagingConfig.accountId}:service/${stagingResources.clusterName}/${stagingConfig.serviceName}`
              ),
              role: iam.Role.fromRoleArn(
                this,
                'StagingEcsDeployCrossAccountRole',
                `arn:aws:iam::${stagingConfig.accountId}:role/CrossAccount_Codepipeline_Role`
              ),
            }),
            new codepipeline_actions.ManualApprovalAction({
              actionName: 'Staging-CommitId',
              additionalInformation: `Approving deployment to staging for commit: #{BuildVariables.COMMIT_ID}`,
              externalEntityLink: 'https://github.com/suno-ai/glockenspiel/commit/#{BuildVariables.COMMIT_ID}',
              notificationTopic: autoApprovalTopic,
              runOrder: 1,
            }),
          ],
        },
        {
          stageName: 'Manual-Approval',
          actions: [
            new codepipeline_actions.ManualApprovalAction({
              actionName: 'Manual-Approval',
              additionalInformation: `Please review staging deployment and approve for production deployment. Commit: #{BuildVariables.COMMIT_ID}`,
              externalEntityLink: 'https://github.com/suno-ai/glockenspiel/commit/#{BuildVariables.COMMIT_ID}',
            }),
          ],
        },
        {
          stageName: 'Production-Build',
          actions: [
            new codepipeline_actions.CodeBuildAction({
              actionName: 'Production-Build',
              project: prodBuildProject,
              input: new codepipeline.Artifact('SourceArtifact'),
              outputs: [new codepipeline.Artifact('ProdBuildArtifact')],
              role: buildActionRole,
            }),
          ],
        },
        {
          stageName: 'Prod-Deploy',
          actions: [
            new codepipeline_actions.EcsDeployAction({
              actionName: 'Prod-Deploy',
              input: new codepipeline.Artifact('ProdBuildArtifact'),
              service: ecs.BaseService.fromServiceArnWithCluster(
                this,
                'ProdCeleryFlowerServiceRef',
                `arn:aws:ecs:${cdk.Aws.REGION}:${productionConfig.accountId}:service/${productionResources.clusterName}/${productionConfig.serviceName}`
              ),
              role: iam.Role.fromRoleArn(
                this,
                'ProdEcsDeployCrossAccountRole',
                `arn:aws:iam::${productionConfig.accountId}:role/CrossAccount_Codepipeline_Role`
              ),
            }),
            new codepipeline_actions.ManualApprovalAction({
              actionName: 'Production-CommitId',
              externalEntityLink: 'https://github.com/suno-ai/glockenspiel/commit/#{BuildVariables.COMMIT_ID}',
              notificationTopic: autoApprovalTopic,
            }),
          ],
        },
      ],
    });

    return pipeline;
  }
}