package eks

import (
	"context"
	"crypto/md5"
	_ "embed"
	"encoding/json"
	"fmt"
	"time"

	"github.com/crossplane-contrib/provider-aws/apis/iam/v1beta1"
	xpv1 "github.com/crossplane/crossplane-runtime/apis/common/v1"
	v1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/utils/pointer"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"

	"github.com/tidbcloud/eks-provider/controllers/dm/task"
	"github.com/tidbcloud/infra-provider/pkg/deploy"
	"github.com/tidbcloud/infra-provider/pkg/k8s"
)

var (
	// AWS managed policies ARNs
	managedPolicies4dmRole = []string{
		"arn:aws:iam::aws:policy/AmazonEC2ContainerRegistryReadOnly",
		"arn:aws:iam::aws:policy/AmazonEKS_CNI_Policy",
		"arn:aws:iam::aws:policy/AmazonEKSWorkerNodePolicy",
		"arn:aws:iam::aws:policy/AmazonRoute53ResolverFullAccess",
		"arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore",
		"arn:aws:iam::aws:policy/AWSDirectConnectFullAccess",
	}
	expectedRoleAttachmentCount = len(managedPolicies4dmRole) + 1 // managed policies + custom node policy
)

var (
	//go:embed dm_node_policy.json
	dmNodePolicy []byte
)

type NodeRole struct {
	*task.BaseTask
	Now func() time.Time
}

const (
	iamReadyPropagationWait = time.Minute
	iamReadyRequeueInterval = 20 * time.Second
)

func NewNodeRole(ctx *task.Context) deploy.Task {
	return &NodeRole{BaseTask: &task.BaseTask{Context: ctx}, Now: time.Now}
}

func (t *NodeRole) Name() string {
	return "NodeRole"
}

func (t *NodeRole) Create(ctx context.Context) deploy.Result {
	ready, err := t.createNodeRole(ctx)
	if err != nil {
		t.Logger.Error(err, "failed to create node role")
		return deploy.Panic(err)
	}
	if !ready {
		return deploy.Wait(iamReadyRequeueInterval, "role or policy is not ready")
	}
	return deploy.Done()
}

func (t *NodeRole) PostCreate(ctx context.Context) deploy.Result {
	roleName := t.getAWSNodeRoleName()
	// fetch role
	role := &v1beta1.Role{
		ObjectMeta: metav1.ObjectMeta{Name: roleName},
	}
	if err := t.Local.Get(ctx, client.ObjectKeyFromObject(role), role); err != nil {
		t.Logger.Error(err, "failed to get role", "roleName", roleName)
		return deploy.Error(err)
	}
	if role.Status.AtProvider.ARN == "" {
		t.Logger.Info("role arn not ready", "roleName", roleName)
		return deploy.Wait(iamReadyRequeueInterval, "role %s ARN is not ready", roleName)
	}

	// list attachments
	attachments := &v1beta1.RolePolicyAttachmentList{}
	if err := t.Local.List(ctx, attachments, client.MatchingLabels{RoleRefLabelKey: roleName}); err != nil {
		t.Logger.Error(err, "failed to list role policy attachments", "roleName", roleName)
		return deploy.Error(err)
	}
	if len(attachments.Items) < expectedRoleAttachmentCount {
		t.Logger.Info("attachments count not match", "have", len(attachments.Items), "want", expectedRoleAttachmentCount)
		return deploy.Wait(iamReadyRequeueInterval, "role %s attachments are not ready: have %d, want %d", roleName, len(attachments.Items), expectedRoleAttachmentCount)
	}
	for _, a := range attachments.Items {
		if !isReadyConditionTrue(a.Status.Conditions) {
			t.Logger.Info("attachment not ready", "name", a.Name)
			return deploy.Wait(iamReadyRequeueInterval, "role policy attachment %s is not ready", a.Name)
		}
	}
	// wait for iamReadyPropagationWait before we can set role ARN back to DMJob status, otherwise nodes in the nodegroup might fail to join eks cluster as the role is not fully propagated in AWS yet
	latest := latestReadyTransition(attachments.Items)
	if latest.IsZero() {
		t.Logger.Info("attachment ready transition is zero, waiting for transition timestamp", "roleName", roleName)
		return deploy.Wait(iamReadyRequeueInterval, "role %s attachment ready transition is unavailable", roleName)
	}
	if t.Now().Sub(latest) < iamReadyPropagationWait {
		t.Logger.Info("waiting propagation cooldown", "since", t.Now().Sub(latest))
		return deploy.Wait(iamReadyRequeueInterval, "role %s is waiting for IAM propagation", roleName)
	}

	// all good, write back ARN
	t.DMJob.Status.AWSNodeRoleARN = role.Status.AtProvider.ARN
	return deploy.Done()
}

