package azure

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"strings"

	"github.com/tidbcloud/observability/pkg/common/message"

	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/storage/armstorage"
	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob"
	"github.com/tidbcloud/observability/pkg/resourceworker/providers"
	"github.com/tidbcloud/observability/pkg/resourceworker/providers/lifecyclecompare"
	"github.com/tidbcloud/observability/pkg/resourceworker/providers/types"

	"bytes"
	"io"
)

// blobServiceURLTmpl - https://o11ydev.blob.core.windows.net/
const blobServiceURLTmpl = "https://%s.blob.core.windows.net/"
const azureLifecycleBlobTypeBlock = "blockBlob"
const azureMaxLifecycleRules = 100

type BlobStorageClient struct {
	cli                      containerCli
	managementPoliciesClient managementCli
	dataCli                  blobDataClient
	storageAccount           string
	resourceGroup            string
}

var _ providers.S3Client = &BlobStorageClient{}

var newClientFactory = armstorage.NewClientFactory
var newBlobClient = azblob.NewClient

func NewBlobStorageClient(subscription, resourceGroup, storageAccount string, cred azcore.TokenCredential) (*BlobStorageClient, error) {
	factory, err := newClientFactory(subscription, cred, nil)
	if err != nil {
		return nil, err
	}
	cli := factory.NewBlobContainersClient()
	managementPoliciesClient := factory.NewManagementPoliciesClient()

	dataCli, err := newBlobClient(fmt.Sprintf(blobServiceURLTmpl, storageAccount), cred, nil)
	if err != nil {
		return nil, err
	}
	return &BlobStorageClient{
		cli:                      cli,
		managementPoliciesClient: managementPoliciesClient,
		dataCli:                  dataCli,
		storageAccount:           storageAccount,
		resourceGroup:            resourceGroup,
	}, nil
}

func (b *BlobStorageClient) CreateBucket(ctx context.Context, name string) error {
	_, err := b.cli.Create(ctx, b.resourceGroup, b.storageAccount, name, armstorage.BlobContainer{}, nil)
	return err
}
func (b *BlobStorageClient) BucketExists(ctx context.Context, name string) (bool, error) {
	_, err := b.cli.Get(ctx, b.resourceGroup, b.storageAccount, name, nil)
	if err != nil {
		return false, err
	}

	return true, nil
}

func (b *BlobStorageClient) Get(ctx context.Context, bucketName string, objectName string) (*io.ReadCloser, error) {
	resp, err := b.dataCli.DownloadStream(ctx, bucketName, objectName, nil)
	if err != nil {
		var respErr *azcore.ResponseError
		if errors.As(err, &respErr) && respErr.StatusCode == http.StatusNotFound {
			return nil, types.ObjectNotFoundError
		}
		return nil, err
	}
	return &resp.Body, nil
}

// ListObjects lists objects in a bucket
func (b *BlobStorageClient) ListObjects(ctx context.Context, bucketName string, prefix string, options types.ListOptions) ([]string, error) {
	pager := b.dataCli.NewListBlobsFlatPager(bucketName, &azblob.ListBlobsFlatOptions{
		Prefix:     &prefix,
		MaxResults: options.MaxKeys,
	})
	objects := make([]string, 0)
	for pager.More() {
		page, err := pager.NextPage(ctx)
		if err != nil {
			return nil, err
		}
		for _, blob := range page.Segment.BlobItems {
			objects = append(objects, *blob.Name)
		}
	}
	return objects, nil
}

func (b *BlobStorageClient) Put(ctx context.Context, bucketName string, objectName string, reader *bytes.Reader) error {
	_, err := b.dataCli.UploadStream(ctx, bucketName, objectName, reader, nil)
	return err
}

