package main

import (
	"fmt"

	"github.com/pulumi/pulumi-aws/sdk/v5/go/aws"
	"github.com/pulumi/pulumi-aws/sdk/v5/go/aws/ec2"
	"github.com/pulumi/pulumi-aws/sdk/v5/go/aws/iam"
	"github.com/pulumi/pulumi-eks/sdk/go/eks"
	"github.com/pulumi/pulumi/sdk/v3/go/pulumi"
)

type EKSCluster struct {
	ClusterName string `json:"clusterName,omitempty" yaml:"clusterName"`
	K8SVersion  string `json:"k8SVersion,omitempty" yaml:"k8SVersion"`

	VPC struct {
		ID               string   `json:"id,omitempty" yaml:"id"`
		PrivateSubnetIds []string `json:"privateSubnetIds,omitempty" yaml:"privateSubnetIds"`
		PublicSubnetIds  []string `json:"publicSubnetIds,omitempty" yaml:"publicSubnetIds"`
	} `json:"vpc,omitempty" yaml:"vpc"`

	StorageClass struct {
		Type string `json:"type,omitempty" yaml:"type"`
	} `json:"storageClass,omitempty" yaml:"storageClass"`

	NodeGroupPolicyArns         []string `json:"nodeGroupPolicyArns,omitempty" yaml:"nodeGroupPolicyArns"`
	ExtraNodeSecurityGroupRules []struct {
		Name       string   `json:"name,omitempty" yaml:"name"`
		Protocol   string   `json:"protocol,omitempty" yaml:"protocol"`
		CIDRBlocks []string `json:"cidrBlocks,omitempty" yaml:"cidrBlocks"`
		FromPort   int      `json:"fromPort,omitempty" yaml:"fromPort"`
		ToPort     int      `json:"toPort,omitempty" yaml:"toPort"`
	} `json:"extraNodeSecurityGroupRules,omitempty" yaml:"extraNodeSecurityGroupRules"`
	NodeGroups []struct {
		Name            string   `json:"name,omitempty" yaml:"name"`
		DesiredCapacity int      `json:"desiredCapacity,omitempty" yaml:"desiredCapacity"`
		InstanceType    string   `json:"instanceType,omitempty" yaml:"instanceType"`
		SSHKeyName      string   `json:"sshKeyName,omitempty" yaml:"sshKeyName"`
		MaxSize         int      `json:"maxSize,omitempty" yaml:"maxSize"`
		MinSize         int      `json:"minSize,omitempty" yaml:"minSize"`
		NodeSubnetIds   []string `json:"nodeSubnetIds,omitempty" yaml:"nodeSubnetIds"`
		AMIID           string   `json:"amiID,omitempty" yaml:"amiID"`
		Taints          []struct {
			Key    string `json:"key,omitempty" yaml:"key"`
			Effect string `json:"effect,omitempty" yaml:"effect"`
			Value  string `json:"value,omitempty" yaml:"value"`
		} `json:"taints,omitempty" yaml:"taints"`
		Labels                   map[string]string // k8s node labels
		KubeletExtraArgs         string            `json:"kubeletExtraArgs,omitempty" yaml:"kubeletExtraArgs"`
		AutoScalingGroupTags     map[string]string `json:"autoScalingGroupTags,omitempty" yaml:"autoScalingGroupTags"`
		ExtraNodeGroupPolicyArns []string          `json:"extraNodeGroupPolicyArns,omitempty" yaml:"extraNodeGroupPolicyArns"`
	} `json:"nodeGroups,omitempty" yaml:"nodeGroups"`

	AdminUsers []string `json:"adminUsers,omitempty" yaml:"adminUsers"`
	AdminRoles []string `json:"adminRoles,omitempty" yaml:"adminRoles"`
}

var ( // some constants
	nodeSecurityGroupTags = map[string]string{
		"test-infra/extra-node-security-rule": "index",
	}

	eksAdminUserGroups = []string{"system:masters"}
)

