package storage

import (
	"compress/gzip"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"strings"
	"sync"

	gcs "cloud.google.com/go/storage"
	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob"
	"github.com/aliyun/alibabacloud-oss-go-sdk-v2/oss"
	"github.com/aliyun/alibabacloud-oss-go-sdk-v2/oss/credentials"
	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/config"
	"github.com/aws/aws-sdk-go-v2/credentials/stscreds"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/aws-sdk-go-v2/service/secretsmanager"
	"github.com/aws/aws-sdk-go-v2/service/sts"
	"github.com/pingcap/log"
	"go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-sdk-go-v2/otelaws"
	"go.uber.org/zap"
	"google.golang.org/api/iterator"
)

type ObjectMeta struct {
	Key  string
	Size int64
}

type Object struct {
	ETag          *string
	ContentType   *string
	ContentLength *int64
	ContentRange  *string
	AcceptRanges  *string
	Reader        io.ReadCloser
}

type Storage interface {
	Object(ctx context.Context, bucket, key string) (*Object, error)
	Objects(ctx context.Context, bucket, prefix string) ([]ObjectMeta, error)
	ObjectsCallback(ctx context.Context, bucket, prefix string, callback func(ObjectMeta) bool) error
	ObjectsCallbackWithStartKey(ctx context.Context, bucket, prefix, startKey string, callback func(ObjectMeta) bool) error
}

type awsStorage struct {
	client *s3.Client
}

var _ Storage = (*awsStorage)(nil)

func (s *awsStorage) Object(ctx context.Context, bucket, key string) (*Object, error) {
	output, err := s.client.GetObject(ctx, &s3.GetObjectInput{
		Bucket: &bucket,
		Key:    &key,
	})
	if err != nil {
		return nil, err
	}

	// Manually decompress. Ref to https://github.com/aws/aws-sdk-go-v2/blob/main/CHANGELOG.md#service-client-highlights-14
	reader := output.Body
	encoding := output.ContentEncoding
	if encoding != nil && *encoding == "gzip" {
		reader, err = gzip.NewReader(reader)
		if err != nil {
			return nil, err
		}
	}

	return &Object{
		ETag:          output.ETag,
		ContentType:   output.ContentType,
		ContentLength: output.ContentLength,
		ContentRange:  output.ContentRange,
		AcceptRanges:  output.AcceptRanges,
		Reader:        reader,
	}, nil
}

func (s *awsStorage) Objects(ctx context.Context, bucket, prefix string) ([]ObjectMeta, error) {
	input := &s3.ListObjectsV2Input{
		Bucket: &bucket,
		Prefix: &prefix,
	}
	p := s3.NewListObjectsV2Paginator(s.client, input)

	var keys []ObjectMeta
	for p.HasMorePages() {
		// next page takes a context
		page, err := p.NextPage(ctx)
		if err != nil {
			return nil, err
		}
		for _, c := range page.Contents {
			keys = append(keys, ObjectMeta{
				Key:  *c.Key,
				Size: *c.Size,
			})
		}
	}
	return keys, nil
}

func (s *awsStorage) ObjectsCallback(ctx context.Context, bucket, prefix string, callback func(ObjectMeta) bool) error {
	input := &s3.ListObjectsV2Input{
		Bucket: &bucket,
		Prefix: &prefix,
	}
	p := s3.NewListObjectsV2Paginator(s.client, input)
	for p.HasMorePages() {
		// next page takes a context
		page, err := p.NextPage(ctx)
		if err != nil {
			return err
		}
		for _, c := range page.Contents {
			if !callback(ObjectMeta{
				Key:  *c.Key,
				Size: *c.Size,
			}) {
				return nil
			}
		}
	}
	return nil
}

func (s *awsStorage) ObjectsCallbackWithStartKey(
	ctx context.Context, bucket, prefix, startKey string, callback func(ObjectMeta) bool,
) error {
	if !callback(ObjectMeta{Key: startKey}) { // Putting it in the first place only for AWS
		return nil
	}
	input := &s3.ListObjectsV2Input{
		Bucket:     &bucket,
		Prefix:     &prefix,
		StartAfter: &startKey,
	}
	p := s3.NewListObjectsV2Paginator(s.client, input)
	for p.HasMorePages() {
		// next page takes a context
		page, err := p.NextPage(ctx)
		if err != nil {
			return err
		}
		for _, c := range page.Contents {
			if !callback(ObjectMeta{
				Key:  *c.Key,
				Size: *c.Size,
			}) {
				return nil
			}
		}
	}
	return nil
}

