package objstore

import (
	"compress/gzip"
	"context"
	"fmt"
	"io"

	gcs "cloud.google.com/go/storage"
	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
	"github.com/Azure/azure-sdk-for-go/sdk/storage/azblob"
	"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/sts"
	"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
}

func NewAwsStorage(ctx context.Context, region, assumeRoleArn string) (Storage, error) {
	cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(region))
	if err != nil {
		return nil, err
	}
	if assumeRoleArn != "" {
		stsSvc := sts.NewFromConfig(cfg)
		creds := stscreds.NewAssumeRoleProvider(stsSvc, assumeRoleArn)
		cfg.Credentials = aws.NewCredentialsCache(creds)
	}
	return &awsStorage{client: s3.NewFromConfig(cfg)}, 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
}

func NewGcpStorage(ctx context.Context) (Storage, error) {
	client, err := gcs.NewClient(ctx)
	if err != nil {
		return nil, err
	}
	return &gcpStorage{client: client}, 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
}

func NewAzureStorage(ctx context.Context, storageAccount string) (Storage, error) {
	cred, err := azidentity.NewDefaultAzureCredential(nil)
	if err != nil {
		return nil, err
	}
	client, err := azblob.NewClient(fmt.Sprintf("https://%s.blob.core.windows.net/", storageAccount), cred, nil)
	if err != nil {
		return nil, err
	}
	return &azureStorage{client: client}, 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 {
	// Azure List Blobs has no StartAfter parameter in the current Go SDK.
	// Marker is an opaque pagination token, not a blob key, so we list by prefix
	// and skip keys lexicographically <= startKey to match S3 StartAfter semantics.
	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 {
			key := *blob.Name
			if key <= startKey {
				continue
			}
			if !callback(ObjectMeta{
				Key:  key,
				Size: *blob.Properties.ContentLength,
			}) {
				return nil
			}
		}
	}
	return nil
}