// NewEKSCluster creates eks cluster and its corresponding resources
func NewEKSCluster(ctx *pulumi.Context, c EKSCluster) error {
	// 1. create worker node ServiceRoles
	var workerNodeServiceRoles []*iam.Role
	var instanceRoles iam.RoleArray // because []*iam.Role and iam.RoleArray are different types...
	for _, nodeGroup := range c.NodeGroups {
		workerNodeServiceRole, err := createWorkerNodeServiceRole(ctx, fmt.Sprintf("%s-%s-workerRole", c.ClusterName, nodeGroup.Name), append(c.NodeGroupPolicyArns, nodeGroup.ExtraNodeGroupPolicyArns...))
		if err != nil {
			return err
		}
		workerNodeServiceRoles = append(workerNodeServiceRoles, workerNodeServiceRole)
		instanceRoles = append(instanceRoles, workerNodeServiceRole)
	}

	// 2. build cluster AdminUsers & AdminRoles
	adminUsers, err := buildClusterAdminUsers(ctx, c.AdminUsers)
	if err != nil {
		return err
	}
	adminRoles, err := buildClusterAdminRole(ctx, c.AdminRoles)
	if err != nil {
		return err
	}

	// 3. create an EKS cluster
	cluster, err := eks.NewCluster(ctx, c.ClusterName, &eks.ClusterArgs{
		Name:                  pulumi.String(c.ClusterName),
		Version:               pulumi.StringPtr(c.K8SVersion),
		CreateOidcProvider:    pulumi.BoolPtr(true),                      // OIDC provider for binding serviceAccount to aws IAM role
		SkipDefaultNodeGroup:  pulumi.BoolPtr(true),                      // deploy custom node group instead
		NodeSecurityGroupTags: pulumi.ToStringMap(nodeSecurityGroupTags), // for attaching extra rules to security group

		// VPC(networking)
		VpcId:            pulumi.StringPtr(c.VPC.ID),
		PublicSubnetIds:  pulumi.ToStringArray(c.VPC.PublicSubnetIds),
		PrivateSubnetIds: pulumi.ToStringArray(c.VPC.PrivateSubnetIds),
		// KubernetesServiceIpAddressRange: // NOTE: default value could be fine
		// StorageClass
		StorageClasses: pulumi.ToArray([]interface{}{
			eks.StorageClass{
				AllowVolumeExpansion: pulumi.BoolRef(true),
				Default:              pulumi.BoolRef(false),
				Type:                 c.StorageClass.Type,
				VolumeBindingMode:    pulumi.StringRef("WaitForFirstConsumer"),
			},
		}),

		// register workerNodeGroupServiceRole to cluster
		InstanceRoles: instanceRoles,

		// cluster admin users & admin roles
		UserMappings: eks.UserMappingArray(adminUsers),
		RoleMappings: eks.RoleMappingArray(adminRoles),
	})
	if err != nil {
		return err
	}

	// 4. create worker node groups to cluster
	for i, nodeGroup := range c.NodeGroups {
		// build workerInstanceProfile using previously generated serviceRole
		workerInstanceProfile, err := iam.NewInstanceProfile(ctx, fmt.Sprintf("%s-%s-instanceProfile", c.ClusterName, nodeGroup.Name),
			&iam.InstanceProfileArgs{Role: workerNodeServiceRoles[i]})
		if err != nil {
			return err
		}

		// build node taints
		taints := make(eks.TaintMap)
		for _, taint := range nodeGroup.Taints {
			taints[taint.Key] = eks.TaintArgs{
				Effect: pulumi.String(taint.Effect),
				Value:  pulumi.String(taint.Value),
			}
		}

		_, err = eks.NewNodeGroup(ctx, nodeGroup.Name, &eks.NodeGroupArgs{
			AutoScalingGroupTags: pulumi.ToStringMap(nodeGroup.AutoScalingGroupTags),
			AmiId:                pulumi.String(nodeGroup.AMIID),
			Cluster:              cluster.Core,
			InstanceType:         pulumi.StringPtr(nodeGroup.InstanceType),
			KeyName:              pulumi.StringPtr(nodeGroup.SSHKeyName),
			DesiredCapacity:      pulumi.IntPtr(nodeGroup.DesiredCapacity),
			MaxSize:              pulumi.IntPtr(nodeGroup.MaxSize),
			MinSize:              pulumi.IntPtr(nodeGroup.MinSize),
			NodeSubnetIds:        pulumi.ToStringArray(nodeGroup.NodeSubnetIds),
			InstanceProfile:      workerInstanceProfile,
			KubeletExtraArgs:     pulumi.StringPtr(nodeGroup.KubeletExtraArgs),
			Taints:               taints,
			Labels:               pulumi.ToStringMap(nodeGroup.Labels),
		})
		if err != nil {
			return err
		}
	}

	// 5. add extra securityGroup rules, because pulumi has issue to attach exist security group: https://github.com/pulumi/pulumi-eks/issues/616
	securityGroups, err := ec2.GetSecurityGroups(ctx, &ec2.GetSecurityGroupsArgs{Tags: nodeSecurityGroupTags})
	if err == nil { // ignores not found error (didn't find robust way to specify notFound error https://github.com/pulumi/pulumi/issues/3364)
		for _, securityGroupID := range securityGroups.Ids {
			for _, extraRule := range c.ExtraNodeSecurityGroupRules {
				_, err := ec2.NewSecurityGroupRule(ctx, fmt.Sprintf("%s-%s", securityGroupID, extraRule.Name), &ec2.SecurityGroupRuleArgs{
					SecurityGroupId: pulumi.String(securityGroupID),
					Type:            pulumi.String("ingress"),

					Protocol:   pulumi.String(extraRule.Protocol),
					CidrBlocks: pulumi.ToStringArray(extraRule.CIDRBlocks),
					FromPort:   pulumi.Int(extraRule.FromPort),
					ToPort:     pulumi.Int(extraRule.ToPort),
				})

				if err != nil {
					return err
				}
			}
		}
	}

	// finally: export the cluster's kubeconfig.
	ctx.Export("kubeconfig", cluster.Kubeconfig)

	return nil

}

