import * as aws from "@pulumi/aws";
import { SubnetGroup } from "@pulumi/aws/dax";
import {
  Cluster,
  GetSubnetGroupResult,
  ProxyDefaultTargetGroup,
} from "@pulumi/aws/rds";
import * as pulumi from "@pulumi/pulumi";

type CreateAuroraDbParams = {
  clusterName: string;
  databaseName: string;
  replicaCount: number;
  instanceType: string;
  publiclyAccessible: boolean;
  awsRegion: Promise<aws.GetRegionResult>;
  availabilityZones: Promise<aws.GetAvailabilityZonesResult>;
  subnetGroup: Promise<GetSubnetGroupResult>;
  securityGroup: aws.ec2.SecurityGroup;
  dbAppUserSecret?: aws.secretsmanager.Secret | null;
};
type createDbReplicaParams = {
  replicaName: string;
  postgresqlCluster: Cluster;
  instanceType: string;
  publiclyAccessible: boolean;
  promotionTier?: number;
};

export const createDbReplica = ({
  replicaName,
  postgresqlCluster,
  instanceType,
  publiclyAccessible,
  promotionTier = 0,
}: createDbReplicaParams) => {
  return new aws.rds.ClusterInstance(replicaName, {
    identifier: replicaName,
    engine: aws.rds.EngineType.AuroraPostgresql,
    clusterIdentifier: postgresqlCluster.id,
    instanceClass: instanceType,
    publiclyAccessible: publiclyAccessible,
    autoMinorVersionUpgrade: true,
    performanceInsightsEnabled: true,
    performanceInsightsRetentionPeriod: 7,
    promotionTier: promotionTier,
  });
};

