package azure

import (
	"context"
	"fmt"
	"regexp"

	azerrors "github.com/Azure/azure-sdk-for-go-extensions/pkg/errors"
	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/authorization/armauthorization/v3"
	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/msi/armmsi"
	"github.com/tidbcloud/observability/pkg/common/utils"
	"github.com/tidbcloud/observability/pkg/resourceworker/providers"
	"github.com/tidbcloud/observability/pkg/resourceworker/providers/types"
)

const (
	audienceKey                     = "api://AzureADTokenExchange"
	ErrorCodeParentResourceNotFound = "ParentResourceNotFound"
	ErrorCodeRoleAssignmentNotFound = "RoleAssignmentNotFound"
	ErrorCodeRoleAssignmentExists   = "RoleAssignmentExists"
)

type IAMClient struct {
	rbacCli         *armauthorization.RoleAssignmentsClient
	fedIdentityCli  *armmsi.FederatedIdentityCredentialsClient
	userIdentityCli *armmsi.UserAssignedIdentitiesClient
	resourceGroup   string
	subscription    string
	location        string
}

var _ providers.IAMClientAzure = &IAMClient{}

func NewIAMClient(subscriptionID, resourceGroup, region string, cred azcore.TokenCredential) (*IAMClient, error) {
	rbacFactory, err := armauthorization.NewClientFactory(subscriptionID, cred, nil)
	if err != nil {
		return nil, err
	}

	identityFactory, err := armmsi.NewClientFactory(subscriptionID, cred, nil)
	if err != nil {
		return nil, err
	}

	rbacCli := rbacFactory.NewRoleAssignmentsClient()
	fedIdCli := identityFactory.NewFederatedIdentityCredentialsClient()
	userIdCli := identityFactory.NewUserAssignedIdentitiesClient()

	return &IAMClient{
		rbacCli:         rbacCli,
		fedIdentityCli:  fedIdCli,
		userIdentityCli: userIdCli,
		resourceGroup:   resourceGroup,
		subscription:    subscriptionID,
		location:        region,
	}, nil
}

func (i *IAMClient) CreateUserIdentity(ctx context.Context, name string) (*armmsi.Identity, error) {
	id := armmsi.Identity{
		Location: &i.location,
	}
	resp, err := i.userIdentityCli.CreateOrUpdate(ctx, i.resourceGroup, name, id, nil)
	if err != nil {
		return nil, err
	}

	return &resp.Identity, nil
}

func (i *IAMClient) CreateOrUpdateRoleAssignment(ctx context.Context, storageID, name string, params *armauthorization.RoleAssignmentProperties) (*armauthorization.RoleAssignment, error) {
	// create role assignments with conditions
	// https://learn.microsoft.com/en-us/azure/storage/blobs/storage-auth-abac-examples?tabs=portal-visual-editor#example-read-blobs-in-named-containers-with-a-path
	resp, err := i.rbacCli.Create(ctx, storageID, name, armauthorization.RoleAssignmentCreateParameters{
		Properties: params,
	}, nil)
	if err != nil {
		return nil, err
	}
	return &resp.RoleAssignment, nil
}

// CreateOrUpdateRoleAssignmentForPrincipal will create a role assignment for a principal
// if there is no role assignment for the principal, create the role assignment with uuid generated automaticlly,
// if the role assignment for the principal is present, update it respectively.
func (i *IAMClient) CreateOrUpdateRoleAssignmentForPrincipal(ctx context.Context, storageID string, params *armauthorization.RoleAssignmentProperties) (*armauthorization.RoleAssignment, error) {
	// Check if role assignment already exists for this principal
	existingAssignment, exists, err := i.GetRoleAssignmentByPrincipalID(ctx, storageID, *params.PrincipalID)
	if err != nil {
		return nil, fmt.Errorf("checking existing role assignment: %w", err)
	}

	if exists {
		// Update existing role assignment
		return i.CreateOrUpdateRoleAssignment(ctx, storageID, *existingAssignment.Name, params)
	}

	// Create new role assignment with UUID
	ra, err := i.CreateOrUpdateRoleAssignment(ctx, storageID, utils.GenerateID(), params)
	if err == nil {
		return ra, nil
	}
	if !IsRoleAssignmentExistsErr(err) {
		return nil, err
	}

	// Handle eventual consistency/race: an assignment may be created between list and create.
	existingAssignment, exists, getErr := i.GetRoleAssignmentByPrincipalID(ctx, storageID, *params.PrincipalID)
	if getErr != nil {
		return nil, fmt.Errorf("create role assignment got conflict and querying existing failed: %w", getErr)
	}
	if exists {
		return existingAssignment, nil
	}
	return nil, err
}

