package aks

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"time"

	xpv1 "github.com/crossplane/crossplane-runtime/apis/common/v1"
	"github.com/tidbcloud/infra-provider/regional-server/api/v1alpha1"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/consts"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/controllers"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/controllers/bucket/task"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/deploy"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/k8s"
	"github.com/tidbcloud/infra-provider/regional-server/pkg/util"
	authv1beta1 "github.com/upbound/provider-azure/apis/authorization/v1beta1"
	storagev1beta1 "github.com/upbound/provider-azure/apis/storage/v1beta1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

type RoleAssignment struct {
	*task.BaseTask
}

func NewRoleAssignment(ctx *task.Context) *RoleAssignment {
	return &RoleAssignment{BaseTask: &task.BaseTask{Context: ctx}}
}

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

func (t *RoleAssignment) Create(ctx context.Context) deploy.Result {
	t.BeginAt = util.Pointer(time.Now())

	_, useBYOCManagedIdentity, err := t.Config.BYOC.AKSManagedIdentityConfig()
	if err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeueInterval)
	}
	if useBYOCManagedIdentity {
		return deploy.Done()
	}

	containerName := t.Bucket.Status.Bucket.Name
	container := &storagev1beta1.Container{
		ObjectMeta: metav1.ObjectMeta{
			Name: containerName,
		},
	}
	if err := t.Local.Get(ctx, client.ObjectKeyFromObject(container), container); err != nil {
		return deploy.Error(fmt.Errorf("failed to get storage container: %w", err))
	}
	scope := container.Status.AtProvider.ResourceManagerID
	if scope == nil || *scope == "" {
		return deploy.Panic(fmt.Errorf("container scope is not ready")).WithRequeuAfter(defaultRequeueInterval)
	}

	permittedDataplaneID2PrincipalID, err := t.getPermittedDataplaneID2PrincipalID(ctx)
	if err != nil {
		return deploy.Error(err)
	}

	desiredKeys := map[string]struct{}{}
	for dpid, principalID := range permittedDataplaneID2PrincipalID {
		desiredKeys[dpid] = struct{}{}

		labels := map[string]string{
			permittedDataplaneLabelKey: dpid,
			principalIDLabelKey:        principalID,
			consts.BucketLabelKey:      t.Bucket.Spec.Id,
		}
		labels, _ = getLabelsAndTags(labels, container.Labels, t.Config)

		roleAssignment := &authv1beta1.RoleAssignment{
			ObjectMeta: metav1.ObjectMeta{
				Name:   roleAssignmentName(t.Bucket.Spec.Id, dpid),
				Labels: labels,
			},
			Spec: authv1beta1.RoleAssignmentSpec{
				ResourceSpec: xpv1.ResourceSpec{
					ProviderConfigReference: &xpv1.Reference{
						Name: t.Bucket.Spec.AccountRef.ProviderConfigName,
					},
					ManagementPolicies: controllers.GetResourceManagementPolicies(),
				},
				ForProvider: authv1beta1.RoleAssignmentParameters{
					PrincipalID:        ptr.To(principalID),
					RoleDefinitionName: ptr.To(roleDefinitionBlobData),
					Scope:              scope,
				},
			},
		}
		if err := k8s.Apply(ctx, t.Local, roleAssignment); err != nil {
			return deploy.Error(fmt.Errorf("failed to apply role assignment: %w", err))
		}
	}

	if err := t.deleteStaleRoleAssignments(ctx, desiredKeys); err != nil {
		return deploy.Error(err)
	}

	return deploy.Done()
}

