import * as aws from "@pulumi/aws";
import * as awsNative from "@pulumi/aws-native";
import * as awsx from "@pulumi/awsx";
import * as archive from "@pulumi/archive";
import * as pulumi from "@pulumi/pulumi";
import { Vpc } from "@pulumi/awsx/ec2";
import { renderAccountPrincipal } from "./config";

export type EndpointConfig = {
  name: string;
  nlbExposedPort: number;
  targetPort: number;
  url: string;
};

export const createAutoUpdateNLB = (
  vpc: Vpc,
  subnetIds: string[],
  endpoints: EndpointConfig[]
) => {
  const functionName = "nlb_updater";
  const functionResourceName = functionName.replace(/_/gi, "-");

  const nlb = new awsx.lb.NetworkLoadBalancer("networkLoadBalancer", {
    subnetIds: subnetIds,
    enableCrossZoneLoadBalancing: true,
  });

  const assumeRole = aws.iam.getPolicyDocument({
    statements: [
      {
        effect: "Allow",
        principals: [
          {
            type: "Service",
            identifiers: ["lambda.amazonaws.com"],
          },
        ],
        actions: ["sts:AssumeRole"],
      },
    ],
  });

  const lambda = archive.getFile({
    type: "zip",
    sourceFile: `${__dirname}/${functionName}.py`,
    outputPath: `${__dirname}/${functionName}.zip`,
  });

  const lambdaExecutionRole = new aws.iam.Role("cache_nlb_updater_lambda", {
    name: functionResourceName,
    assumeRolePolicy: assumeRole.then((assumeRole) => assumeRole.json),
  });

  const lambdaBasicExecutionPolicyAttachment = new aws.iam.RolePolicyAttachment(
    "lambdaBasicExecutionPolicyAttachment",
    {
      role: lambdaExecutionRole.name,
      policyArn:
        "arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole",
    }
  );

  const policyDoc = aws.iam.getPolicyDocument({
    statements: [
      {
        effect: "Allow",
        actions: [
          "logs:CreateLogGroup",
          "logs:CreateLogStream",
          "logs:PutLogEvents",
          "rds:DescribeDBInstances",
          "elasticloadbalancing:DescribeTargetHealth",
          "elasticloadbalancing:RegisterTargets",
          "elasticloadbalancing:DeregisterTargets",
        ],
        resources: ["*"],
      },
    ],
  });

  policyDoc.then((doc) => {
    const policy = new aws.iam.Policy(`${functionResourceName}-policy`, {
      name: `${functionResourceName}-policy`,
      description: "NLB Updater Function Policy",
      policy: doc.json,
    });

    new aws.iam.RolePolicyAttachment(
      `${functionResourceName}-policy-attachment`,
      {
        role: lambdaExecutionRole.name,
        policyArn: policy.arn,
      }
    );
  });

  // For each resource create a target group, listener, and lambda to keep them up to date
  endpoints.forEach((endpoint) => {
    const targetGroup = new aws.lb.TargetGroup(`${endpoint.name}-tg`, {
      name: `${endpoint.name}-tg`,
      port: endpoint.targetPort,
      protocol: "TCP",
      targetType: "ip",
      vpcId: vpc.vpcId,
    });

    new aws.lb.Listener(`${endpoint.name}-listener`, {
      loadBalancerArn: nlb.loadBalancer.arn,
      port: endpoint.nlbExposedPort,
      protocol: "TCP",
      defaultActions: [
        {
          type: "forward",
          targetGroupArn: targetGroup.arn,
        },
      ],
    });

    const lambdaFunction = new aws.lambda.Function(
      `${endpoint.name}-updater-function`,
      {
        code: new pulumi.asset.FileArchive(`${__dirname}/${functionName}.zip`),
        name: `${endpoint.name}-updater-function`,
        role: lambdaExecutionRole.arn,
        handler: `${functionName}.lambda_handler`,
        sourceCodeHash: lambda.then((lambda) => lambda.outputBase64sha256),
        runtime: aws.lambda.Runtime.Python3d12,
        environment: {
          variables: {
            CLUSTER_ENDPOINT: endpoint.url,
            TARGET_GROUP: targetGroup.arn,
          },
        },
      }
    );

    // Make a cloudwatch trigger that will invoke lambda function regularly
    const rule = new aws.cloudwatch.EventRule(`${endpoint.name}-event-rule`, {
      scheduleExpression: "rate(5 minutes)",
    });

    const target = new aws.cloudwatch.EventTarget(
      `${endpoint.name}-event-target`,
      {
        rule: rule.name,
        arn: lambdaFunction.arn,
      }
    );

    const permission = new aws.lambda.Permission(
      `${endpoint.name}-event-permission`,
      {
        action: "lambda:InvokeFunction",
        function: lambdaFunction.name,
        principal: "events.amazonaws.com",
        sourceArn: rule.arn,
      }
    );
  });

  // Create VPC Endpoint Service for NLB
  const vpcEndpointService = new aws.ec2.VpcEndpointService(
    "vpc-endpoint-service",
    {
      acceptanceRequired: true,
      networkLoadBalancerArns: [nlb.loadBalancer.arn],
      tags: {
        Name: "nlb-render-vpc-endpoint",
      },
    }
  );

  if (renderAccountPrincipal) {
    new awsNative.ec2.VpcEndpointServicePermissions(
      "vpcEndpointServicePermissionsResource",
      {
        serviceId: vpcEndpointService.id,
        allowedPrincipals: [renderAccountPrincipal],
      }
    );
  }

  return { nlb, vpcEndpointService };
};