package aks

import (
	"context"
	"fmt"
	"time"

	xpv1 "github.com/crossplane/crossplane-runtime/apis/common/v1"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	authv1beta1 "github.com/upbound/provider-azure/apis/authorization/v1beta1"
	managedidentityv1beta1 "github.com/upbound/provider-azure/apis/managedidentity/v1beta1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"

	"github.com/tidbcloud/infra-provider/pkg/deploy"
	"github.com/tidbcloud/infra-provider/pkg/k8s"
	"github.com/tidbcloud/infra-provider/pkg/util"
	"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"
)

const (
	managedIdentityVNetRoleAssignmentSuffix = "-cp-mi-vnet-ra"
	managedIdentityDESRoleAssignmentSuffix  = "-cp-mi-des-ra"
	managedIdentityRoleAssignmentLabel      = "cloud.pingcap.com/aks-managed-identity-role-assignment"
	managedIdentityRoleAssignmentValue      = "true"
	servicePrincipalType                    = "ServicePrincipal"
	diskEncryptionSetOperatorRole           = "Disk Encryption Set Operator for Managed Disks"
)

type AKSManagedIdentity struct {
	*task.BaseTask
}

func NewAKSManagedIdentity(ctx *task.Context) *AKSManagedIdentity {
	return &AKSManagedIdentity{
		BaseTask: &task.BaseTask{Context: ctx},
	}
}

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

func (t *AKSManagedIdentity) Create(ctx context.Context) deploy.Result {
	t.BeginAt = ptr.To(time.Now())

	enable, err := t.shouldCreateManagedIdentity(ctx)
	if err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}
	if !enable {
		return deploy.Done()
	}

	if err := k8s.Apply(ctx, t.Local, t.newControlPlaneIdentity()); err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}
	return deploy.Done()
}

func (t *AKSManagedIdentity) PostCreate(ctx context.Context) deploy.Result {
	enable, err := t.shouldCreateManagedIdentity(ctx)
	if err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}
	if !enable {
		return deploy.Done()
	}

	identity := &managedidentityv1beta1.UserAssignedIdentity{}
	if err := t.Local.Get(ctx, types.NamespacedName{Name: aksControlPlaneManagedIdentityName(t.APPBox)}, identity); err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}

	condition := identity.Status.GetCondition(xpv1.TypeReady)
	if condition.Type != xpv1.TypeReady || condition.Status != corev1.ConditionTrue {
		return deploy.Panic(fmt.Errorf("managed identity %s is not ready: %s", identity.Name, condition.Reason)).WithRequeuAfter(defaultRequeuInterval)
	}
	if identity.Status.AtProvider.PrincipalID == nil || *identity.Status.AtProvider.PrincipalID == "" {
		return deploy.Panic(fmt.Errorf("managed identity %s principal ID is not ready", identity.Name)).WithRequeuAfter(defaultRequeuInterval)
	}

	roleAssignments, err := t.preAKSRoleAssignments(identity.Status.AtProvider.PrincipalID)
	if err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}

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

	for _, ra := range roleAssignments {
		roleAssignment := &authv1beta1.RoleAssignment{}
		if err := t.Local.Get(ctx, types.NamespacedName{Name: ra.Name}, roleAssignment); err != nil {
			return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
		}
		condition := roleAssignment.Status.GetCondition(xpv1.TypeReady)
		if condition.Type != xpv1.TypeReady || condition.Status != corev1.ConditionTrue {
			return deploy.Panic(fmt.Errorf("role assignment %s is not ready: %s", ra.Name, condition.Reason)).WithRequeuAfter(defaultRequeuInterval)
		}
	}

	return deploy.Done()
}

func (t *AKSManagedIdentity) Delete(ctx context.Context) deploy.Result {
	if err := t.deleteRoleAssignmentsBySelector(ctx); err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}

	return deploy.PanicOrDone(k8s.DeleteObject(ctx, t.Local, &managedidentityv1beta1.UserAssignedIdentity{
		ObjectMeta: metav1.ObjectMeta{Name: aksControlPlaneManagedIdentityName(t.APPBox)},
	}))
}

func (t *AKSManagedIdentity) PostDelete(ctx context.Context) deploy.Result {
	roleAssignments, err := t.managedIdentityRoleAssignments(ctx)
	if err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}
	if len(roleAssignments.Items) > 0 {
		return deploy.Panic(fmt.Errorf("managed identity role assignments are still present: %d", len(roleAssignments.Items))).WithRequeuAfter(defaultRequeuInterval)
	}

	if err := k8s.GetDeletedObject(ctx, t.Local, types.NamespacedName{Name: aksControlPlaneManagedIdentityName(t.APPBox)}, &managedidentityv1beta1.UserAssignedIdentity{}); err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeuInterval)
	}
	return deploy.Done()
}

func (t *AKSManagedIdentity) shouldCreateManagedIdentity(ctx context.Context) (bool, error) {
	cluster, exists, err := getExistingAKS(ctx, t.Local, aksClusterName(t.APPBox))
	if err != nil {
		return false, err
	}
	return shouldEnableAPIServerVNetIntegration(t.APPBox, cluster, exists), nil
}

func (t *AKSManagedIdentity) deleteRoleAssignmentsBySelector(ctx context.Context) error {
	roleAssignments, err := t.managedIdentityRoleAssignments(ctx)
	if err != nil {
		return err
	}
	for i := range roleAssignments.Items {
		if err := k8s.DeleteObject(ctx, t.Local, &roleAssignments.Items[i]); err != nil {
			return err
		}
	}
	return nil
}