func (t *RoleAssignment) PostCreate(ctx context.Context) deploy.Result {
	_, useBYOCManagedIdentity, err := t.Config.BYOC.AKSManagedIdentityConfig()
	if err != nil {
		return deploy.Panic(err).WithRequeuAfter(defaultRequeueInterval)
	}
	if useBYOCManagedIdentity {
		t.Bucket.Status.Phase = v1alpha1.BucketPhaseReady
		return deploy.Done()
	}

	permittedDataplaneID2PrincipalID, err := t.getPermittedDataplaneID2PrincipalID(ctx)
	if err != nil {
		return deploy.Error(err)
	}

	for dpid := range permittedDataplaneID2PrincipalID {
		roleAssignment := &authv1beta1.RoleAssignment{
			ObjectMeta: metav1.ObjectMeta{
				Name: roleAssignmentName(t.Bucket.Spec.Id, dpid),
			},
		}
		if err := t.Local.Get(ctx, client.ObjectKeyFromObject(roleAssignment), roleAssignment); err != nil {
			return deploy.Error(err)
		}
		if !controllers.IsResourceReady(roleAssignment.Status.ResourceStatus) {
			return deploy.Error(fmt.Errorf("role assignment %s not ready", roleAssignment.Name))
		}
	}

	t.Bucket.Status.Phase = v1alpha1.BucketPhaseReady
	return deploy.Done()
}

func (t *RoleAssignment) Delete(ctx context.Context) deploy.Result {
	roleAssignmentList := &authv1beta1.RoleAssignmentList{}
	listOpts := &client.ListOptions{
		LabelSelector: labels.SelectorFromSet(map[string]string{
			consts.BucketLabelKey: t.Bucket.Spec.Id,
		}),
	}
	if err := t.Local.List(ctx, roleAssignmentList, listOpts); err != nil {
		return deploy.Error(fmt.Errorf("list role assignments error: %w", err))
	}
	for _, roleAssignment := range roleAssignmentList.Items {
		ra := roleAssignment
		if err := controllers.DeleteUpboundResource(ctx, t.Local, &ra); err != nil {
			return deploy.Error(fmt.Errorf("failed to delete role assignment %s: %w", roleAssignment.Name, err))
		}
	}
	return deploy.Done()
}

func (t *RoleAssignment) getPermittedDataplaneID2PrincipalID(ctx context.Context) (map[string]string, error) {
	permittedDataplaneIDs := append([]string{t.Bucket.Spec.DataplaneId}, t.Bucket.Spec.GrantedDataplaneIds...)
	result := map[string]string{}
	for _, dpid := range permittedDataplaneIDs {
		principalID, err := getNodePrincipalID(ctx, t.Local, dpid, t.Config.Namespace)
		if err != nil {
			return nil, fmt.Errorf("get dataplane '%s' principal ids error: %w", dpid, err)
		}
		if principalID == "" {
			continue
		}
		result[dpid] = principalID
	}
	return result, nil
}

func (t *RoleAssignment) deleteStaleRoleAssignments(ctx context.Context, desired map[string]struct{}) error {
	roleAssignmentList := &authv1beta1.RoleAssignmentList{}
	listOpts := &client.ListOptions{
		LabelSelector: labels.SelectorFromSet(map[string]string{
			consts.BucketLabelKey: t.Bucket.Spec.Id,
		}),
	}
	if err := t.Local.List(ctx, roleAssignmentList, listOpts); err != nil {
		return fmt.Errorf("list role assignments error: %w", err)
	}
	for _, roleAssignment := range roleAssignmentList.Items {
		dpID := roleAssignment.Labels[permittedDataplaneLabelKey]
		if _, ok := desired[dpID]; ok {
			continue
		}
		ra := roleAssignment
		if err := controllers.DeleteUpboundResource(ctx, t.Local, &ra); err != nil {
			return fmt.Errorf("failed to delete role assignment %s: %w", roleAssignment.Name, err)
		}
	}
	return nil
}

const (
	maxRoleAssignmentNameLength  = 253
	roleAssignmentNameHashLength = 16
)

func roleAssignmentName(bucketID, dpid string) string {
	hash := sha256.Sum256([]byte(bucketID + ":" + dpid))
	name := fmt.Sprintf("%s.%s", bucketID, dpid)
	if len(name) > maxRoleAssignmentNameLength {
		name = name[:(maxRoleAssignmentNameLength - roleAssignmentNameHashLength)]
		name = fmt.Sprintf("%s.%s", name, hex.EncodeToString(hash[:])[:roleAssignmentNameHashLength])
	}
	return name
}