type gcpStorage struct {
	client *gcs.Client
}

var _ Storage = (*gcpStorage)(nil)

func (s *gcpStorage) Object(ctx context.Context, bucket, key string) (*Object, error) {
	obj := s.client.Bucket(bucket).Object(key)

	attrs, err := obj.Attrs(ctx)
	if err != nil {
		return nil, err
	}

	reader, err := obj.NewReader(ctx)
	if err != nil {
		return nil, err
	}

	return &Object{
		ETag:          &attrs.Etag,
		ContentLength: &attrs.Size,
		ContentType:   &attrs.ContentType,
		Reader:        reader,
	}, nil
}

func (s *gcpStorage) Objects(ctx context.Context, bucket, prefix string) ([]ObjectMeta, error) {
	iter := s.client.Bucket(bucket).Objects(ctx, &gcs.Query{
		Prefix: prefix,
	})
	var keys []ObjectMeta
	for {
		attr, err := iter.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			return nil, err
		}
		if attr == nil {
			break
		}
		keys = append(keys, ObjectMeta{
			Key:  attr.Name,
			Size: attr.Size,
		})
	}
	return keys, nil
}

func (s *gcpStorage) ObjectsCallback(ctx context.Context, bucket, prefix string, callback func(ObjectMeta) bool) error {
	iter := s.client.Bucket(bucket).Objects(ctx, &gcs.Query{
		Prefix: prefix,
	})
	for {
		attr, err := iter.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			return err
		}
		if attr == nil {
			break
		}
		if !callback(ObjectMeta{
			Key:  attr.Name,
			Size: attr.Size,
		}) {
			return nil
		}
	}
	return nil
}

func (s *gcpStorage) ObjectsCallbackWithStartKey(
	ctx context.Context, bucket, prefix, startKey string, callback func(ObjectMeta) bool,
) error {
	iter := s.client.Bucket(bucket).Objects(ctx, &gcs.Query{
		Prefix:      prefix,
		StartOffset: startKey,
	})
	for {
		attr, err := iter.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			return err
		}
		if attr == nil {
			break
		}
		if !callback(ObjectMeta{
			Key:  attr.Name,
			Size: attr.Size,
		}) {
			return nil
		}
	}
	return nil
}

type azureStorage struct {
	client *azblob.Client
}

var _ Storage = (*azureStorage)(nil)

func (s *azureStorage) Object(ctx context.Context, bucket, key string) (*Object, error) {
	client := s.client.ServiceClient().NewContainerClient(bucket).NewBlobClient(key)
	resp, err := client.DownloadStream(ctx, nil)
	if err != nil {
		return nil, err
	}
	return &Object{
		ETag:          (*string)(resp.ETag),
		ContentType:   resp.ContentType,
		ContentLength: resp.ContentLength,
		ContentRange:  resp.ContentRange,
		AcceptRanges:  resp.AcceptRanges,
		Reader:        resp.Body,
	}, nil
}

func (s *azureStorage) Objects(ctx context.Context, bucket, prefix string) ([]ObjectMeta, error) {
	var objects []ObjectMeta
	client := s.client.ServiceClient().NewContainerClient(bucket)
	pager := client.NewListBlobsFlatPager(&azblob.ListBlobsFlatOptions{Prefix: &prefix})
	for pager.More() {
		resp, err := pager.NextPage(ctx)
		if err != nil {
			return nil, err
		}
		for _, blob := range resp.Segment.BlobItems {
			objects = append(objects, ObjectMeta{
				Key:  *blob.Name,
				Size: *blob.Properties.ContentLength,
			})
		}
	}
	return objects, nil
}