func buildClusterAdminUsers(ctx *pulumi.Context, userNames []string) (eks.UserMappingArray, error) {
	current, err := aws.GetCallerIdentity(ctx, nil, nil)
	if err != nil {
		return nil, err
	}

	var adminUsers []eks.UserMappingInput
	for _, userName := range userNames {
		adminUsers = append(adminUsers, eks.UserMappingArgs{
			Groups:   pulumi.ToStringArray(eksAdminUserGroups),
			Username: pulumi.String(userName),
			UserArn:  pulumi.Sprintf("arn:aws:iam::%s:user/%s", current.AccountId, userName),
		})
	}

	return adminUsers, nil
}

func buildClusterAdminRole(ctx *pulumi.Context, roleNames []string) (eks.RoleMappingArray, error) {
	current, err := aws.GetCallerIdentity(ctx, nil, nil)
	if err != nil {
		return nil, err
	}

	var adminRoles []eks.RoleMappingInput
	for _, roleName := range roleNames {
		adminRoles = append(adminRoles, eks.RoleMappingArgs{
			Groups:   pulumi.ToStringArray(eksAdminUserGroups),
			Username: pulumi.String(roleName),
			RoleArn:  pulumi.Sprintf("arn:aws:iam::%s:role/%s", current.AccountId, roleName),
		})
	}

	return adminRoles, nil
}

// createWorkerNodeServiceRole creates a role and attaches the EKS worker node IAM managed policies.
func createWorkerNodeServiceRole(ctx *pulumi.Context, roleName string, nodeGroupPolicyArns []string) (*iam.Role, error) {
	role, err := iam.NewRole(ctx, roleName, &iam.RoleArgs{
		AssumeRolePolicy: pulumi.String(`{
	"Version": "2012-10-17",
	"Statement": [
		{
		"Effect": "Allow",
		"Principal": {
			"Service": "ec2.amazonaws.com"
		},
		"Action": "sts:AssumeRole"
		}
	]
}
`),
		ManagedPolicyArns: pulumi.ToStringArray(nodeGroupPolicyArns),
	})
	return role, err
}