package azure

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"net/http"
	"testing"

	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/blob"
	"github.com/stretchr/testify/assert"

	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
	"github.com/Azure/azure-sdk-for-go/sdk/azcore/arm"
	azfake "github.com/Azure/azure-sdk-for-go/sdk/azcore/fake"
	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/storage/armstorage"
	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/storage/armstorage/fake"
	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob"
	. "github.com/onsi/gomega"
	azmock "github.com/tidbcloud/observability/test/mockresourceworker/providers/azure"
	"go.uber.org/mock/gomock"

	"github.com/tidbcloud/observability/pkg/common/message"
	"github.com/tidbcloud/observability/pkg/resourceworker/providers/types"
)

var errBoom = errors.New("fake")

func TestNewBlobStorageClient(t *testing.T) {
	g := NewWithT(t)
	_, err := NewBlobStorageClient("fake", "fake", "fake", &azfake.TokenCredential{})
	g.Expect(err).To(Succeed())

	// Test case for NewClientFactory error
	originalNewClientFactory := newClientFactory
	newClientFactory = func(subscriptionID string, cred azcore.TokenCredential, options *arm.ClientOptions) (*armstorage.ClientFactory, error) {
		return nil, errBoom
	}
	_, err = NewBlobStorageClient("fake", "fake", "fake", &azfake.TokenCredential{})
	assert.Error(t, err)
	newClientFactory = originalNewClientFactory

	// Test case for azblob.NewClient error
	originalNewBlobClient := newBlobClient
	newBlobClient = func(credentialURL string, cred azcore.TokenCredential, options *azblob.ClientOptions) (*azblob.Client, error) {
		return nil, errBoom
	}
	_, err = NewBlobStorageClient("fake", "fake", "fake", &azfake.TokenCredential{})
	g.Expect(err).To(HaveOccurred())
	g.Expect(err).To(MatchError(errBoom))
	newBlobClient = originalNewBlobClient
}

func TestCreateBucket(t *testing.T) {
	g := NewWithT(t)
	fServer := fake.BlobContainersServer{
		Create: func(ctx context.Context, resourceGroupName string, storageAccountName string, blobContainerName string, blobContainer armstorage.BlobContainer, options *armstorage.BlobContainersClientCreateOptions) (resp azfake.Responder[armstorage.BlobContainersClientCreateResponse], errResp azfake.ErrorResponder) {
			resp.SetResponse(http.StatusOK, armstorage.BlobContainersClientCreateResponse{}, nil)
			return
		},
	}
	cli, err := armstorage.NewBlobContainersClient("fake", &azfake.TokenCredential{}, &arm.ClientOptions{
		ClientOptions: azcore.ClientOptions{
			Transport: fake.NewBlobContainersServerTransport(&fServer),
		},
	})
	g.Expect(err).To(Succeed())
	bsCli := &BlobStorageClient{
		cli:            cli,
		resourceGroup:  "fake",
		storageAccount: "fake",
	}
	g.Expect(bsCli.CreateBucket(context.TODO(), "fake")).To(Succeed())

	// Test case for Create error
	fServerWithError := fake.BlobContainersServer{
		Create: func(ctx context.Context, resourceGroupName string, storageAccountName string, blobContainerName string, blobContainer armstorage.BlobContainer, options *armstorage.BlobContainersClientCreateOptions) (resp azfake.Responder[armstorage.BlobContainersClientCreateResponse], errResp azfake.ErrorResponder) {
			errResp.SetError(errBoom)
			return
		},
	}
	cliWithError, err := armstorage.NewBlobContainersClient("fake", &azfake.TokenCredential{}, &arm.ClientOptions{
		ClientOptions: azcore.ClientOptions{
			Transport: fake.NewBlobContainersServerTransport(&fServerWithError),
		},
	})
	g.Expect(err).To(Succeed())
	bsCliWithError := &BlobStorageClient{
		cli:            cliWithError,
		resourceGroup:  "fake",
		storageAccount: "fake",
	}
	g.Expect(bsCliWithError.CreateBucket(context.TODO(), "fake")).To(HaveOccurred())
	g.Expect(bsCliWithError.CreateBucket(context.TODO(), "fake")).To(MatchError(errBoom))
}

func TestBucketExists(t *testing.T) {
	g := NewWithT(t)

	f := func(name string, server *fake.BlobContainersServer, expErr error, expExist bool) {
		t.Run(name, func(t *testing.T) {
			cli, err := armstorage.NewBlobContainersClient("fake", &azfake.TokenCredential{}, &arm.ClientOptions{
				ClientOptions: azcore.ClientOptions{
					Transport: fake.NewBlobContainersServerTransport(server),
				},
			})
			g.Expect(err).To(Succeed())
			b := &BlobStorageClient{
				cli:            cli,
				resourceGroup:  "fake",
				storageAccount: "fake",
			}
			exist, err := b.BucketExists(context.TODO(), "fake")
			g.Expect(errors.Is(err, expErr)).To(BeTrue())
			g.Expect(exist == expExist).To(BeTrue())
		})
	}

	f("error when calling bucket exists", &fake.BlobContainersServer{
		Get: func(ctx context.Context, resourceGroupName, accountName, containerName string, options *armstorage.BlobContainersClientGetOptions) (resp azfake.Responder[armstorage.BlobContainersClientGetResponse], errResp azfake.ErrorResponder) {
			errResp.SetError(errBoom)
			return
		},
	}, errBoom, false)

	f("bucket exist when calling bucket exist", &fake.BlobContainersServer{
		Get: func(ctx context.Context, resourceGroupName string, storageAccountName string, blobContainerName string, options *armstorage.BlobContainersClientGetOptions) (resp azfake.Responder[armstorage.BlobContainersClientGetResponse], errResp azfake.ErrorResponder) {
			resp.SetResponse(http.StatusOK, armstorage.BlobContainersClientGetResponse{}, nil)
			return
		},
	}, nil, true)

}