// buildAzureLifecycleRules converts a message.LifecycleRuleParam into one or two armstorage.LifecycleRule objects.
// Each rule will be scoped to the specific bucketName and its relative prefix.
// These rules are for the Management Plane API.
func buildAzureLifecycleRules(bucketName string, param message.LifecycleRuleParam) []*armstorage.ManagementPolicyRule {
	var rules []*armstorage.ManagementPolicyRule

	fullPrefix := bucketName
	if param.Prefix != "" {
		if !strings.HasSuffix(fullPrefix, "/") {
			fullPrefix += "/"
		}
		fullPrefix += strings.TrimPrefix(param.Prefix, "/")
	}

	if param.ExpirationDays > 0 {
		ruleName := fmt.Sprintf("%s_delete", param.RuleName)
		rules = append(rules, &armstorage.ManagementPolicyRule{
			Type:    toPtr(armstorage.RuleTypeLifecycle),
			Name:    toPtr(ruleName),
			Enabled: toPtr(true),
			Definition: &armstorage.ManagementPolicyDefinition{
				Filters: &armstorage.ManagementPolicyFilter{
					BlobTypes:   []*string{toPtr(azureLifecycleBlobTypeBlock)},
					PrefixMatch: []*string{toPtr(fullPrefix)},
				},
				Actions: &armstorage.ManagementPolicyAction{
					BaseBlob: &armstorage.ManagementPolicyBaseBlob{
						Delete: &armstorage.DateAfterModification{
							DaysAfterModificationGreaterThan: toPtr(float32(param.ExpirationDays)),
						},
					},
				},
			},
		})
	}

	if param.ColdDays > 0 {
		ruleName := fmt.Sprintf("%s_cool", param.RuleName)
		rules = append(rules, &armstorage.ManagementPolicyRule{
			Type:    toPtr(armstorage.RuleTypeLifecycle),
			Name:    toPtr(ruleName),
			Enabled: toPtr(true),
			Definition: &armstorage.ManagementPolicyDefinition{
				Filters: &armstorage.ManagementPolicyFilter{
					BlobTypes:   []*string{toPtr(azureLifecycleBlobTypeBlock)},
					PrefixMatch: []*string{toPtr(fullPrefix)},
				},
				Actions: &armstorage.ManagementPolicyAction{
					BaseBlob: &armstorage.ManagementPolicyBaseBlob{
						TierToCool: &armstorage.DateAfterModification{
							DaysAfterModificationGreaterThan: toPtr(float32(param.ColdDays)),
						},
					},
				},
			},
		})
	}
	return rules
}

func (b *BlobStorageClient) getExistingAzureRules(ctx context.Context) ([]*armstorage.ManagementPolicyRule, error) {
	getPolicyRes, err := b.managementPoliciesClient.Get(ctx, b.resourceGroup, b.storageAccount, armstorage.ManagementPolicyNameDefault, nil)

	if err != nil {
		var httpErr *azcore.ResponseError
		if errors.As(err, &httpErr) && httpErr.StatusCode == http.StatusNotFound {
			return []*armstorage.ManagementPolicyRule{}, nil
		}
		return nil, fmt.Errorf("failed to get existing management policy for storage account %s: %w", b.storageAccount, err)
	}

	if getPolicyRes.ManagementPolicy.Properties == nil || getPolicyRes.ManagementPolicy.Properties.Policy == nil {
		return []*armstorage.ManagementPolicyRule{}, nil
	}

	if getPolicyRes.ManagementPolicy.Properties.Policy.Rules == nil {
		return []*armstorage.ManagementPolicyRule{}, nil
	}

	return getPolicyRes.ManagementPolicy.Properties.Policy.Rules, nil
}

func collectNewAzureRules(bucketName string, rules []message.LifecycleRuleParam) (map[string]*armstorage.ManagementPolicyRule, []string) {
	newAzureRulesMap := make(map[string]*armstorage.ManagementPolicyRule)
	newAzureRuleOrder := make([]string, 0, len(rules)*2)
	for _, param := range rules {
		azureRules := buildAzureLifecycleRules(bucketName, param)
		for _, azureRule := range azureRules {
			if azureRule.Name == nil {
				continue
			}
			ruleName := *azureRule.Name
			if _, exists := newAzureRulesMap[ruleName]; !exists {
				newAzureRuleOrder = append(newAzureRuleOrder, ruleName)
			}
			newAzureRulesMap[ruleName] = azureRule
		}
	}

	return newAzureRulesMap, newAzureRuleOrder
}

