// Package azblob add by blacknc <blacknc@163.com> at 2024/09/18 14:03:46
package azblob

import (
	"context"
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"io"
	"net/url"
	"path"
	"regexp"
	"strings"

	"github.com/Azure/azure-sdk-for-go/sdk/azcore"
	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob"
	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/blob"
	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/bloberror"
	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob/container"
	"github.com/pingcap/errors"
	"github.com/pingcap/log"
	"github.com/tidbcloud/dataflow-service/internal/constants"
	"go.uber.org/zap"
	"k8s.io/utils/ptr"
)

var (
	azblobHostRegexp = regexp.MustCompile(`^([^\.]+)\.blob\.core\.windows\.net$`)
	azblobPathRegexp = regexp.MustCompile(`^/([^\/]+)/(.*)?$`)
)

// CheckAzblobAccess check if azblobURL can be accessed
// azblobURL: https://<account-name>.blob.core.windows.net/<container-name>/<blob-prefix>
func CheckAzblobAccess(ctx context.Context, azblobURL, sasToken string) error {
	abs, err := NewAzureBlobStorage(ctx, azblobURL, sasToken, nil)
	if err != nil {
		return err
	}

	expectedErr := errors.New("Stop Iter")
	err = abs.WalkDir(ctx, &WalkOption{ListCount: 1}, func(string, int64) error {
		// return an error when meet the first regular file to break the walk loop
		return expectedErr
	})
	if !errors.ErrorEqual(err, expectedErr) {
		if err == nil {
			return errors.Errorf("there is no file in this container")
		}
		return errors.AddStack(err)
	}
	return nil
}

// GetFileContent get azblobURL's content with range [startOffset, endOffset]
// azblobURL: https://<account-name>.blob.core.windows.net/<container-name>/<blob-name>
func GetFileContent(ctx context.Context, azblobURL, sasToken string, startOffset, endOffset int64) (string, error) {

	abs, err := NewAzureBlobStorage(ctx, azblobURL, sasToken, nil)
	if err != nil {
		return "", err
	}

	data, err := abs.ReadFile(ctx, "", startOffset, endOffset)
	if err != nil {
		return "", err
	}
	return string(data), nil
}

type AzureBlobStorage interface {
	// FileExists return true if file exists
	FileExists(ctx context.Context, name string) (bool, error)
	WalkDir(ctx context.Context, opt *WalkOption, fn func(path string, size int64) error) error
	Container(ctx context.Context, container string) (*container.Client, error)
	GetAccountName() string
	ReadFile(ctx context.Context, name string, startOffset, endOffset int64) ([]byte, error)
}

var ErrStorageInvalidConfig = errors.Normalize("invalid external storage config", errors.RFCCodeText("Import:ExternalStorage:ErrStorageInvalidConfig"))

// NewAzureBlobStorage new azure blob storage handler
// azblobURL: https://<account-name>.blob.core.windows.net/<container-name>/<blob-prefix>
func NewAzureBlobStorage(ctx context.Context, sourceURL, sasToken string, options *AzblobBackendOptions) (AzureBlobStorage, error) {

	if sourceURL == "" || sasToken == "" {
		return nil, errors.Annotate(ErrStorageInvalidConfig, "empty source url or token is not allowed")
	}

	azurl, err := ParseAzblobURL(sourceURL)
	if err != nil {
		return nil, err
	}

	if options == nil {
		options = &AzblobBackendOptions{}
	}
	options.Bucket = azurl.Container()
	options.Prefix = strings.Trim(azurl.BlobPath(), "/")
	options.AccountName = azurl.AccountName()
	options.SASToken = sasToken
	return newAzureBlobStorageWithOptions(ctx, options)
}

type AzblobURL struct {
	accountName   string
	containerName string
	blobPath      string
	sasToken      string
}

func (u *AzblobURL) ToInnerURL(needQuery bool) string {
	innerUrl, _ := url.Parse(fmt.Sprintf("azure://%s/%s", u.containerName, u.blobPath))
	if needQuery {
		query := url.Values{}
		if u.accountName != "" {
			query.Set("account-name", u.accountName)
		}
		if u.sasToken != "" {
			query.Set("sas-token", u.sasToken)
		}
		innerUrl.RawQuery = query.Encode()
	}
	return innerUrl.String()
}