func TestGet(t *testing.T) {
	g := NewWithT(t)
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()
	mockCli := azmock.NewMockblobDataClient(ctrl)
	bsCli := &BlobStorageClient{
		dataCli: mockCli,
	}
	t.Run("success", func(t *testing.T) {
		mockCli.EXPECT().DownloadStream(context.TODO(), "fake", "fake", gomock.Any()).Return(azblob.DownloadStreamResponse{
			DownloadResponse: blob.DownloadResponse{
				Body: io.NopCloser(bytes.NewBufferString("fake content")),
			},
		}, nil).Times(1)
		reader, err := bsCli.Get(context.TODO(), "fake", "fake")
		g.Expect(err).To(Succeed())
		g.Expect(reader).ToNot(BeNil())

		content, _ := io.ReadAll(*reader)
		g.Expect(string(content)).To(Equal("fake content"))
	})

	t.Run("failed", func(t *testing.T) {
		mockCli.EXPECT().DownloadStream(context.TODO(), "fake", "fake", gomock.Any()).Return(azblob.DownloadStreamResponse{}, errBoom).Times(1)
		_, err := bsCli.Get(context.TODO(), "fake", "fake")
		g.Expect(err).To(HaveOccurred())
	})
}

func TestListObjects(t *testing.T) {
	g := NewWithT(t)
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()
	mockCli := azmock.NewMockblobDataClient(ctrl)
	bsCli := &BlobStorageClient{
		dataCli: mockCli,
	}

	t.Run("failed", func(t *testing.T) {
		mockCli.EXPECT().NewListBlobsFlatPager("fakeBucket", gomock.Any()).Return(runtime.NewPager[azblob.ListBlobsFlatResponse](runtime.PagingHandler[azblob.ListBlobsFlatResponse]{
			Fetcher: func(ctx context.Context, _ *azblob.ListBlobsFlatResponse) (azblob.ListBlobsFlatResponse, error) {
				return azblob.ListBlobsFlatResponse{}, errors.New("failed")
			},
		}))

		_, err := bsCli.ListObjects(context.TODO(), "fakeBucket", "prefix", types.ListOptions{})
		g.Expect(err).To(HaveOccurred())
	})
}

func TestPut(t *testing.T) {
	g := NewWithT(t)
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()
	mockCli := azmock.NewMockblobDataClient(ctrl)
	mockCli.EXPECT().UploadStream(context.TODO(), "fake", "fake", gomock.Any(), gomock.Any()).Return(azblob.UploadStreamResponse{}, nil)
	bsCli := &BlobStorageClient{
		dataCli: mockCli,
	}
	g.Expect(bsCli.Put(context.TODO(), "fake", "fake", bytes.NewReader([]byte("fake")))).To(Succeed())

	// Test case for UploadStream error
	mockCliWithError := azmock.NewMockblobDataClient(ctrl)
	mockCliWithError.EXPECT().UploadStream(context.TODO(), "fake", "fake", gomock.Any(), gomock.Any()).Return(azblob.UploadStreamResponse{}, errBoom).AnyTimes()
	bsCliWithError := &BlobStorageClient{
		dataCli: mockCliWithError,
	}
	g.Expect(bsCliWithError.Put(context.TODO(), "fake", "fake", bytes.NewReader([]byte("fake")))).To(HaveOccurred())
	g.Expect(bsCliWithError.Put(context.TODO(), "fake", "fake", bytes.NewReader([]byte("fake")))).To(MatchError(errBoom))
}

func TestDelete(t *testing.T) {
	g := NewWithT(t)
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()
	mockCli := azmock.NewMockblobDataClient(ctrl)
	mockCli.EXPECT().DeleteBlob(context.TODO(), "fake", "fake", gomock.Any()).Return(azblob.DeleteBlobResponse{}, nil).Times(1)
	bsCli := &BlobStorageClient{
		dataCli: mockCli,
	}
	g.Expect(bsCli.Delete(context.TODO(), "fake", "fake")).To(Succeed())

	t.Run("multiple objects", func(t *testing.T) {
		mockCliMulti := azmock.NewMockblobDataClient(ctrl)
		mockCliMulti.EXPECT().DeleteBlob(context.TODO(), "fake", "object1", gomock.Any()).Return(azblob.DeleteBlobResponse{}, nil).Times(1)
		mockCliMulti.EXPECT().DeleteBlob(context.TODO(), "fake", "object2", gomock.Any()).Return(azblob.DeleteBlobResponse{}, nil).Times(1)
		bsCliMulti := &BlobStorageClient{
			dataCli: mockCliMulti,
		}
		g.Expect(bsCliMulti.Delete(context.TODO(), "fake", "object1", "object2")).To(Succeed())
	})

	t.Run("delete error", func(t *testing.T) {
		mockCliError := azmock.NewMockblobDataClient(ctrl)
		mockCliError.EXPECT().DeleteBlob(context.TODO(), "fake", "object1", gomock.Any()).Return(azblob.DeleteBlobResponse{}, errBoom).AnyTimes()
		bsCliError := &BlobStorageClient{
			dataCli: mockCliError,
		}
		g.Expect(bsCliError.Delete(context.TODO(), "fake", "object1")).To(HaveOccurred())
		g.Expect(bsCliError.Delete(context.TODO(), "fake", "object1")).To(MatchError(errBoom))
	})
}

