package aks

import (
	"context"
	"fmt"
	"sort"
	"strings"

	xpv1 "github.com/crossplane/crossplane-runtime/apis/common/v1"
	authv1beta1 "github.com/upbound/provider-azure/apis/authorization/v1beta1"
	"github.com/upbound/provider-azure/apis/containerservice/v1beta1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/utils/ptr"

	"github.com/tidbcloud/infra-provider/pkg/deploy"
	"github.com/tidbcloud/infra-provider/pkg/k8s"
	"github.com/tidbcloud/infra-provider/regional-server/api/v1alpha1"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/controllers"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/controllers/appbox/task"
)

type NodeRoleAssignment struct {
	*task.BaseTask
}

type RoleAssignment struct {
	Name               string
	RoleDefinitionName string
	Scope              string
	ProviderConfigName string
	Identity           *string
}

func NewNodeRoleAssignment(ctx *task.Context) *NodeRoleAssignment {
	return &NodeRoleAssignment{
		BaseTask: &task.BaseTask{Context: ctx},
	}
}

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

// VNET role assignment name suffix and format
const (
	vnetRoleAssignmentSuffix = "-vnet-ra"
	// desRoleAssignmentSuffix is Disk Encryption Set role assignment suffix
	desRoleAssignmentSuffix = "-des-ra"
	networkContributorRole  = "Network Contributor"
)

func getVNetRoleAssignmentScope(subscription, resourceGroup, vnetName string) string {
	return fmt.Sprintf("/subscriptions/%s/resourceGroups/%s/providers/Microsoft.Network/virtualNetworks/%s", subscription, resourceGroup, vnetName)
}

func getVNetRoleAssignmentName(aksName string) string {
	return fmt.Sprintf("%s%s", aksName, vnetRoleAssignmentSuffix)
}

// ACR role assignment name suffix and format
const (
	acrRoleAssignmentSuffix = "-acr-ra"
	acrPullRole             = "AcrPull"
)

func getACRRoleAssignmentName(aksName string) string {
	return fmt.Sprintf("%s%s", aksName, acrRoleAssignmentSuffix)
}

func getACRRoleAssigmentScope(subscription, resourceGroup string) string {
	return fmt.Sprintf("/subscriptions/%s/resourceGroups/%s", subscription, resourceGroup)
}

// Disk role assignment name suffix and format
const (
	diskRoleAssignmentSuffix = "-disk-ra"
	vmRoleAssignmentSuffix   = "-vm-ra"
)

func getDiskRoleAssignmentName(aksName string) string {
	return fmt.Sprintf("%s%s", aksName, diskRoleAssignmentSuffix)
}

func getDiskEncryptionSetRoleAssignmentName(aksName string) string {
	return fmt.Sprintf("%s%s", aksName, desRoleAssignmentSuffix)
}

func getDiskRoleAssigmentScope(subscription, resourceGroup string) string {
	return fmt.Sprintf("/subscriptions/%s/resourceGroups/%s", subscription, resourceGroup)
}

func getDiskEncryptionSetRoleAssigmentScope(subscription, resourceGroup string) string {
	return getDiskRoleAssigmentScope(subscription, resourceGroup)
}

func getVMRoleAssignmentName(aksName string) string {
	return fmt.Sprintf("%s%s", aksName, vmRoleAssignmentSuffix)
}

func getVMRoleAssignmentScope(subscription, resourceGroup string) string {
	return fmt.Sprintf("/subscriptions/%s/resourceGroups/%s", subscription, resourceGroup)
}

func (t *NodeRoleAssignment) getNodeResourceGroupName() string {
	return fmt.Sprintf("%s%s", t.AKSCluster.Name, nodeResourceGroupSuffix)
}