func (u *AzblobURL) AccountName() string {
	return u.accountName
}

func (u *AzblobURL) Container() string {
	return u.containerName
}

func (u *AzblobURL) BlobPath() string {
	return u.blobPath
}

func (u *AzblobURL) SASToken() string {
	return u.sasToken
}
func (u *AzblobURL) SetSASToken(token string) {
	u.sasToken = token
}

func IsAzureScheme(scheme string) bool {
	switch scheme {
	case constants.AzblobScheme:
		return true
	default:
		return false
	}
}

func ParseAzblobURL(sourceURL string) (*AzblobURL, error) {

	errmsg := fmt.Sprintf("source url '%s' is not valid. Please provide correct url, such as the template format '%s'",
		sourceURL, "https://<account-name>.blob.core.windows.net/<container-name>/<blob-prefix>/")
	comerr := errors.Annotate(ErrStorageInvalidConfig, errmsg)

	u, err := url.Parse(sourceURL)
	if err != nil {
		return nil, comerr
	}
	if !IsAzureScheme(u.Scheme) {
		return nil, comerr
	}

	azurl := new(AzblobURL)
	if matchs := azblobHostRegexp.FindStringSubmatch(u.Host); len(matchs) != 2 {
		return nil, comerr
	} else {
		azurl.accountName = matchs[1]
	}

	if matchs := azblobPathRegexp.FindStringSubmatch(u.Path); len(matchs) != 3 {
		return nil, comerr
	} else {
		azurl.containerName = matchs[1]
		azurl.blobPath = strings.TrimLeft(matchs[2], "/")
	}

	azurl.sasToken = u.Query().Get("sas-token")
	return azurl, nil
}

// ConvertAzblobURL convert raw azblob url to lightning azblob url format
// from: https://<account-name>.blob.core.windows.net/<contoainer-name>/<blob-name>
// to:   azure://<container-name>/<blob-name>?account-name=<account-name>
func ConvertAzblobURL(sourceURL string, needQuery bool) (string, error) {
	azurl, err := ParseAzblobURL(sourceURL)
	if err != nil {
		return "", err
	}
	return azurl.ToInnerURL(needQuery), nil
}

// AzblobBackendOptions is the options for Azure Blob storage.
type AzblobBackendOptions struct {
	Bucket          string `json:"bucket"`
	Prefix          string `json:"prefix"`
	Endpoint        string `json:"endpoint"`
	AccountName     string `json:"account-name"`
	AccountKey      string `json:"account-key"`
	AccessTier      string `json:"access-tier"`
	SASToken        string `json:"sas-token"`
	EncryptionScope string `json:"encryption-scope"`
	EncryptionKey   string `json:"encryption-key"`
}

const azblobRetryTimes int32 = 5

func getDefaultClientOptions() *azblob.ClientOptions {
	return &azblob.ClientOptions{
		ClientOptions: azcore.ClientOptions{
			Retry: policy.RetryOptions{
				MaxRetries: azblobRetryTimes,
			},
		},
	}
}

// AzureBlobStorage is a storage engine that stores data in Azure Blob Storage.
type azureBlobStorageImpl struct {
	options         *AzblobBackendOptions
	serviceClient   *azblob.Client
	containerClient *container.Client
	accessTier      blob.AccessTier
	cpkScope        *blob.CPKScopeInfo
	cpkInfo         *blob.CPKInfo
}