func TestBuildAzureLifecycleRules(t *testing.T) {
	g := NewWithT(t)

	// Test case: ExpirationDays only
	param1 := message.LifecycleRuleParam{
		RuleName:       "expire_rule",
		Prefix:         "logs/",
		ExpirationDays: 30,
		ColdDays:       0,
	}
	rules1 := buildAzureLifecycleRules("mybucket", param1)
	g.Expect(rules1).To(HaveLen(1))
	g.Expect(*rules1[0].Name).To(Equal("expire_rule_delete"))
	g.Expect(*rules1[0].Enabled).To(BeTrue())
	g.Expect(*rules1[0].Definition.Filters.PrefixMatch[0]).To(Equal("mybucket/logs/"))
	g.Expect(*rules1[0].Definition.Actions.BaseBlob.Delete.DaysAfterModificationGreaterThan).To(Equal(float32(30)))
	g.Expect(rules1[0].Definition.Actions.BaseBlob.TierToCool).To(BeNil())

	// Test case: ColdDays only
	param2 := message.LifecycleRuleParam{
		RuleName:       "cold_rule",
		Prefix:         "archive",
		ExpirationDays: 0,
		ColdDays:       90,
	}
	rules2 := buildAzureLifecycleRules("mybucket", param2)
	g.Expect(rules2).To(HaveLen(1))
	g.Expect(*rules2[0].Name).To(Equal("cold_rule_cool"))
	g.Expect(*rules2[0].Enabled).To(BeTrue())
	g.Expect(*rules2[0].Definition.Filters.PrefixMatch[0]).To(Equal("mybucket/archive"))
	g.Expect(*rules2[0].Definition.Actions.BaseBlob.TierToCool.DaysAfterModificationGreaterThan).To(Equal(float32(90)))
	g.Expect(rules2[0].Definition.Actions.BaseBlob.Delete).To(BeNil())

	// Test case: Both ExpirationDays and ColdDays
	param3 := message.LifecycleRuleParam{
		RuleName:       "full_rule",
		Prefix:         "",
		ExpirationDays: 7,
		ColdDays:       30,
	}
	rules3 := buildAzureLifecycleRules("mybucket", param3)
	g.Expect(rules3).To(HaveLen(2))
	g.Expect(*rules3[0].Name).To(Equal("full_rule_delete"))
	g.Expect(*rules3[0].Definition.Actions.BaseBlob.Delete.DaysAfterModificationGreaterThan).To(Equal(float32(7)))
	g.Expect(*rules3[1].Name).To(Equal("full_rule_cool"))
	g.Expect(*rules3[1].Definition.Actions.BaseBlob.TierToCool.DaysAfterModificationGreaterThan).To(Equal(float32(30)))
	g.Expect(*rules3[0].Definition.Filters.PrefixMatch[0]).To(Equal("mybucket"))
	g.Expect(*rules3[1].Definition.Filters.PrefixMatch[0]).To(Equal("mybucket"))

	// Test case: No prefix
	param4 := message.LifecycleRuleParam{
		RuleName:       "no_prefix",
		ExpirationDays: 10,
	}
	rules4 := buildAzureLifecycleRules("mybucket", param4)
	g.Expect(rules4).To(HaveLen(1))
	g.Expect(*rules4[0].Definition.Filters.PrefixMatch[0]).To(Equal("mybucket"))

	// Test case: Prefix with leading slash
	param5 := message.LifecycleRuleParam{
		RuleName:       "leading_slash",
		Prefix:         "/path/to/data",
		ExpirationDays: 10,
	}
	rules5 := buildAzureLifecycleRules("mybucket", param5)
	g.Expect(rules5).To(HaveLen(1))
	g.Expect(*rules5[0].Definition.Filters.PrefixMatch[0]).To(Equal("mybucket/path/to/data"))

	// Test case: Prefix with trailing slash
	param6 := message.LifecycleRuleParam{
		RuleName:       "trailing_slash",
		Prefix:         "path/to/data/",
		ExpirationDays: 10,
	}
	rules6 := buildAzureLifecycleRules("mybucket", param6)
	g.Expect(rules6).To(HaveLen(1))
	g.Expect(*rules6[0].Definition.Filters.PrefixMatch[0]).To(Equal("mybucket/path/to/data/"))
}