func (t *NodeRole) Delete(ctx context.Context) deploy.Result {
	// if err := t.deleteNodeRole(ctx); err != nil {
	// 	return deploy.Panic(err)
	// }
	return deploy.Done()
}

func (t *NodeRole) createNodeRole(ctx context.Context) (bool, error) {
	roleName := t.getAWSNodeRoleName()
	trustPolicy, err := t.getTrustPolicy()
	if err != nil {
		t.Logger.Error(err, "failed to get trust policy")
		return false, fmt.Errorf("get trust policy: %w", err)
	}

	role, roleReady, err := t.createRole(ctx, roleName, trustPolicy)
	if err != nil {
		t.Logger.Error(err, "failed to create role", "roleName", roleName)
		return false, fmt.Errorf("create role: %w", err)
	}

	// Create custom node policy
	policyARN, policyReady, err := t.createNodeRolePolicy(ctx, roleName)
	if err != nil {
		t.Logger.Error(err, "failed to create node role policy", "roleName", roleName)
		return false, fmt.Errorf("create node role policy: %w", err)
	}

	if roleReady && policyReady {
		// Append custom policy to managed policies
		allPolicies := append([]string{policyARN}, managedPolicies4dmRole...)
		if err := t.createRoleAttachment(ctx, allPolicies, role); err != nil {
			t.Logger.Error(err, "failed to create role attachment", "roleName", roleName)
			return false, fmt.Errorf("create role attachment: %w", err)
		}
		return true, nil
	}

	return false, nil
}

// func (t *NodeRole) deleteNodeRole(ctx context.Context) error {
// 	roleName := t.getAWSNodeRoleName()

// 	if err := client.IgnoreNotFound(t.Local.DeleteAllOf(ctx,
// 		&v1beta1.RolePolicyAttachment{}, client.MatchingLabels{RoleRefLabelKey: roleName})); err != nil {
// 		return fmt.Errorf("delete rolePolicyAttachment: %w", err)
// 	}

// 	if err := client.IgnoreNotFound(t.Local.DeleteAllOf(ctx,
// 		&v1beta1.Policy{}, client.MatchingLabels{RoleRefLabelKey: roleName})); err != nil {
// 		return fmt.Errorf("delete policy: %w", err)
// 	}

// 	role := &v1beta1.Role{
// 		ObjectMeta: metav1.ObjectMeta{
// 			Name: roleName,
// 		},
// 	}
// 	if err := client.IgnoreNotFound(t.Local.Delete(ctx, role)); err != nil {
// 		return fmt.Errorf("delete IAM role: %w", err)
// 	}

// 	return nil
// }

func (t *NodeRole) getAWSNodeRoleName() string {
	return fmt.Sprintf("dm-%s-%s", t.Config.Network.Region, t.ProjectID)
}

func (t *NodeRole) getTrustPolicy() (string, error) {
	trustPolicy := map[string]interface{}{
		"Version": "2012-10-17",
		"Statement": []map[string]interface{}{
			{
				"Effect": "Allow",
				"Principal": map[string]interface{}{
					"Service": "ec2.amazonaws.com",
				},
				"Action": []string{"sts:AssumeRole"},
			},
		},
	}

	policyBytes, err := json.Marshal(trustPolicy)
	if err != nil {
		return "", fmt.Errorf("marshal trust policy: %w", err)
	}

	return string(policyBytes), nil
}

