import * as eks from "@pulumi/eks";
import * as aws from "@pulumi/aws";
import * as k8s from "@pulumi/kubernetes";
import * as pulumi from "@pulumi/pulumi";
import {
  desiredClusterSize,
  eksNodeInstanceType,
  maxClusterSize,
  minClusterSize,
} from "./config";
import { Vpc } from "@pulumi/awsx/ec2";

export const createEksCluster = (vpc: Vpc) => {
  const subnetsWithEksClusterTag = aws.ec2.getSubnets({
    tags: {
      "eks-cluster": "true",
    },
  });
  const subnetIds = subnetsWithEksClusterTag.then((subnets) => subnets.ids);
  // Create EKS cluster with the specified configuration
  const cluster = new eks.Cluster("eks-cluster", {
    name: "suno-eks-cluster",
    vpcId: vpc.vpcId,
    publicSubnetIds: vpc.publicSubnetIds,
    privateSubnetIds: subnetIds,
    instanceType: eksNodeInstanceType,
    desiredCapacity: desiredClusterSize,
    minSize: minClusterSize,
    maxSize: maxClusterSize,
    nodeAssociatePublicIpAddress: false,
    endpointPrivateAccess: false,
    endpointPublicAccess: true,
    createOidcProvider: true,
    roleMappings: getClusterRoleMapping(),
  });

  const ebsCSIDriverPolicy = new aws.iam.Policy("ebsCSI-Policy", {
    policy: {
      Version: "2012-10-17",
      Statement: [
        {
          Effect: "Allow",
          Action: [
            "ec2:CreateSnapshot",
            "ec2:AttachVolume",
            "ec2:DetachVolume",
            "ec2:ModifyVolume",
            "ec2:DescribeAvailabilityZones",
            "ec2:DescribeInstances",
            "ec2:DescribeSnapshots",
            "ec2:DescribeTags",
            "ec2:DescribeVolumes",
            "ec2:DescribeVolumesModifications",
          ],
          Resource: "*",
        },
        {
          Effect: "Allow",
          Action: ["ec2:CreateTags"],
          Resource: ["arn:aws:ec2:*:*:volume/*", "arn:aws:ec2:*:*:snapshot/*"],
          Condition: {
            StringEquals: {
              "ec2:CreateAction": ["CreateVolume", "CreateSnapshot"],
            },
          },
        },
        {
          Effect: "Allow",
          Action: ["ec2:DeleteTags"],
          Resource: ["arn:aws:ec2:*:*:volume/*", "arn:aws:ec2:*:*:snapshot/*"],
        },
        {
          Effect: "Allow",
          Action: ["ec2:CreateVolume"],
          Resource: "*",
          Condition: {
            StringLike: {
              "aws:RequestTag/ebs.csi.aws.com/cluster": "true",
            },
          },
        },
        {
          Effect: "Allow",
          Action: ["ec2:CreateVolume"],
          Resource: "*",
          Condition: {
            StringLike: {
              "aws:RequestTag/CSIVolumeName": "*",
            },
          },
        },
        {
          Effect: "Allow",
          Action: ["ec2:DeleteVolume"],
          Resource: "*",
          Condition: {
            StringLike: {
              "ec2:ResourceTag/ebs.csi.aws.com/cluster": "true",
            },
          },
        },
        {
          Effect: "Allow",
          Action: ["ec2:DeleteVolume"],
          Resource: "*",
          Condition: {
            StringLike: {
              "ec2:ResourceTag/CSIVolumeName": "*",
            },
          },
        },
        {
          Effect: "Allow",
          Action: ["ec2:DeleteVolume"],
          Resource: "*",
          Condition: {
            StringLike: {
              "ec2:ResourceTag/kubernetes.io/created-for/pvc/name": "*",
            },
          },
        },
        {
          Effect: "Allow",
          Action: ["ec2:DeleteSnapshot"],
          Resource: "*",
          Condition: {
            StringLike: {
              "ec2:ResourceTag/CSIVolumeSnapshotName": "*",
            },
          },
        },
        {
          Effect: "Allow",
          Action: ["ec2:DeleteSnapshot"],
          Resource: "*",
          Condition: {
            StringLike: {
              "ec2:ResourceTag/ebs.csi.aws.com/cluster": "true",
            },
          },
        },
      ],
    },
  });

  const serviceAccount = new k8s.core.v1.ServiceAccount(
    "ebs-csi-driver",
    {
      metadata: {
        name: "ebs-csi-driver",
        namespace: "kube-system",
      },
    },
    { provider: cluster.provider }
  );

  const ebsCSIDriverRole = serviceAccount.metadata.apply((metadata) => {
    return new aws.iam.Role("eks-ebs-csi-role", {
      assumeRolePolicy: pulumi
        .all([cluster.core.oidcProvider?.url, cluster.core.oidcProvider?.arn])
        .apply(([url, arn]) => {
          return JSON.stringify({
            Version: "2012-10-17",
            Statement: [
              {
                Effect: "Allow",
                Principal: {
                  Federated: arn,
                },
                Action: "sts:AssumeRoleWithWebIdentity",
                Condition: {
                  StringEquals: {
                    [`${url}:sub`]: `system:serviceaccount:${metadata.namespace}:${metadata.name}`,
                  },
                },
              },
            ],
          });
        }),
    });
  });

  new aws.iam.RolePolicyAttachment("ebsModifyPolicy", {
    role: ebsCSIDriverRole.name,
    policyArn: ebsCSIDriverPolicy.arn,
  });

  new k8s.core.v1.ServiceAccountPatch(
    "ebs-csi-driver-patch",
    {
      metadata: {
        name: serviceAccount.metadata.name,
        namespace: serviceAccount.metadata.namespace,
        annotations: {
          "eks.amazonaws.com/role-arn": ebsCSIDriverRole.arn,
        },
      },
    },
    { provider: cluster.provider }
  );

  return cluster;
};

const getClusterRoleMapping = () => {
  const stack = pulumi.getStack();
  if (stack === "staging") {
    const current = aws.getCallerIdentity({});
    const accountId = current.then((current) => current.accountId);
    return [
      {
        roleArn: pulumi.interpolate`arn:aws:iam::${accountId}:role/deployer`,
        username: `deployer`,
        groups: ["system:masters"],
      },
    ];
  }
  return undefined;
};