func (s *azureStorage) ObjectsCallback(ctx context.Context, bucket, prefix string, callback func(ObjectMeta) bool) error {
	client := s.client.ServiceClient().NewContainerClient(bucket)
	pager := client.NewListBlobsFlatPager(&azblob.ListBlobsFlatOptions{Prefix: &prefix})
	for pager.More() {
		resp, err := pager.NextPage(ctx)
		if err != nil {
			return err
		}
		for _, blob := range resp.Segment.BlobItems {
			if !callback(ObjectMeta{
				Key:  *blob.Name,
				Size: *blob.Properties.ContentLength,
			}) {
				return nil
			}
		}
	}
	return nil
}

func (s *azureStorage) ObjectsCallbackWithStartKey(ctx context.Context, bucket, prefix, startKey string, callback func(ObjectMeta) bool) error {
	client := s.client.ServiceClient().NewContainerClient(bucket)
	pager := client.NewListBlobsFlatPager(&azblob.ListBlobsFlatOptions{
		Prefix: &prefix,
		Marker: &startKey,
	})
	for pager.More() {
		resp, err := pager.NextPage(ctx)
		if err != nil {
			return err
		}
		for _, blob := range resp.Segment.BlobItems {
			if !callback(ObjectMeta{
				Key:  *blob.Name,
				Size: *blob.Properties.ContentLength,
			}) {
				return nil
			}
		}
	}
	return nil
}

type aliyunStorage struct {
	client    *oss.Client
	devClient *oss.Client
}

var _ Storage = (*aliyunStorage)(nil)

func (s *aliyunStorage) clientOrDevClient(bucket string) *oss.Client {
	if s.devClient != nil && strings.HasPrefix(bucket, "o11y-dev-shared-") {
		return s.devClient
	}
	return s.client
}

func (s *aliyunStorage) Object(ctx context.Context, bucket, key string) (*Object, error) {
	log.Info("AliyunStorage Get Object", zap.String("bucket", bucket), zap.String("key", key))
	resp, err := s.clientOrDevClient(bucket).GetObject(ctx, &oss.GetObjectRequest{
		Bucket: oss.Ptr(bucket),
		Key:    oss.Ptr(key),
	})
	if err != nil {
		return nil, err
	}
	return &Object{
		ETag:          resp.ETag,
		ContentType:   resp.ContentType,
		ContentLength: &resp.ContentLength,
		ContentRange:  resp.ContentRange,
		Reader:        resp.Body,
	}, nil
}

func (s *aliyunStorage) Objects(ctx context.Context, bucket, prefix string) ([]ObjectMeta, error) {
	var keys []ObjectMeta
	var continuationToken *string
	log.Info("AliyunStorage Get Objects", zap.String("bucket", bucket), zap.String("prefix", prefix))
	for {
		input := &oss.ListObjectsV2Request{
			Bucket:            oss.Ptr(bucket),
			Prefix:            oss.Ptr(prefix),
			ContinuationToken: continuationToken,
		}
		resp, err := s.clientOrDevClient(bucket).ListObjectsV2(ctx, input)
		if err != nil {
			return nil, err
		}
		for _, obj := range resp.Contents {
			keys = append(keys, ObjectMeta{
				Key:  *obj.Key,
				Size: obj.Size,
			})
		}
		if resp.IsTruncated {
			continuationToken = resp.NextContinuationToken
		} else {
			break
		}
	}
	return keys, nil
}

func (s *aliyunStorage) ObjectsCallback(ctx context.Context, bucket, prefix string, callback func(ObjectMeta) bool) error {
	var continuationToken *string
	for {
		input := &oss.ListObjectsV2Request{
			Bucket:            oss.Ptr(bucket),
			Prefix:            oss.Ptr(prefix),
			ContinuationToken: continuationToken,
		}
		resp, err := s.clientOrDevClient(bucket).ListObjectsV2(ctx, input)
		if err != nil {
			return err
		}
		for _, obj := range resp.Contents {
			if !callback(ObjectMeta{
				Key:  *obj.Key,
				Size: obj.Size,
			}) {
				return nil
			}
		}
		if resp.IsTruncated {
			continuationToken = resp.NextContinuationToken
		} else {
			break
		}
	}
	return nil
}