func newAzureBlobStorageWithOptions(ctx context.Context, options *AzblobBackendOptions) (*azureBlobStorageImpl, error) {

	if len(options.Bucket) == 0 {
		return nil, errors.New("bucket(container) cannot be empty to access azure blob storage")
	}

	if len(options.AccountName) == 0 || len(options.SASToken) == 0 {
		return nil, errors.New("azure blob account name and access sign is required!")
	}
	serviceURL := options.Endpoint
	if len(serviceURL) == 0 {
		if strings.HasPrefix(options.SASToken, "?") {
			serviceURL = fmt.Sprintf("https://%s.blob.core.windows.net/%s", options.AccountName, options.SASToken)
		} else {
			serviceURL = fmt.Sprintf("https://%s.blob.core.windows.net/?%s", options.AccountName, options.SASToken)
		}
	}

	clientOptions := getDefaultClientOptions()
	serviceClient, err := azblob.NewClientWithNoCredential(serviceURL, clientOptions)
	if err != nil {
		return nil, errors.Annotate(err, "Failed to create azure service client")
	}

	containerClient := serviceClient.ServiceClient().NewContainerClient(options.Bucket)
	if _, err = containerClient.GetProperties(ctx, &container.GetPropertiesOptions{}); err != nil {
		return nil, errors.Trace(err)
	}

	var cpkScope *blob.CPKScopeInfo = nil
	if len(options.EncryptionScope) > 0 {
		cpkScope = &blob.CPKScopeInfo{
			EncryptionScope: &options.EncryptionScope,
		}
	}

	var cpkInfo *blob.CPKInfo = nil
	if len(options.EncryptionKey) > 0 {
		defaultAlgorithm := blob.EncryptionAlgorithmTypeAES256
		keySlice := []byte(options.EncryptionKey)
		keySha256 := sha256.Sum256(keySlice)
		cpkInfo = &blob.CPKInfo{
			EncryptionAlgorithm: &defaultAlgorithm,
			EncryptionKey:       ptr.To(base64.StdEncoding.EncodeToString(keySlice)),
			EncryptionKeySHA256: ptr.To(base64.StdEncoding.EncodeToString(keySha256[:])),
		}
	}

	// parse storage access-tier
	var accessTier blob.AccessTier
	switch options.AccessTier {
	case "Archive", "archive":
		accessTier = blob.AccessTierArchive
	case "Cool", "cool":
		accessTier = blob.AccessTierCool
	case "Hot", "hot":
		accessTier = blob.AccessTierHot
	default:
		accessTier = blob.AccessTier(options.AccessTier)
	}

	log.Debug("select accessTier", zap.String("accessTier", string(accessTier)))

	return &azureBlobStorageImpl{
		options,
		serviceClient,
		containerClient,
		accessTier,
		cpkScope,
		cpkInfo,
	}, nil
}

// FileExists checks if file with name exists in Azure Blob Storage.
// name: blob base name
func (s *azureBlobStorageImpl) FileExists(ctx context.Context, name string) (bool, error) {
	client := s.containerClient.NewBlockBlobClient(s.withPrefix(name))
	_, err := client.GetProperties(ctx, nil)
	if err != nil {
		if bloberror.HasCode(err, bloberror.BlobNotFound) {
			return false, nil
		}
		return false, errors.Trace(err)
	}
	return true, nil
}

// WalkOption is the option of storage.WalkDir.
type WalkOption struct {
	// walk on SubDir of base directory, i.e. if the base dir is '/path/to/base'
	// then we're walking '/path/to/base/<SubDir>'
	SubDir string
	// whether subdirectory under the walk dir is skipped, only works for LOCAL storage now.
	// default is false, i.e. we walk recursively.
	SkipSubDir bool
	// ObjPrefix used fo prefix search in storage. Note that only part of storage
	// support it.
	// It can save lots of time when we want find specify prefix objects in storage.
	// For example. we have 10000 <Hash>.sst files and 10 backupmeta.(\d+) files.
	// we can use ObjPrefix = "backupmeta" to retrieve all meta files quickly.
	ObjPrefix string
	// ListCount is the number of entries per page.
	//
	// In cloud storages such as S3 and GCS, the files listed and sent in pages.
	// Typically a page contains 1000 files, and if a folder has 3000 descendant
	// files, one would need 3 requests to retrieve all of them. This parameter
	// controls this size. Note that both S3 and GCS limits the maximum to 1000.
	//
	// Typically you want to leave this field unassigned (zero) to use the
	// default value (1000) to minimize the number of requests, unless you want
	// to reduce the possibility of timeout on an extremely slow connection, or
	// perform testing.
	ListCount int64
}