func (t *AKSManagedIdentity) managedIdentityRoleAssignments(ctx context.Context) (*authv1beta1.RoleAssignmentList, error) {
	roleAssignments := &authv1beta1.RoleAssignmentList{}
	if err := t.Local.List(ctx, roleAssignments, t.managedIdentityRoleAssignmentSelector()); err != nil {
		return nil, err
	}
	return roleAssignments, nil
}

func (t *AKSManagedIdentity) managedIdentityRoleAssignmentSelector() client.MatchingLabels {
	return client.MatchingLabels{
		v1alpha1.APPBoxNameKey:             t.APPBox.Name,
		managedIdentityRoleAssignmentLabel: managedIdentityRoleAssignmentValue,
	}
}

func (t *AKSManagedIdentity) newControlPlaneIdentity() *managedidentityv1beta1.UserAssignedIdentity {
	tags := cloudv1.TagsFromLabels(t.APPBox.Labels)
	tags = util.AddResourceName(tags, aksControlPlaneManagedIdentityName(t.APPBox))

	return &managedidentityv1beta1.UserAssignedIdentity{
		ObjectMeta: metav1.ObjectMeta{
			Name: aksControlPlaneManagedIdentityName(t.APPBox),
			Labels: map[string]string{
				v1alpha1.APPBoxNameKey: t.APPBox.Name,
			},
		},
		Spec: managedidentityv1beta1.UserAssignedIdentitySpec{
			ResourceSpec: xpv1.ResourceSpec{
				ProviderConfigReference: &xpv1.Reference{
					Name: t.ProviderConfigName(),
				},
			},
			ForProvider: managedidentityv1beta1.UserAssignedIdentityParameters{
				Name:              ptr.To(aksControlPlaneManagedIdentityName(t.APPBox)),
				Location:          ptr.To(t.APPBox.Spec.Region),
				ResourceGroupName: ptr.To(controllers.GetResourceGroupName(t.VPC)),
				Tags:              util.ToMapPointer(tags),
			},
		},
	}
}

func (t *AKSManagedIdentity) preAKSRoleAssignments(principalID *string) ([]RoleAssignment, error) {
	if t.VPC == nil || t.VPC.Spec.AccountRef == nil {
		return nil, fmt.Errorf("azure account reference is required for AKS managed identity")
	}
	vnetName := t.VPC.Status.VPC.Name
	if vnetName == "" {
		return nil, fmt.Errorf("vnet name is not ready")
	}

	aksName := aksClusterName(t.APPBox)
	subscriptionID := t.VPC.Spec.AccountRef.AccountID
	resourceGroupName := controllers.GetResourceGroupName(t.VPC)

	return []RoleAssignment{
		{
			Name:               getManagedIdentityVNetRoleAssignmentName(aksName),
			RoleDefinitionName: networkContributorRole,
			Scope:              getVNetRoleAssignmentScope(subscriptionID, resourceGroupName, vnetName),
			ProviderConfigName: t.ProviderConfigName(),
			Identity:           principalID,
		},
	}, nil
}

func getManagedIdentityVNetRoleAssignmentName(aksName string) string {
	return fmt.Sprintf("%s%s", aksName, managedIdentityVNetRoleAssignmentSuffix)
}

func getManagedIdentityDESRoleAssignmentName(aksName string) string {
	return fmt.Sprintf("%s%s", aksName, managedIdentityDESRoleAssignmentSuffix)
}

func controlPlaneDESRoleAssignment(aksName, subscriptionID, nodeResourceGroup, providerConfig string, principalID *string) RoleAssignment {
	return RoleAssignment{
		Name:               getManagedIdentityDESRoleAssignmentName(aksName),
		RoleDefinitionName: diskEncryptionSetOperatorRole,
		Scope:              getDiskEncryptionSetRoleAssigmentScope(subscriptionID, nodeResourceGroup),
		ProviderConfigName: providerConfig,
		Identity:           principalID,
	}
}

func getControlPlaneManagedIdentityPrincipalID(ctx context.Context, cli client.Client, appBox *v1alpha1.APPBox) (*string, error) {
	identity := &managedidentityv1beta1.UserAssignedIdentity{}
	name := aksControlPlaneManagedIdentityName(appBox)
	if err := cli.Get(ctx, types.NamespacedName{Name: name}, identity); err != nil {
		return nil, fmt.Errorf("get control plane managed identity %s: %w", name, err)
	}
	principalID := identity.Status.AtProvider.PrincipalID
	if principalID == nil || *principalID == "" {
		return nil, fmt.Errorf("control plane managed identity %s principal ID is not ready", name)
	}
	return principalID, nil
}

func newManagedIdentityRoleAssignmentObject(appBoxName string, ra RoleAssignment) *authv1beta1.RoleAssignment {
	return &authv1beta1.RoleAssignment{
		ObjectMeta: metav1.ObjectMeta{
			Name: ra.Name,
			Labels: map[string]string{
				v1alpha1.APPBoxNameKey:             appBoxName,
				managedIdentityRoleAssignmentLabel: managedIdentityRoleAssignmentValue,
			},
		},
		Spec: authv1beta1.RoleAssignmentSpec{
			ResourceSpec: xpv1.ResourceSpec{
				ProviderConfigReference: &xpv1.Reference{
					Name: ra.ProviderConfigName,
				},
			},
			ForProvider: authv1beta1.RoleAssignmentParameters{
				PrincipalID:                  ra.Identity,
				PrincipalType:                ptr.To(servicePrincipalType),
				RoleDefinitionName:           ptr.To(ra.RoleDefinitionName),
				Scope:                        ptr.To(ra.Scope),
				SkipServicePrincipalAADCheck: ptr.To(true),
			},
		},
	}
}