func (i *IAMClient) CreateOrUpdateFederatedIdentity(ctx context.Context, identityName string, oidcIssuer string, workloadServiceAccounts []string) error {
	saSet := map[string]struct{}{}
	for _, sa := range workloadServiceAccounts {
		saSet[sa] = struct{}{}
		_, err := i.fedIdentityCli.CreateOrUpdate(ctx, i.resourceGroup, identityName, normalizeFederatedIdentityName(sa), armmsi.FederatedIdentityCredential{
			Properties: &armmsi.FederatedIdentityCredentialProperties{
				Audiences: []*string{to.Ptr(audienceKey)},
				Issuer:    to.Ptr(oidcIssuer),
				Subject:   to.Ptr(sa),
			},
		}, nil)
		if err != nil {
			return err
		}
	}

	pager := i.fedIdentityCli.NewListPager(i.resourceGroup, identityName, nil)
	for pager.More() {
		page, err := pager.NextPage(ctx)
		if err != nil {
			return fmt.Errorf("finding federated identity credential: %w", err)
		}
		for _, cred := range page.Value {
			// delete the federated identities not present in the workloadSA.
			if _, ok := saSet[*cred.Properties.Subject]; !ok {
				_, err := i.fedIdentityCli.Delete(ctx, i.resourceGroup, identityName, *cred.Name, nil)
				if err != nil {
					return fmt.Errorf("deleting federated identity credential %s : %w", *cred.Name, err)
				}
			}
		}
	}
	return nil
}

// The name of a federated identity credentials must be within 3 and 120 characters,
// only contains letters (A-Z, a-z), numbers, hyphens and dashes and must start with a number or letter.
// returns the normalized federated identity name
func normalizeFederatedIdentityName(s string) string {
	// Remove any characters that are not letters, numbers, hyphens or dashes
	reg := regexp.MustCompile("[^a-zA-Z0-9-]+")
	normalized := reg.ReplaceAllString(s, "-")
	// Limit the length to between 3 and 120 characters
	if len(normalized) > 120 {
		return normalized[:120]
	}
	return normalized
}

func (i *IAMClient) GetUserIdentity(ctx context.Context, name string) (*armmsi.Identity, bool, error) {
	resp, err := i.userIdentityCli.Get(ctx, i.resourceGroup, name, nil)
	if err != nil {
		if azerrors.IsNotFoundErr(err) {
			return nil, false, nil
		}
		return nil, false, err
	}
	return &resp.Identity, true, nil
}

func (i *IAMClient) GetWorkloadIdentity(ctx context.Context, identityName string) (*types.AzureWorkloadIdentity, bool, error) {
	identity, exist, err := i.GetUserIdentity(ctx, identityName)
	if !exist || err != nil {
		return nil, exist, err
	}

	result := []*armmsi.FederatedIdentityCredentialProperties{}
	pager := i.fedIdentityCli.NewListPager(i.resourceGroup, identityName, nil)
	for pager.More() {
		page, err := pager.NextPage(ctx)
		if err != nil {
			return nil, exist, err
		}
		for _, cred := range page.Value {
			result = append(result, cred.Properties)
		}
	}

	return &types.AzureWorkloadIdentity{
		Identity:            identity,
		FederatedIdentities: result,
	}, exist, nil
}

func (i *IAMClient) GetRoleAssignment(ctx context.Context, storageID, resourcename string) (*armauthorization.RoleAssignment, bool, error) {
	resp, err := i.rbacCli.Get(ctx, storageID, resourcename, nil)
	if err != nil {
		if azerrors.IsNotFoundErr(err) {
			return nil, false, nil
		}
		return nil, false, err
	}
	return &resp.RoleAssignment, true, err
}

func (i *IAMClient) GetRoleAssignmentByPrincipalID(ctx context.Context, scope, principalID string) (*armauthorization.RoleAssignment, bool, error) {
	filter := fmt.Sprintf("assignedTo('%s')", principalID)
	pager := i.rbacCli.NewListForScopePager(scope, &armauthorization.RoleAssignmentsClientListForScopeOptions{
		Filter: &filter,
	})
	for pager.More() {
		page, err := pager.NextPage(ctx)
		if err != nil {
			return nil, false, err
		}
		for _, roleAssignment := range page.Value {
			if roleAssignment.ID == nil {
				continue
			}
			return roleAssignment, true, nil
		}
	}
	return nil, false, nil
}

func (i *IAMClient) DeleteRoleAssignment(ctx context.Context, storageId, resourcename string) error {
	_, err := i.rbacCli.Delete(ctx, storageId, resourcename, nil)
	if err != nil && !IsNotFoundErr(err) {
		return err
	}
	return nil
}

func (i *IAMClient) DeleteRoleAssignmentByPrincipalID(ctx context.Context, storageId, principalID string) error {
	ra, exists, err := i.GetRoleAssignmentByPrincipalID(ctx, storageId, principalID)
	if err != nil {
		return err
	}
	if !exists {
		return nil
	}
	_, err = i.rbacCli.Delete(ctx, storageId, *ra.Name, nil)
	if err != nil && !IsNotFoundErr(err) {
		return fmt.Errorf("deleting role assignment %s: %w", *ra.Name, err)
	}

	return nil
}

func (i *IAMClient) DeleteUserIdentity(ctx context.Context, name string) error {
	_, err := i.userIdentityCli.Delete(ctx, i.resourceGroup, name, nil)
	if err != nil && !azerrors.IsNotFoundErr(err) {
		return err
	}
	return nil
}

func IsNotFoundErr(err error) bool {
	azErr := azerrors.IsResponseError(err)
	return azErr != nil && (azErr.ErrorCode == azerrors.ResourceNotFound || azErr.ErrorCode == ErrorCodeParentResourceNotFound || azErr.ErrorCode == ErrorCodeRoleAssignmentNotFound)
}

func IsRoleAssignmentExistsErr(err error) bool {
	azErr := azerrors.IsResponseError(err)
	return azErr != nil && azErr.ErrorCode == ErrorCodeRoleAssignmentExists
}