// getRoleAssignmentList returns a list of role assignments for the node pool.
// if onlyName is true, only return the role assignment name list, it's used in delete task.
func (t *NodeRoleAssignment) getRoleAssignmentList(ctx context.Context, onlyRoleName bool) ([]RoleAssignment, error) {
	aksName := t.APPBox.Status.KubernetesClusterName
	userAssignedControlPlane := isAKSControlPlaneUserAssignedIdentity(t.AKSCluster)
	if onlyRoleName {
		// used in delete task
		roleAssignmentList := []RoleAssignment{
			{
				Name: getACRRoleAssignmentName(aksName),
			},
			{
				Name: getDiskRoleAssignmentName(aksName),
			},
			{
				Name: getVMRoleAssignmentName(aksName),
			},
		}
		if !userAssignedControlPlane {
			roleAssignmentList = append(roleAssignmentList,
				RoleAssignment{Name: getDiskEncryptionSetRoleAssignmentName(aksName)},
				RoleAssignment{Name: getVNetRoleAssignmentName(aksName)},
			)
		} else {
			roleAssignmentList = append(roleAssignmentList, RoleAssignment{
				Name: getManagedIdentityDESRoleAssignmentName(aksName),
			})
		}
		return roleAssignmentList, nil
	}

	kubeletIdentityObjectId, err := t.getKubeletIdentity(ctx)
	if err != nil {
		return nil, err
	}

	var clusterPrincipalID *string
	if userAssignedControlPlane {
		clusterPrincipalID, err = getControlPlaneManagedIdentityPrincipalID(ctx, t.Local, t.APPBox)
		if err != nil {
			return nil, err
		}
	} else {
		clusterPrincipalID, err = t.getAKSIdentity(ctx)
		if err != nil {
			return nil, err
		}
	}

	// Role assignments bound to the kubelet identity.
	roleAssignmentList := []RoleAssignment{
		{
			Name:               getACRRoleAssignmentName(aksName),
			RoleDefinitionName: "AcrPull",
			Scope: func() string {
				centralSubscrition := t.Config.ProviderConfig.AccountID
				acrResourceGroup := t.Config.ACRResourceGroup
				return getACRRoleAssigmentScope(centralSubscrition, acrResourceGroup)
			}(),
			ProviderConfigName: t.Config.ProviderConfig.CrossplaneProvider,
			Identity:           kubeletIdentityObjectId,
		},
		{
			Name: getDiskRoleAssignmentName(aksName),
			// Disk Operator for Managed Disk role definition doesn't work here.
			RoleDefinitionName: "Disk Pool Operator",
			Scope: func() string {
				accountID := t.VPC.Spec.AccountRef.AccountID
				resourceGroup := t.getNodeResourceGroupName()
				return getDiskRoleAssigmentScope(accountID, resourceGroup)
			}(),
			ProviderConfigName: t.ProviderConfigName(),
			Identity:           kubeletIdentityObjectId,
		},
		{
			Name:               getVMRoleAssignmentName(aksName),
			RoleDefinitionName: "Virtual Machine Contributor",
			Scope: func() string {
				accountID := t.VPC.Spec.AccountRef.AccountID
				resourceGroup := t.getNodeResourceGroupName()
				return getVMRoleAssignmentScope(accountID, resourceGroup)
			}(),
			ProviderConfigName: t.ProviderConfigName(),
			Identity:           kubeletIdentityObjectId,
		},
	}

	// Role assignments bound to the control plane identity.
	if userAssignedControlPlane {
		// Required for CMEK scenarios where the csi-driver needs to access disk encryption sets.
		roleAssignmentList = append(roleAssignmentList, controlPlaneDESRoleAssignment(
			aksName,
			t.VPC.Spec.AccountRef.AccountID,
			t.getNodeResourceGroupName(),
			t.ProviderConfigName(),
			clusterPrincipalID,
		))
	} else {
		roleAssignmentList = append(roleAssignmentList,
			RoleAssignment{
				Name:               getDiskEncryptionSetRoleAssignmentName(aksName),
				RoleDefinitionName: diskEncryptionSetOperatorRole,
				Scope: func() string {
					subscriptionID := t.VPC.Spec.AccountRef.AccountID
					resourceGroup := t.getNodeResourceGroupName()
					return getDiskEncryptionSetRoleAssigmentScope(subscriptionID, resourceGroup)
				}(),
				ProviderConfigName: t.ProviderConfigName(),
				Identity:           clusterPrincipalID,
			},
			RoleAssignment{
				Name:               getVNetRoleAssignmentName(aksName),
				RoleDefinitionName: "Network Contributor",
				Scope: func() string {
					subscriptionID := t.VPC.Spec.AccountRef.AccountID
					vnetName := t.VPC.Status.VPC.Name
					resourceGroupName := controllers.GetResourceGroupName(t.VPC)
					return getVNetRoleAssignmentScope(subscriptionID, resourceGroupName, vnetName)
				}(),
				ProviderConfigName: t.ProviderConfigName(),
				Identity:           clusterPrincipalID,
			},
		)
	}

	return roleAssignmentList, nil
}

func (t *NodeRoleAssignment) Create(ctx context.Context) deploy.Result {
	roleAssignments, err := t.getRoleAssignmentList(ctx, false)

	if err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}

	for _, ra := range roleAssignments {
		roleAssignment := t.newRoleAssignmentObject(ra)
		if err := k8s.Apply(ctx, t.Local, roleAssignment); err != nil {
			return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
		}
	}
	return deploy.Done()
}