func (s *aliyunStorage) ObjectsCallbackWithStartKey(
	ctx context.Context, bucket, prefix, startKey string, callback func(ObjectMeta) bool,
) error {
	var continuationToken *string
	for {
		input := &oss.ListObjectsV2Request{
			Bucket:            oss.Ptr(bucket),
			Prefix:            oss.Ptr(prefix),
			StartAfter:        oss.Ptr(startKey),
			ContinuationToken: continuationToken,
		}
		resp, err := s.clientOrDevClient(bucket).ListObjectsV2(ctx, input)
		if err != nil {
			return err
		}
		for _, obj := range resp.Contents {
			if !callback(ObjectMeta{
				Key:  *obj.Key,
				Size: obj.Size,
			}) {
				return nil
			}
		}
		if resp.IsTruncated {
			continuationToken = resp.NextContinuationToken
		} else {
			break
		}
	}
	return nil
}

type StorageFactory interface {
	Storage(ctx context.Context, provider, region string) (Storage, error)
}

type storageFactory struct {
	sync.Mutex
	gcp    Storage
	aws    *awsStorages
	aliyun *aliyunStorages
	azure  map[string]Storage
}

var _ StorageFactory = (*storageFactory)(nil)

func NewStorageFactory(s3AssumeRoleArn string) (StorageFactory, error) {
	gcsClient, err := gcs.NewClient(context.Background())
	if err != nil {
		return nil, err
	}
	aliyunStores, err := newAliyunStorages("us-west-2")
	if err != nil {
		return nil, err
	}
	eastus2Token, err := azureBlobSasToken("eastus2")
	if err != nil {
		return nil, err
	}
	eastus2Client, err := azblob.NewClientFromConnectionString(eastus2Token, nil)
	if err != nil {
		return nil, err
	}
	southeastasiaToken, err := azureBlobSasToken("southeastasia")
	if err != nil {
		return nil, err
	}
	southeastasiaClient, err := azblob.NewClientFromConnectionString(southeastasiaToken, nil)
	if err != nil {
		return nil, err
	}
	westeuropeToken, err := azureBlobSasToken("westeurope")
	if err != nil {
		return nil, err
	}
	westeuropeClient, err := azblob.NewClientFromConnectionString(westeuropeToken, nil)
	if err != nil {
		return nil, err
	}
	japaneastToken, err := azureBlobSasToken("japaneast")
	if err != nil {
		return nil, err
	}
	japaneastClient, err := azblob.NewClientFromConnectionString(japaneastToken, nil)
	if err != nil {
		return nil, err
	}
	return &storageFactory{
		gcp: &gcpStorage{
			client: gcsClient,
		},
		aws:    newAWSStorages(s3AssumeRoleArn),
		aliyun: aliyunStores,
		azure: map[string]Storage{
			"eastus2": &azureStorage{
				client: eastus2Client,
			},
			"southeastasia": &azureStorage{
				client: southeastasiaClient,
			},
			"westeurope": &azureStorage{
				client: westeuropeClient,
			},
			"japaneast": &azureStorage{
				client: japaneastClient,
			},
		},
	}, nil
}

func NewStorageFactoryForCN(s3AssumeRoleArn string) (StorageFactory, error) {
	aliyunStores, err := newAliyunStorages("cn-north-1")
	if err != nil {
		return nil, err
	}
	return &storageFactory{
		aliyun: aliyunStores,
	}, nil
}

func (f *storageFactory) Storage(ctx context.Context, provider, region string) (Storage, error) {
	switch provider {
	case "aws":
		return f.aws.storage(ctx, region)
	case "gcp":
		return f.gcp, nil
	case "alicloud", "aliyun":
		return f.aliyun.storage(ctx, region)
	case "azure":
		if store, ok := f.azure[region]; ok {
			return store, nil
		}
	}
	return nil, fmt.Errorf("unsupported provider '%s' and region '%s'", provider, region)
}

type awsStorages struct {
	sync.Mutex
	assumeRoleArn string
	stores        map[string]Storage
}

func newAWSStorages(assumeRoleArn string) *awsStorages {
	return &awsStorages{
		assumeRoleArn: assumeRoleArn,
		stores:        map[string]Storage{},
	}
}