func (t *NodeRole) createRole(ctx context.Context, roleName, trustPolicy string) (*v1beta1.Role, bool, error) {
	role := &v1beta1.Role{
		ObjectMeta: metav1.ObjectMeta{
			Name: roleName,
		},
		Spec: v1beta1.RoleSpec{
			ForProvider: v1beta1.RoleParameters{
				Description:              pointer.String("Node role for DM clusters"),
				AssumeRolePolicyDocument: trustPolicy,
			},
			ResourceSpec: xpv1.ResourceSpec{
				DeletionPolicy:          xpv1.DeletionDelete,
				ProviderConfigReference: &xpv1.Reference{Name: t.GetProviderConfigName()},
			},
		},
	}

	err := k8s.Apply(ctx, t.Local, role)
	if err != nil {
		return role, false, err
	}

	return role, role.Status.AtProvider.ARN != "", nil
}

func (t *NodeRole) createNodeRolePolicy(ctx context.Context, roleName string) (string, bool, error) {
	policyName := policyName(roleName)
	policy := &v1beta1.Policy{
		ObjectMeta: metav1.ObjectMeta{
			Name:   policyName,
			Labels: map[string]string{RoleRefLabelKey: roleName},
		},
		Spec: v1beta1.PolicySpec{
			ForProvider: v1beta1.PolicyParameters{
				Description: pointer.String("Node role permissions for DM clusters"),
				Name:        policyName,
				Document:    string(dmNodePolicy),
			},
			ResourceSpec: xpv1.ResourceSpec{
				DeletionPolicy:          xpv1.DeletionDelete,
				ProviderConfigReference: &xpv1.Reference{Name: t.GetProviderConfigName()},
			},
		},
	}

	_, err := controllerutil.CreateOrUpdate(ctx, t.Local, policy, func() error {
		return nil
	})
	if err != nil {
		return "", false, err
	}

	return policy.Status.AtProvider.ARN, policy.Status.AtProvider.ARN != "", nil
}

func (t *NodeRole) createRoleAttachment(ctx context.Context, policyARNs []string, role *v1beta1.Role) error {
	for _, policyARN := range policyARNs {
		attachment := &v1beta1.RolePolicyAttachment{
			ObjectMeta: metav1.ObjectMeta{
				Name:   policyAttachmentName(role.Name, policyARN),
				Labels: map[string]string{RoleRefLabelKey: role.Name},
			},
			Spec: v1beta1.RolePolicyAttachmentSpec{
				ForProvider: v1beta1.RolePolicyAttachmentParameters{
					RoleNameRef: &xpv1.Reference{
						Name: role.Name,
					},
					PolicyARN: policyARN,
				},
				ResourceSpec: xpv1.ResourceSpec{
					DeletionPolicy:          xpv1.DeletionDelete,
					ProviderConfigReference: &xpv1.Reference{Name: t.GetProviderConfigName()},
				},
			},
		}

		err := controllerutil.SetOwnerReference(role, attachment, t.Local.Scheme())
		if err != nil {
			return err
		}

		_, err = controllerutil.CreateOrUpdate(ctx, t.Local, attachment, func() error {
			return nil
		})
		if err != nil {
			return err
		}
	}

	return nil
}

// helpers for readiness timestamps

func isReadyConditionTrue(conds []xpv1.Condition) bool {
	for _, c := range conds {
		if c.Type == xpv1.TypeReady && c.Status == v1.ConditionTrue {
			return true
		}
	}
	return false
}

func getReadyTransition(conds []xpv1.Condition) time.Time {
	var latest time.Time
	for _, c := range conds {
		if c.Type == xpv1.TypeReady && c.Status == v1.ConditionTrue {
			if c.LastTransitionTime.IsZero() {
				continue
			}
			if c.LastTransitionTime.Time.After(latest) {
				latest = c.LastTransitionTime.Time
			}
		}
	}
	return latest
}

func latestReadyTransition(attachments []v1beta1.RolePolicyAttachment) time.Time {
	latest := time.Time{}
	for _, a := range attachments {
		aT := getReadyTransition(a.Status.Conditions)
		if aT.After(latest) {
			latest = aT
		}
	}
	return latest
}

func policyAttachmentName(name, policyARN string) string {
	hash := fmt.Sprintf("%x", md5.Sum([]byte(policyARN)))[:8]
	return fmt.Sprintf("%s-%s", name, hash)
}

func policyName(roleName string) string {
	return fmt.Sprintf("%s-node-policy", roleName)
}

const RoleRefLabelKey = "role-reference"