func shouldDropExistingAzureRule(oldRule *armstorage.ManagementPolicyRule, bucketName string, newAzureRulesMap map[string]*armstorage.ManagementPolicyRule, unusedRulePrefixes []string) bool {
	if oldRule.Name == nil {
		return true
	}
	if newAzureRulesMap[*oldRule.Name] != nil {
		return true
	}

	// Match AWS S3: unusedRulePrefixes are rule ID/name prefixes (e.g. profile_deletion_).
	for _, unusedPrefix := range unusedRulePrefixes {
		if strings.HasPrefix(*oldRule.Name, unusedPrefix) {
			return true
		}
	}

	if oldRule.Definition == nil || oldRule.Definition.Filters == nil || len(oldRule.Definition.Filters.PrefixMatch) == 0 {
		return false
	}

	oldRuleFullPrefix := *oldRule.Definition.Filters.PrefixMatch[0]
	if !strings.HasPrefix(oldRuleFullPrefix, bucketName) {
		return false
	}

	relativePrefix := strings.TrimPrefix(oldRuleFullPrefix, bucketName)
	if strings.HasPrefix(relativePrefix, "/") {
		relativePrefix = strings.TrimPrefix(relativePrefix, "/")
	}

	for _, unusedPrefix := range unusedRulePrefixes {
		if strings.HasPrefix(relativePrefix, unusedPrefix) {
			return true
		}
	}

	return false
}

func (b *BlobStorageClient) AddLifecycle(ctx context.Context, bucketName string, rules []message.LifecycleRuleParam, unusedRulePrefixes []string) error {
	existingAzureRules, err := b.getExistingAzureRules(ctx)
	if err != nil {
		return err
	}

	newAzureRulesMap, newAzureRuleOrder := collectNewAzureRules(bucketName, rules)
	rulesToApply := make([]*armstorage.ManagementPolicyRule, 0, len(newAzureRulesMap)+len(existingAzureRules))
	for _, ruleName := range newAzureRuleOrder {
		rulesToApply = append(rulesToApply, newAzureRulesMap[ruleName])
	}

	for _, oldRule := range existingAzureRules {
		if shouldDropExistingAzureRule(oldRule, bucketName, newAzureRulesMap, unusedRulePrefixes) {
			continue
		}

		rulesToApply = append(rulesToApply, oldRule)
	}

	if len(rulesToApply) == 0 {
		return fmt.Errorf("no lifecycle rules to apply for storage account %s", b.storageAccount)
	}

	existingKept := len(rulesToApply) - len(newAzureRulesMap)
	if len(rulesToApply) > azureMaxLifecycleRules {
		fmt.Printf(
			"azure lifecycle: truncating rules for storage account %s from %d to %d (new=%d, existing_kept=%d)\n",
			b.storageAccount, len(rulesToApply), azureMaxLifecycleRules, len(newAzureRulesMap), existingKept,
		)
		// New rules are appended first; drop oldest retained rules from the tail.
		rulesToApply = rulesToApply[:azureMaxLifecycleRules]
	} else {
		fmt.Printf(
			"azure lifecycle: applying %d rules to storage account %s (new=%d, existing_kept=%d)\n",
			len(rulesToApply), b.storageAccount, len(newAzureRulesMap), existingKept,
		)
	}

	if lifecyclecompare.AzureLifecycleRulesEqual(rulesToApply, existingAzureRules) {
		fmt.Printf("lifecycle rules unchanged for storage account %s, skip update\n", b.storageAccount)
		return nil
	}

	managementPolicy := armstorage.ManagementPolicy{
		Properties: &armstorage.ManagementPolicyProperties{
			Policy: &armstorage.ManagementPolicySchema{
				Rules: rulesToApply,
			},
		},
	}
	_, err = b.managementPoliciesClient.CreateOrUpdate(ctx, b.resourceGroup, b.storageAccount, armstorage.ManagementPolicyNameDefault, managementPolicy, nil)

	if err != nil {
		return fmt.Errorf("failed to CreateOrUpdate management policy: %w", err)
	}

	return nil
}

func (b *BlobStorageClient) RemoveLifecycle(ctx context.Context, bucketName string, name string) error {
	panic("implement me")
}

func (b *BlobStorageClient) Delete(ctx context.Context, bucketName string, objectNames ...string) error {
	for _, objectName := range objectNames {
		_, err := b.dataCli.DeleteBlob(ctx, bucketName, objectName, nil)
		if err != nil {
			return err
		}
	}
	return nil
}

func toPtr[T any](v T) *T {
	return &v
}