export const createAuroraDb = ({
  clusterName,
  databaseName,
  replicaCount,
  instanceType,
  publiclyAccessible,
  awsRegion,
  availabilityZones,
  subnetGroup,
  securityGroup,
  dbAppUserSecret = null,
}: CreateAuroraDbParams) => {
  const paramGroup = new aws.rds.ClusterParameterGroup(`${clusterName}-pg`, {
    name: `${clusterName}-pg`,
    family: "aurora-postgresql15",
    description: `${clusterName} cluster parameter group`,
    parameters: [
      {
        name: "max_sync_workers_per_subscription",
        value: "4",
        applyMethod: "pending-reboot",
      },
      {
        name: "max_logical_replication_workers",
        value: "5",
        applyMethod: "pending-reboot",
      },
      {
        name: "shared_preload_libraries",
        value: "pg_stat_statements",
        applyMethod: "pending-reboot",
      },
      {
        name: "track_activity_query_size",
        value: "4096",
        applyMethod: "pending-reboot",
      },
      {
        name: "pg_stat_statements.track",
        value: "ALL",
        applyMethod: "pending-reboot",
      },
      {
        name: "pg_stat_statements.max",
        value: "10000",
        applyMethod: "pending-reboot",
      },
      {
        name: "pg_stat_statements.track_utility",
        value: "off",
        applyMethod: "pending-reboot",
      },
      {
        name: "track_io_timing",
        value: "on",
        applyMethod: "pending-reboot",
      },
    ],
  });

  const postgresqlCluster = new aws.rds.Cluster(clusterName, {
    clusterIdentifier: clusterName,
    engine: aws.rds.EngineType.AuroraPostgresql,
    engineVersion: "15.4",
    availabilityZones: availabilityZones.then((available) => available.names),
    dbSubnetGroupName: subnetGroup.then((subnetGroup) => subnetGroup.name),
    vpcSecurityGroupIds: [securityGroup.id],
    databaseName: databaseName,
    manageMasterUserPassword: true,
    masterUsername: "postgres",
    backupRetentionPeriod: 15,
    storageType: "aurora-iopt1",
    deletionProtection: true,
    skipFinalSnapshot: false,
    dbClusterParameterGroupName: paramGroup.name.apply((name) => name),
    storageEncrypted: true,
    enabledCloudwatchLogsExports: ["postgresql"],
  });

  const dbReplicas = [];

  for (let i = 0; i < replicaCount; i++) {
    dbReplicas.push(
      createDbReplica({
        replicaName: `${clusterName}-instance-${i}`,
        postgresqlCluster,
        instanceType,
        publiclyAccessible,
      })
    );
  }

  const dbUserSecretArn = dbAppUserSecret
    ? dbAppUserSecret.arn
    : postgresqlCluster.masterUserSecrets.apply((masterUserSecrets) => {
        return masterUserSecrets[0].secretArn;
      });

  const pgKMSKey = new aws.kms.Key(`${clusterName}-kms`, {
    description: "KMS key to encrypt PostgreSQL secrets",
    enableKeyRotation: true,
  });

  const kmsRoleForProxy = new aws.iam.Role(`${clusterName}-proxy-role`, {
    name: `${clusterName}-proxy-role1`,
    assumeRolePolicy: JSON.stringify({
      Version: "2012-10-17",
      Statement: [
        {
          Action: "sts:AssumeRole",
          Effect: "Allow",
          Principal: {
            Service: "rds.amazonaws.com",
          },
        },
      ],
    }),
    inlinePolicies: [
      {
        name: "AllowSecretsManagerAccess",
        policy: pulumi
          .all([pgKMSKey.arn, dbUserSecretArn, awsRegion])
          .apply(([arn, dbUserSecretArn, awsRegion]) => {
            return JSON.stringify({
              Version: "2012-10-17",
              Statement: [
                {
                  Sid: "GetSecretValue",
                  Action: ["secretsmanager:GetSecretValue"],
                  Effect: "Allow",
                  Resource: [dbUserSecretArn],
                },
                {
                  Sid: "DecryptSecretValue",
                  Action: ["kms:Decrypt"],
                  Effect: "Allow",
                  Resource: [arn],
                  Condition: {
                    StringEquals: {
                      "kms:ViaService": `secretsmanager.${awsRegion.name}.amazonaws.com`,
                    },
                  },
                },
              ],
            });
          }),
      },
    ],
  });

  const proxy = new aws.rds.Proxy(`${clusterName}-pgdb-proxy`, {
    name: `${clusterName}-pgdb-proxy`,
    debugLogging: false,
    engineFamily: "POSTGRESQL",
    idleClientTimeout: 1800,
    requireTls: true,
    roleArn: kmsRoleForProxy.arn,
    vpcSecurityGroupIds: [securityGroup.id],
    vpcSubnetIds: subnetGroup.then((subnetGroup) => subnetGroup.subnetIds),
    auths: [
      {
        authScheme: "SECRETS",
        description: "secrets access only",
        iamAuth: "DISABLED",
        secretArn: dbUserSecretArn,
      },
    ],
  });

  const defaultProxyGroup = new aws.rds.ProxyDefaultTargetGroup(
    `${clusterName}-pgdb-proxy-group`,
    {
      dbProxyName: proxy.name,
      connectionPoolConfig: {
        connectionBorrowTimeout: 120,
        // Ensure we don't use 100% of connections in case something goes
        // rogue someone can manually still get into the DB directly
        maxConnectionsPercent: 98,
        maxIdleConnectionsPercent: 50,
      },
    }
  );

  const proxyTarget = new aws.rds.ProxyTarget(
    `${clusterName}-pgdb-proxy-target`,
    {
      dbClusterIdentifier: postgresqlCluster.id,
      dbProxyName: proxy.name,
      targetGroupName: defaultProxyGroup.name,
    }
  );

  const readOnlyProxy = new aws.rds.ProxyEndpoint(
    `${clusterName}-pgdb-proxy-readonly`,
    {
      dbProxyName: proxy.name.apply((name) => name),
      dbProxyEndpointName: `${clusterName}-pgdb-proxy-readonly`,
      vpcSubnetIds: subnetGroup.then((subnetGroup) => subnetGroup.subnetIds),
      vpcSecurityGroupIds: [securityGroup.id],
      targetRole: "READ_ONLY",
    }
  );

  return {
    postgresqlCluster,
    dbReplicas,
    proxy,
    readOnlyProxy,
  };
};