func (s *awsStorages) storage(ctx context.Context, region string) (Storage, error) {
	s.Lock()
	defer s.Unlock()

	if client, ok := s.stores[region]; ok {
		return client, nil
	}

	cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region))
	if err != nil {
		return nil, err
	}

	if s.assumeRoleArn != "" {
		stsSvc := sts.NewFromConfig(cfg)
		creds := stscreds.NewAssumeRoleProvider(stsSvc, s.assumeRoleArn)
		cfg.Credentials = aws.NewCredentialsCache(creds)
	}

	// Instrument with OpenTelemetry
	otelaws.AppendMiddlewares(&cfg.APIOptions)

	client := &awsStorage{client: s3.NewFromConfig(cfg)}
	s.stores[region] = client
	return client, nil
}

type aliyunStorages struct {
	sync.Mutex
	stores             map[string]Storage
	accessKeyID        string
	accessKeySecret    string
	devAccessKeyID     string
	devAccessKeySecret string
}

func newAliyunStorages(clincRegion string) (*aliyunStorages, error) {
	ak, sk, err := aliyunToken(clincRegion, "aliyun-oss-token")
	if err != nil {
		return nil, err
	}
	devAk, devSk, _ := aliyunToken(clincRegion, "aliyun-oss-token-dev")
	return &aliyunStorages{
		stores:             map[string]Storage{},
		accessKeyID:        ak,
		accessKeySecret:    sk,
		devAccessKeyID:     devAk,
		devAccessKeySecret: devSk,
	}, nil
}

func (s *aliyunStorages) storage(_ context.Context, region string) (Storage, error) {
	s.Lock()
	defer s.Unlock()
	if client, ok := s.stores[region]; ok {
		return client, nil
	}
	client := &aliyunStorage{
		client: oss.NewClient(oss.NewConfig().
			WithRegion(region).
			WithCredentialsProvider(
				credentials.NewStaticCredentialsProvider(s.accessKeyID, s.accessKeySecret),
			)),
	}

	if s.devAccessKeyID != "" && s.devAccessKeySecret != "" {
		client.devClient = oss.NewClient(oss.NewConfig().
			WithRegion(region).
			WithCredentialsProvider(
				credentials.NewStaticCredentialsProvider(s.devAccessKeyID, s.devAccessKeySecret),
			))
	}
	s.stores[region] = client
	return client, nil
}

type azureBlobSasTokenSecret struct {
	Token string `json:"token"`
}

func azureBlobSasToken(region string) (string, error) {
	secretName := "azure-blob-sas-token-o11yprod" + region
	config, err := config.LoadDefaultConfig(context.TODO(), config.WithRegion("us-west-2"))
	if err != nil {
		return "", err
	}
	svc := secretsmanager.NewFromConfig(config)
	input := &secretsmanager.GetSecretValueInput{
		SecretId:     aws.String(secretName),
		VersionStage: aws.String("AWSCURRENT"),
	}
	result, err := svc.GetSecretValue(context.TODO(), input)
	if err != nil {
		return "", err
	}
	if result.SecretString == nil {
		return "", fmt.Errorf("secret string is nil")
	}
	var s azureBlobSasTokenSecret
	if err := json.Unmarshal([]byte(*result.SecretString), &s); err != nil {
		return "", err
	}
	return s.Token, nil
}

type aliyunTokenSecret struct {
	AccessKeyID     string `json:"access_key_id"`
	AccessKeySecret string `json:"access_key_secret"`
}

func aliyunToken(clincRegion, secretKeyName string) (string, string, error) {
	config, err := config.LoadDefaultConfig(context.TODO(), config.WithRegion(clincRegion))
	if err != nil {
		return "", "", err
	}
	svc := secretsmanager.NewFromConfig(config)
	input := &secretsmanager.GetSecretValueInput{
		SecretId:     aws.String(secretKeyName),
		VersionStage: aws.String("AWSCURRENT"),
	}
	result, err := svc.GetSecretValue(context.TODO(), input)
	if err != nil {
		return "", "", err
	}
	if result.SecretString == nil {
		return "", "", fmt.Errorf("secret string is nil")
	}
	var s aliyunTokenSecret
	if err := json.Unmarshal([]byte(*result.SecretString), &s); err != nil {
		return "", "", err
	}
	return s.AccessKeyID, s.AccessKeySecret, nil
}