func (t *NodeRoleAssignment) newRoleAssignmentObject(ra RoleAssignment) *authv1beta1.RoleAssignment {
	if strings.HasSuffix(ra.Name, managedIdentityDESRoleAssignmentSuffix) {
		return newManagedIdentityRoleAssignmentObject(t.APPBox.Name, ra)
	}
	return &authv1beta1.RoleAssignment{
		ObjectMeta: metav1.ObjectMeta{
			Name: ra.Name,
			Labels: map[string]string{
				v1alpha1.APPBoxNameKey: t.APPBox.Name,
			},
		},
		Spec: authv1beta1.RoleAssignmentSpec{
			ResourceSpec: xpv1.ResourceSpec{
				ProviderConfigReference: &xpv1.Reference{
					Name: ra.ProviderConfigName,
				},
			},
			ForProvider: authv1beta1.RoleAssignmentParameters{
				PrincipalID:        ra.Identity,
				RoleDefinitionName: ptr.To(ra.RoleDefinitionName),
				Scope:              ptr.To(ra.Scope),
			},
		},
	}
}

func (t *NodeRoleAssignment) getKubeletIdentity(ctx context.Context) (*string, error) {
	aks := t.AKSCluster
	if aks == nil {
		return nil, fmt.Errorf("aks cluster not found")
	}
	if len(aks.Status.AtProvider.KubeletIdentity) == 0 {
		return nil, fmt.Errorf("kubeletIdentity.objectId not found")
	}
	return aks.Status.AtProvider.KubeletIdentity[0].ObjectID, nil
}

func (t *NodeRoleAssignment) getAKSIdentity(ctx context.Context) (*string, error) {
	aks := t.AKSCluster
	if aks == nil {
		return nil, fmt.Errorf("aks cluster not found")
	}
	if len(aks.Status.AtProvider.Identity) == 0 {
		return nil, fmt.Errorf("no identity found in aks %s", aks.Name)
	}
	principalID := aks.Status.AtProvider.Identity[0].PrincipalID
	if principalID == nil || *principalID == "" {
		return nil, fmt.Errorf("aks %s system-assigned principal ID is not ready", aks.Name)
	}
	return principalID, nil
}

func isAKSControlPlaneUserAssignedIdentity(cluster *v1beta1.KubernetesCluster) bool {
	if cluster == nil {
		return false
	}
	for _, identity := range cluster.Spec.ForProvider.Identity {
		if ptr.Deref(identity.Type, "") == UserAssignedIdentityType {
			return true
		}
	}
	return false
}

func (t *NodeRoleAssignment) PostCreate(ctx context.Context) (res deploy.Result) {
	roleAssignments, err := t.getRoleAssignmentList(ctx, false)
	if err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}

	roleAssignmentStatus := []v1alpha1.RoleAssignmentStatus{}
	for _, ra := range roleAssignments {
		key := types.NamespacedName{Name: ra.Name}
		roleAssignment := &authv1beta1.RoleAssignment{}
		err := t.Local.Get(ctx, key, roleAssignment)
		if err != nil {
			return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
		}

		condition := roleAssignment.Status.ConditionedStatus.GetCondition(xpv1.TypeReady)
		if condition.Status != corev1.ConditionTrue {
			return deploy.Panic(fmt.Errorf("role assignment status not ready: %v", condition.Reason)).WithRequeuAfter(defaultRequeuInterval)
		}

		roleAssignmentStatus = append(roleAssignmentStatus, v1alpha1.RoleAssignmentStatus{
			Name:               ra.Name,
			ID:                 ptr.Deref(roleAssignment.Status.AtProvider.ID, ""),
			PrincipalID:        ptr.Deref(roleAssignment.Spec.ForProvider.PrincipalID, ""),
			RoleDefinitionName: ptr.Deref(roleAssignment.Spec.ForProvider.RoleDefinitionName, ""),
			Scope:              ptr.Deref(roleAssignment.Spec.ForProvider.Scope, ""),
		})
	}

	// sort the role assignment status by name
	t.APPBox.Status.NodeRoleAssignments = func() []v1alpha1.RoleAssignmentStatus {
		sort.Slice(roleAssignmentStatus, func(i, j int) bool {
			return roleAssignmentStatus[i].Name < roleAssignmentStatus[j].Name
		})
		return roleAssignmentStatus
	}()

	return deploy.PanicOrDone(t.APPBoxRW.UpdateStatus(ctx, t.APPBox))
}

func (t *NodeRoleAssignment) Delete(ctx context.Context) deploy.Result {
	roleAssignments, err := t.getRoleAssignmentList(ctx, true)
	if err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}

	for _, ra := range roleAssignments {
		roleAssignment := &authv1beta1.RoleAssignment{
			ObjectMeta: metav1.ObjectMeta{
				Name: ra.Name,
			},
		}
		if err := k8s.DeleteObject(ctx, t.Local, roleAssignment); err != nil {
			return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
		}
	}

	return deploy.Done()
}

func (t *NodeRoleAssignment) PostDelete(ctx context.Context) deploy.Result {
	return deploy.Done()
}