// WalkDir implements the StorageReader interface.
// check has file
// // return expectedErr means at least meet one file
// expectedErr := errors.New("Stop Iter")
//
//	walkErr := s.WalkDir(ctx, &storage.WalkOption{ListCount: 1}, func(string, int64) error {
//		// return an error when meet the first regular file to break the walk loop
//		return expectedErr
//	})
//
//	if !errors.ErrorEqual(walkErr, expectedErr) {
//		if walkErr == nil {
//			return common.ErrEmptySourceDir.GenWithStackByArgs(taskCfg.Mydumper.SourceDir)
//		}
//		return common.NormalizeOrWrapErr(common.ErrStorageUnknown, walkErr)
//	}
func (s *azureBlobStorageImpl) WalkDir(ctx context.Context, opt *WalkOption, fn func(path string, size int64) error) error {
	if opt == nil {
		opt = &WalkOption{}
	}
	prefix := path.Join(s.options.Prefix, opt.SubDir)
	if len(prefix) > 0 && !strings.HasSuffix(prefix, "/") {
		prefix += "/"
	}
	if len(opt.ObjPrefix) != 0 {
		prefix += opt.ObjPrefix
	}

	pager := s.containerClient.NewListBlobsFlatPager(&container.ListBlobsFlatOptions{
		Prefix: &prefix,
	})
	for pager.More() {
		page, err := pager.NextPage(ctx)
		if err != nil {
			return errors.Annotatef(err, "Failed to list azure blobs, bucket(container)='%s'", s.options.Bucket)
		}

		for _, blob := range page.Segment.BlobItems {
			// when walk on specify directory, the result include storage.Prefix,
			// which can not be reuse in other API(Open/Read) directly.
			// so we use TrimPrefix to filter Prefix for next Open/Read.
			path := strings.TrimPrefix((*blob.Name), s.options.Prefix)
			// trim the prefix '/' to ensure that the path returned is consistent with the local storage
			path = strings.TrimPrefix(path, "/")
			if err := fn(path, *blob.Properties.ContentLength); err != nil {
				return errors.Trace(err)
			}
		}
	}

	return nil
}

func (s *azureBlobStorageImpl) Container(ctx context.Context, containerName string) (*container.Client, error) {
	containerClient := s.serviceClient.ServiceClient().NewContainerClient(containerName)
	if _, err := containerClient.GetProperties(ctx, &container.GetPropertiesOptions{}); err != nil {
		return nil, errors.Trace(err)
	}
	return containerClient, nil
}

func (s *azureBlobStorageImpl) GetAccountName() string {
	return s.options.AccountName
}

func (s *azureBlobStorageImpl) withPrefix(name string) string {
	return path.Join(s.options.Prefix, name)
}

// ReadFile reads a file from Azure Blob Storage.
// name: blob base name
func (s *azureBlobStorageImpl) ReadFile(ctx context.Context, name string, startOffset, endOffset int64) ([]byte, error) {
	client := s.containerClient.NewBlockBlobClient(s.withPrefix(name))
	resp, err := client.DownloadStream(ctx, &blob.DownloadStreamOptions{
		CPKInfo: s.cpkInfo,
		Range: azblob.HTTPRange{
			Offset: startOffset,
			Count:  endOffset - startOffset + 1,
		},
	})
	if err != nil {
		return nil, errors.Annotatef(err, "Failed to download azure blob file, file info: bucket(container)='%s', key='%s'", s.options.Bucket, s.withPrefix(name))
	}
	body := resp.NewRetryReader(ctx, &blob.RetryReaderOptions{
		MaxRetries: azblobRetryTimes,
	})
	data, err := io.ReadAll(body)
	if err != nil {
		return nil, errors.Annotatef(err, "Failed to read azure blob file, file info: bucket(container)='%s', key='%s'", s.options.Bucket, s.withPrefix(name))
	}
	return data, body.Close()
}