func TestAddLifecycle(t *testing.T) {
	g := NewWithT(t)
	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	bsCli := &BlobStorageClient{
		resourceGroup:  "fakeResourceGroup",
		storageAccount: "fakeStorageAccount",
	}

	t.Run("add new rules when no existing policy", func(t *testing.T) {
		mockManagementPoliciesClient := azmock.NewMockmanagementCli(ctrl)
		mockManagementPoliciesClient.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(armstorage.ManagementPoliciesClientGetResponse{}, &azcore.ResponseError{StatusCode: http.StatusNotFound, RawResponse: &http.Response{Status: "404 Not Found"}})
		mockManagementPoliciesClient.EXPECT().CreateOrUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
			func(_ context.Context, _ string, _ string, _ armstorage.ManagementPolicyName, policy armstorage.ManagementPolicy, _ *armstorage.ManagementPoliciesClientCreateOrUpdateOptions) (armstorage.ManagementPoliciesClientCreateOrUpdateResponse, error) {
				g.Expect(policy.Properties.Policy.Rules).To(HaveLen(2)) // one for delete, one for cool
				g.Expect(*policy.Properties.Policy.Rules[0].Name).To(ContainSubstring("testRule_delete"))
				g.Expect(*policy.Properties.Policy.Rules[1].Name).To(ContainSubstring("testRule_cool"))
				return armstorage.ManagementPoliciesClientCreateOrUpdateResponse{}, nil
			},
		)
		rules := []message.LifecycleRuleParam{
			{RuleName: "testRule", Prefix: "test/", ExpirationDays: 30, ColdDays: 90},
		}
		bsCli.managementPoliciesClient = mockManagementPoliciesClient
		err := bsCli.AddLifecycle(context.TODO(), "fakeBucket", rules, []string{})
		g.Expect(err).To(Succeed())
	})

	t.Run("add new rules with existing policy and no conflicts", func(t *testing.T) {
		mockManagementPoliciesClient := azmock.NewMockmanagementCli(ctrl)
		existingRules := []*armstorage.ManagementPolicyRule{
			{
				Name:    toPtr("existingRule_delete"),
				Enabled: toPtr(true),
				Definition: &armstorage.ManagementPolicyDefinition{
					Filters: &armstorage.ManagementPolicyFilter{
						PrefixMatch: []*string{toPtr("fakeBucket/existing/")},
					},
				},
			},
		}
		mockManagementPoliciesClient.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(armstorage.ManagementPoliciesClientGetResponse{
			ManagementPolicy: armstorage.ManagementPolicy{
				Properties: &armstorage.ManagementPolicyProperties{
					Policy: &armstorage.ManagementPolicySchema{
						Rules: existingRules,
					},
				},
			},
		}, nil)
		mockManagementPoliciesClient.EXPECT().CreateOrUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
			func(_ context.Context, _ string, _ string, _ armstorage.ManagementPolicyName, policy armstorage.ManagementPolicy, _ *armstorage.ManagementPoliciesClientCreateOrUpdateOptions) (armstorage.ManagementPoliciesClientCreateOrUpdateResponse, error) {
				g.Expect(policy.Properties.Policy.Rules).To(HaveLen(3)) // existing + 2 new
				return armstorage.ManagementPoliciesClientCreateOrUpdateResponse{}, nil
			},
		)

		rules := []message.LifecycleRuleParam{
			{RuleName: "testRule", Prefix: "test/", ExpirationDays: 30, ColdDays: 90},
		}
		bsCli.managementPoliciesClient = mockManagementPoliciesClient
		err := bsCli.AddLifecycle(context.TODO(), "fakeBucket", rules, []string{})
		g.Expect(err).To(Succeed())
	})

	// Temporarily disable: this UT is experiencing false negative.
	// t.Run("replace existing rules with the same name", func(t *testing.T) {
	// 	mockManagementPoliciesClient := azmock.NewMockmanagementCli(ctrl)
	// 	existingRules := []*armstorage.ManagementPolicyRule{
	// 		{
	// 			Name:    toPtr("testRule_delete"),
	// 			Enabled: toPtr(true),
	// 			Definition: &armstorage.ManagementPolicyDefinition{
	// 				Filters: &armstorage.ManagementPolicyFilter{
	// 					PrefixMatch: []*string{toPtr("fakeBucket/old_test/")},
	// 				},
	// 				Actions: &armstorage.ManagementPolicyAction{
	// 					BaseBlob: &armstorage.ManagementPolicyBaseBlob{
	// 						Delete: &armstorage.DateAfterModification{
	// 							DaysAfterModificationGreaterThan: toPtr(float32(10)),
	// 						},
	// 					},
	// 				},
	// 			},
	// 		},
	// 	}
	// 	mockManagementPoliciesClient.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(armstorage.ManagementPoliciesClientGetResponse{
	// 		ManagementPolicy: armstorage.ManagementPolicy{
	// 			Properties: &armstorage.ManagementPolicyProperties{
	// 				Policy: &armstorage.ManagementPolicySchema{
	// 					Rules: existingRules,
	// 				},
	// 			},
	// 		},
	// 	}, nil)
	// 	mockManagementPoliciesClient.EXPECT().CreateOrUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
	// 		func(_ context.Context, _ string, _ string, _ armstorage.ManagementPolicyName, policy armstorage.ManagementPolicy, _ *armstorage.ManagementPoliciesClientCreateOrUpdateOptions) (armstorage.ManagementPoliciesClientCreateOrUpdateResponse, error) {
	// 			g.Expect(policy.Properties.Policy.Rules).To(HaveLen(2)) // new delete + new cool
	// 			g.Expect(*policy.Properties.Policy.Rules[0].Name).To(Equal("testRule_delete"))
	// 			g.Expect(*policy.Properties.Policy.Rules[0].Definition.Actions.BaseBlob.Delete.DaysAfterModificationGreaterThan).To(Equal(float32(30)))
	// 			g.Expect(*policy.Properties.Policy.Rules[1].Name).To(Equal("testRule_cool"))
	// 			return armstorage.ManagementPoliciesClientCreateOrUpdateResponse{}, nil
	// 		},
	// 	)

	// 	rules := []message.LifecycleRuleParam{
	// 		{RuleName: "testRule", Prefix: "test/", ExpirationDays: 30, ColdDays: 90},
	// 	}
	// 	bsCli.managementPoliciesClient = mockManagementPoliciesClient
	// 	err := bsCli.AddLifecycle(context.TODO(), "fakeBucket", rules, []string{})
	// 	g.Expect(err).To(Succeed())
	// })

	t.Run("filter out rules with unused prefixes", func(t *testing.T) {
		mockManagementPoliciesClient := azmock.NewMockmanagementCli(ctrl)
		existingRules := []*armstorage.ManagementPolicyRule{
			{
				Name:    toPtr("used_rule_delete"),
				Enabled: toPtr(true),
				Definition: &armstorage.ManagementPolicyDefinition{
					Filters: &armstorage.ManagementPolicyFilter{
						PrefixMatch: []*string{toPtr("fakeBucket/used/")},
					},
				},
			},
			{
				Name:    toPtr("unused_rule_delete"),
				Enabled: toPtr(true),
				Definition: &armstorage.ManagementPolicyDefinition{
					Filters: &armstorage.ManagementPolicyFilter{
						PrefixMatch: []*string{toPtr("fakeBucket/unused/")},
					},
				},
			},
		}
		mockManagementPoliciesClient.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(armstorage.ManagementPoliciesClientGetResponse{
			ManagementPolicy: armstorage.ManagementPolicy{
				Properties: &armstorage.ManagementPolicyProperties{
					Policy: &armstorage.ManagementPolicySchema{
						Rules: existingRules,
					},
				},
			},
		}, nil)
		mockManagementPoliciesClient.EXPECT().CreateOrUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
			func(_ context.Context, _ string, _ string, _ armstorage.ManagementPolicyName, policy armstorage.ManagementPolicy, _ *armstorage.ManagementPoliciesClientCreateOrUpdateOptions) (armstorage.ManagementPoliciesClientCreateOrUpdateResponse, error) {
				g.Expect(policy.Properties.Policy.Rules).To(HaveLen(2)) // one new + one existing (used)
				g.Expect(*policy.Properties.Policy.Rules[0].Name).To(ContainSubstring("newRule"))
				g.Expect(*policy.Properties.Policy.Rules[1].Name).To(Equal("used_rule_delete"))
				return armstorage.ManagementPoliciesClientCreateOrUpdateResponse{}, nil
			},
		)

		rules := []message.LifecycleRuleParam{
			{RuleName: "newRule", Prefix: "new/", ExpirationDays: 10},
		}
		unusedRulePrefixes := []string{"unused/"}
		bsCli.managementPoliciesClient = mockManagementPoliciesClient
		err := bsCli.AddLifecycle(context.TODO(), "fakeBucket", rules, unusedRulePrefixes)
		g.Expect(err).To(Succeed())
	})

	t.Run("Get management policy error (not 404)", func(t *testing.T) {
		mockManagementPoliciesClient := azmock.NewMockmanagementCli(ctrl)
		mockManagementPoliciesClient.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(armstorage.ManagementPoliciesClientGetResponse{}, errBoom)

		rules := []message.LifecycleRuleParam{
			{RuleName: "testRule", ExpirationDays: 30},
		}
		bsCli.managementPoliciesClient = mockManagementPoliciesClient
		err := bsCli.AddLifecycle(context.TODO(), "fakeBucket", rules, []string{})
		g.Expect(err).To(HaveOccurred())
		g.Expect(err.Error()).To(ContainSubstring("failed to get existing management policy"))
	})

	t.Run("CreateOrUpdate management policy error", func(t *testing.T) {
		mockManagementPoliciesClient := azmock.NewMockmanagementCli(ctrl)
		mockManagementPoliciesClient.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(armstorage.ManagementPoliciesClientGetResponse{}, &azcore.ResponseError{StatusCode: http.StatusNotFound, RawResponse: &http.Response{Status: "404 Not Found"}})
		mockManagementPoliciesClient.EXPECT().CreateOrUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(armstorage.ManagementPoliciesClientCreateOrUpdateResponse{}, errBoom)

		rules := []message.LifecycleRuleParam{
			{RuleName: "testRule", ExpirationDays: 30},
		}
		bsCli.managementPoliciesClient = mockManagementPoliciesClient
		err := bsCli.AddLifecycle(context.TODO(), "fakeBucket", rules, []string{})
		g.Expect(err).To(HaveOccurred())
		g.Expect(err.Error()).To(ContainSubstring("failed to CreateOrUpdate management policy"))
	})

	t.Run("existing policy with nil properties or policy", func(t *testing.T) {
		mockManagementPoliciesClient := azmock.NewMockmanagementCli(ctrl)
		mockManagementPoliciesClient.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(armstorage.ManagementPoliciesClientGetResponse{
			ManagementPolicy: armstorage.ManagementPolicy{
				Properties: &armstorage.ManagementPolicyProperties{
					Policy: nil, // Simulating nil policy
				},
			},
		}, nil)
		mockManagementPoliciesClient.EXPECT().CreateOrUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
			func(_ context.Context, _ string, _ string, _ armstorage.ManagementPolicyName, policy armstorage.ManagementPolicy, _ *armstorage.ManagementPoliciesClientCreateOrUpdateOptions) (armstorage.ManagementPoliciesClientCreateOrUpdateResponse, error) {
				g.Expect(policy.Properties.Policy.Rules).To(HaveLen(1))
				return armstorage.ManagementPoliciesClientCreateOrUpdateResponse{}, nil
			},
		)
		rules := []message.LifecycleRuleParam{
			{RuleName: "testRule", ExpirationDays: 30},
		}
		bsCli.managementPoliciesClient = mockManagementPoliciesClient
		err := bsCli.AddLifecycle(context.TODO(), "fakeBucket", rules, []string{})
		g.Expect(err).To(Succeed())
	})

	t.Run("existing policy with empty rules", func(t *testing.T) {
		mockManagementPoliciesClient := azmock.NewMockmanagementCli(ctrl)
		mockManagementPoliciesClient.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(armstorage.ManagementPoliciesClientGetResponse{
			ManagementPolicy: armstorage.ManagementPolicy{
				Properties: &armstorage.ManagementPolicyProperties{
					Policy: &armstorage.ManagementPolicySchema{
						Rules: []*armstorage.ManagementPolicyRule{},
					},
				},
			},
		}, nil)
		mockManagementPoliciesClient.EXPECT().CreateOrUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
			func(_ context.Context, _ string, _ string, _ armstorage.ManagementPolicyName, policy armstorage.ManagementPolicy, _ *armstorage.ManagementPoliciesClientCreateOrUpdateOptions) (armstorage.ManagementPoliciesClientCreateOrUpdateResponse, error) {
				g.Expect(policy.Properties.Policy.Rules).To(HaveLen(1))
				return armstorage.ManagementPoliciesClientCreateOrUpdateResponse{}, nil
			},
		)
		rules := []message.LifecycleRuleParam{
			{RuleName: "testRule", ExpirationDays: 30},
		}
		bsCli.managementPoliciesClient = mockManagementPoliciesClient
		err := bsCli.AddLifecycle(context.TODO(), "fakeBucket", rules, []string{})
		g.Expect(err).To(Succeed())
	})

	t.Run("existing policy with rule without name should be dropped", func(t *testing.T) {
		mockManagementPoliciesClient := azmock.NewMockmanagementCli(ctrl)
		existingRules := []*armstorage.ManagementPolicyRule{
			{ // This rule should be dropped
				Name:    nil,
				Enabled: toPtr(true),
				Definition: &armstorage.ManagementPolicyDefinition{
					Filters: &armstorage.ManagementPolicyFilter{
						PrefixMatch: []*string{toPtr("fakeBucket/old_test/")},
					},
				},
			},
			{ // This rule should be kept
				Name:    toPtr("keep_rule"),
				Enabled: toPtr(true),
				Definition: &armstorage.ManagementPolicyDefinition{
					Filters: &armstorage.ManagementPolicyFilter{
						PrefixMatch: []*string{toPtr("fakeBucket/keep/")},
					},
				},
			},
		}
		mockManagementPoliciesClient.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(armstorage.ManagementPoliciesClientGetResponse{
			ManagementPolicy: armstorage.ManagementPolicy{
				Properties: &armstorage.ManagementPolicyProperties{
					Policy: &armstorage.ManagementPolicySchema{
						Rules: existingRules,
					},
				},
			},
		}, nil)
		mockManagementPoliciesClient.EXPECT().CreateOrUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
			func(_ context.Context, _ string, _ string, _ armstorage.ManagementPolicyName, policy armstorage.ManagementPolicy, _ *armstorage.ManagementPoliciesClientCreateOrUpdateOptions) (armstorage.ManagementPoliciesClientCreateOrUpdateResponse, error) {
				g.Expect(policy.Properties.Policy.Rules).To(HaveLen(2)) // one new + "keep_rule"
				g.Expect(*policy.Properties.Policy.Rules[0].Name).To(ContainSubstring("newRule"))
				g.Expect(*policy.Properties.Policy.Rules[1].Name).To(Equal("keep_rule"))
				return armstorage.ManagementPoliciesClientCreateOrUpdateResponse{}, nil
			},
		)

		rules := []message.LifecycleRuleParam{
			{RuleName: "newRule", ExpirationDays: 10},
		}
		bsCli.managementPoliciesClient = mockManagementPoliciesClient
		err := bsCli.AddLifecycle(context.TODO(), "fakeBucket", rules, []string{})
		g.Expect(err).To(Succeed())
	})

	t.Run("existing policy with rule with matching new rule name should be dropped (replaced)", func(t *testing.T) {
		mockManagementPoliciesClient := azmock.NewMockmanagementCli(ctrl)
		existingRules := []*armstorage.ManagementPolicyRule{
			{ // This rule should be dropped as "newRule_delete" will replace it
				Name:    toPtr("newRule_delete"),
				Enabled: toPtr(true),
				Definition: &armstorage.ManagementPolicyDefinition{
					Filters: &armstorage.ManagementPolicyFilter{
						PrefixMatch: []*string{toPtr("fakeBucket/old_prefix/")},
					},
				},
			},
		}
		mockManagementPoliciesClient.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(armstorage.ManagementPoliciesClientGetResponse{
			ManagementPolicy: armstorage.ManagementPolicy{
				Properties: &armstorage.ManagementPolicyProperties{
					Policy: &armstorage.ManagementPolicySchema{
						Rules: existingRules,
					},
				},
			},
		}, nil)
		mockManagementPoliciesClient.EXPECT().CreateOrUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
			func(_ context.Context, _ string, _ string, _ armstorage.ManagementPolicyName, policy armstorage.ManagementPolicy, _ *armstorage.ManagementPoliciesClientCreateOrUpdateOptions) (armstorage.ManagementPoliciesClientCreateOrUpdateResponse, error) {
				g.Expect(policy.Properties.Policy.Rules).To(HaveLen(1))
				g.Expect(*policy.Properties.Policy.Rules[0].Name).To(Equal("newRule_delete"))
				g.Expect(*policy.Properties.Policy.Rules[0].Definition.Filters.PrefixMatch[0]).To(Equal("fakeBucket/new_prefix"))
				return armstorage.ManagementPoliciesClientCreateOrUpdateResponse{}, nil
			},
		)

		rules := []message.LifecycleRuleParam{
			{RuleName: "newRule", Prefix: "new_prefix", ExpirationDays: 10},
		}
		bsCli.managementPoliciesClient = mockManagementPoliciesClient
		err := bsCli.AddLifecycle(context.TODO(), "fakeBucket", rules, []string{})
		g.Expect(err).To(Succeed())
	})

	t.Run("existing policy with rule name matching unused rule prefix should be dropped", func(t *testing.T) {
		mockManagementPoliciesClient := azmock.NewMockmanagementCli(ctrl)
		existingRules := []*armstorage.ManagementPolicyRule{
			{
				Name:    toPtr("profile_deletion_oldapp_delete"),
				Enabled: toPtr(true),
				Definition: &armstorage.ManagementPolicyDefinition{
					Filters: &armstorage.ManagementPolicyFilter{
						PrefixMatch: []*string{toPtr("fakeBucket/0/t/p/oldapp/profiles/")},
					},
				},
			},
		}
		mockManagementPoliciesClient.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(armstorage.ManagementPoliciesClientGetResponse{
			ManagementPolicy: armstorage.ManagementPolicy{
				Properties: &armstorage.ManagementPolicyProperties{
					Policy: &armstorage.ManagementPolicySchema{
						Rules: existingRules,
					},
				},
			},
		}, nil)
		mockManagementPoliciesClient.EXPECT().CreateOrUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
			func(_ context.Context, _ string, _ string, _ armstorage.ManagementPolicyName, policy armstorage.ManagementPolicy, _ *armstorage.ManagementPoliciesClientCreateOrUpdateOptions) (armstorage.ManagementPoliciesClientCreateOrUpdateResponse, error) {
				g.Expect(policy.Properties.Policy.Rules).To(HaveLen(1))
				g.Expect(*policy.Properties.Policy.Rules[0].Name).To(Equal("new_rule_delete"))
				return armstorage.ManagementPoliciesClientCreateOrUpdateResponse{}, nil
			},
		)

		rules := []message.LifecycleRuleParam{
			{RuleName: "new_rule", Prefix: "new/", ExpirationDays: 10},
		}
		unusedRulePrefixes := []string{"profile_deletion_"}
		bsCli.managementPoliciesClient = mockManagementPoliciesClient
		err := bsCli.AddLifecycle(context.TODO(), "fakeBucket", rules, unusedRulePrefixes)
		g.Expect(err).To(Succeed())
	})

	t.Run("existing policy with rule whose prefix matches an unused prefix should be dropped", func(t *testing.T) {
		mockManagementPoliciesClient := azmock.NewMockmanagementCli(ctrl)
		existingRules := []*armstorage.ManagementPolicyRule{
			{
				Name:    toPtr("unused_prefix_rule"),
				Enabled: toPtr(true),
				Definition: &armstorage.ManagementPolicyDefinition{
					Filters: &armstorage.ManagementPolicyFilter{
						PrefixMatch: []*string{toPtr("fakeBucket/unused_data/")},
					},
				},
			},
		}
		mockManagementPoliciesClient.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(armstorage.ManagementPoliciesClientGetResponse{
			ManagementPolicy: armstorage.ManagementPolicy{
				Properties: &armstorage.ManagementPolicyProperties{
					Policy: &armstorage.ManagementPolicySchema{
						Rules: existingRules,
					},
				},
			},
		}, nil)
		mockManagementPoliciesClient.EXPECT().CreateOrUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
			func(_ context.Context, _ string, _ string, _ armstorage.ManagementPolicyName, policy armstorage.ManagementPolicy, _ *armstorage.ManagementPoliciesClientCreateOrUpdateOptions) (armstorage.ManagementPoliciesClientCreateOrUpdateResponse, error) {
				g.Expect(policy.Properties.Policy.Rules).To(HaveLen(1))
				g.Expect(*policy.Properties.Policy.Rules[0].Name).To(Equal("new_rule_delete"))
				return armstorage.ManagementPoliciesClientCreateOrUpdateResponse{}, nil
			},
		)

		rules := []message.LifecycleRuleParam{
			{RuleName: "new_rule", ExpirationDays: 10},
		}
		unusedRulePrefixes := []string{"unused_data/"}
		bsCli.managementPoliciesClient = mockManagementPoliciesClient
		err := bsCli.AddLifecycle(context.TODO(), "fakeBucket", rules, unusedRulePrefixes)
		g.Expect(err).To(Succeed())
	})

	t.Run("existing policy with rule whose prefix does not contain bucketName", func(t *testing.T) {
		mockManagementPoliciesClient := azmock.NewMockmanagementCli(ctrl)
		existingRules := []*armstorage.ManagementPolicyRule{
			{
				Name:    toPtr("other_bucket_rule"),
				Enabled: toPtr(true),
				Definition: &armstorage.ManagementPolicyDefinition{
					Filters: &armstorage.ManagementPolicyFilter{
						PrefixMatch: []*string{toPtr("otherBucket/data/")},
					},
				},
			},
		}
		mockManagementPoliciesClient.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(armstorage.ManagementPoliciesClientGetResponse{
			ManagementPolicy: armstorage.ManagementPolicy{
				Properties: &armstorage.ManagementPolicyProperties{
					Policy: &armstorage.ManagementPolicySchema{
						Rules: existingRules,
					},
				},
			},
		}, nil)
		mockManagementPoliciesClient.EXPECT().CreateOrUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
			func(_ context.Context, _ string, _ string, _ armstorage.ManagementPolicyName, policy armstorage.ManagementPolicy, _ *armstorage.ManagementPoliciesClientCreateOrUpdateOptions) (armstorage.ManagementPoliciesClientCreateOrUpdateResponse, error) {
				g.Expect(policy.Properties.Policy.Rules).To(HaveLen(2))
				g.Expect(*policy.Properties.Policy.Rules[0].Name).To(ContainSubstring("new_rule"))
				g.Expect(*policy.Properties.Policy.Rules[1].Name).To(Equal("other_bucket_rule"))
				return armstorage.ManagementPoliciesClientCreateOrUpdateResponse{}, nil
			},
		)

		rules := []message.LifecycleRuleParam{
			{RuleName: "new_rule", ExpirationDays: 10},
		}
		unusedRulePrefixes := []string{"unused_data/"}
		bsCli.managementPoliciesClient = mockManagementPoliciesClient
		err := bsCli.AddLifecycle(context.TODO(), "fakeBucket", rules, unusedRulePrefixes)
		g.Expect(err).To(Succeed())
	})

	t.Run("existing policy with rule with nil Definition or Filters or PrefixMatch", func(t *testing.T) {
		mockManagementPoliciesClient := azmock.NewMockmanagementCli(ctrl)
		existingRules := []*armstorage.ManagementPolicyRule{
			{
				Name:       toPtr("rule_nil_definition"),
				Enabled:    toPtr(true),
				Definition: nil,
			},
			{
				Name:    toPtr("rule_nil_filters"),
				Enabled: toPtr(true),
				Definition: &armstorage.ManagementPolicyDefinition{
					Filters: nil,
				},
			},
			{
				Name:    toPtr("rule_nil_prefix_match"),
				Enabled: toPtr(true),
				Definition: &armstorage.ManagementPolicyDefinition{
					Filters: &armstorage.ManagementPolicyFilter{
						PrefixMatch: nil,
					},
				},
			},
			{
				Name:    toPtr("rule_empty_prefix_match"),
				Enabled: toPtr(true),
				Definition: &armstorage.ManagementPolicyDefinition{
					Filters: &armstorage.ManagementPolicyFilter{
						PrefixMatch: []*string{},
					},
				},
			},
		}
		mockManagementPoliciesClient.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(armstorage.ManagementPoliciesClientGetResponse{
			ManagementPolicy: armstorage.ManagementPolicy{
				Properties: &armstorage.ManagementPolicyProperties{
					Policy: &armstorage.ManagementPolicySchema{
						Rules: existingRules,
					},
				},
			},
		}, nil)
		mockManagementPoliciesClient.EXPECT().CreateOrUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
			func(_ context.Context, _ string, _ string, _ armstorage.ManagementPolicyName, policy armstorage.ManagementPolicy, _ *armstorage.ManagementPoliciesClientCreateOrUpdateOptions) (armstorage.ManagementPoliciesClientCreateOrUpdateResponse, error) {
				g.Expect(policy.Properties.Policy.Rules).To(HaveLen(5)) // 1 new + 4 existing
				g.Expect(*policy.Properties.Policy.Rules[0].Name).To(ContainSubstring("new_rule"))
				g.Expect(*policy.Properties.Policy.Rules[1].Name).To(Equal("rule_nil_definition"))
				g.Expect(*policy.Properties.Policy.Rules[2].Name).To(Equal("rule_nil_filters"))
				g.Expect(*policy.Properties.Policy.Rules[3].Name).To(Equal("rule_nil_prefix_match"))
				g.Expect(*policy.Properties.Policy.Rules[4].Name).To(Equal("rule_empty_prefix_match"))
				return armstorage.ManagementPoliciesClientCreateOrUpdateResponse{}, nil
			},
		)

		rules := []message.LifecycleRuleParam{
			{RuleName: "new_rule", ExpirationDays: 10},
		}
		bsCli.managementPoliciesClient = mockManagementPoliciesClient
		err := bsCli.AddLifecycle(context.TODO(), "fakeBucket", rules, []string{})
		g.Expect(err).To(Succeed())
	})

	t.Run("truncate to azure max lifecycle rules keeping new rules first", func(t *testing.T) {
		mockManagementPoliciesClient := azmock.NewMockmanagementCli(ctrl)
		existingRules := make([]*armstorage.ManagementPolicyRule, 0, azureMaxLifecycleRules)
		for i := 0; i < azureMaxLifecycleRules; i++ {
			name := fmt.Sprintf("other_bucket_rule_%d", i)
			existingRules = append(existingRules, &armstorage.ManagementPolicyRule{
				Name:    toPtr(name),
				Enabled: toPtr(true),
				Definition: &armstorage.ManagementPolicyDefinition{
					Filters: &armstorage.ManagementPolicyFilter{
						PrefixMatch: []*string{toPtr("otherBucket/data/")},
					},
				},
			})
		}
		mockManagementPoliciesClient.EXPECT().Get(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(armstorage.ManagementPoliciesClientGetResponse{
			ManagementPolicy: armstorage.ManagementPolicy{
				Properties: &armstorage.ManagementPolicyProperties{
					Policy: &armstorage.ManagementPolicySchema{
						Rules: existingRules,
					},
				},
			},
		}, nil)
		mockManagementPoliciesClient.EXPECT().CreateOrUpdate(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
			func(_ context.Context, _ string, _ string, _ armstorage.ManagementPolicyName, policy armstorage.ManagementPolicy, _ *armstorage.ManagementPoliciesClientCreateOrUpdateOptions) (armstorage.ManagementPoliciesClientCreateOrUpdateResponse, error) {
				g.Expect(policy.Properties.Policy.Rules).To(HaveLen(azureMaxLifecycleRules))
				g.Expect(*policy.Properties.Policy.Rules[0].Name).To(Equal("new_rule_delete"))
				return armstorage.ManagementPoliciesClientCreateOrUpdateResponse{}, nil
			},
		)

		rules := []message.LifecycleRuleParam{
			{RuleName: "new_rule", ExpirationDays: 10},
		}
		bsCli.managementPoliciesClient = mockManagementPoliciesClient
		err := bsCli.AddLifecycle(context.TODO(), "fakeBucket", rules, []string{})
		g.Expect(err).To(Succeed())
	})
}

func TestRemoveLifecycle(t *testing.T) {
	g := NewWithT(t)
	bsCli := &BlobStorageClient{}
	g.Expect(func() { bsCli.RemoveLifecycle(context.TODO(), "fake", "fake") }).To(PanicWith("implement me"))
}

// there is no good way to fake ListObjects for now
// refer to https://github.com/Azure/azure-sdk-for-go/issues/16613