package trafficjobclient

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

	"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/s3/types"
	"github.com/aws/aws-sdk-go-v2/service/sts"
)

var ErrNotFound = errors.New("object not found")

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)
	ObjectExists(ctx context.Context, bucket, key string) (bool, error)
	ObjectSize(ctx context.Context, bucket, key string) (int64, error)
	Objects(ctx context.Context, bucket, prefix string) ([]ObjectMeta, error)
	ObjectsDelimiter(ctx context.Context, bucket, prefix string) ([]string, 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
	SaveObject(ctx context.Context, bucket, key string, body io.Reader) error
}

type awsStorage struct {
	client *s3.Client
}

var _ Storage = (*awsStorage)(nil)

func NewAWSStorage(ctx context.Context, region, assumeRoleARN string) Storage {
	cfg, err := config.LoadDefaultConfig(ctx, config.WithRegion("us-west-2"))
	if err != nil {
		panic(err)
	}
	if assumeRoleARN != "" {
		stsSvc := sts.NewFromConfig(cfg)
		creds := stscreds.NewAssumeRoleProvider(stsSvc, assumeRoleARN)
		cfg.Credentials = aws.NewCredentialsCache(creds)
	}
	return &awsStorage{
		client: s3.NewFromConfig(cfg),
	}
}

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 {
		var t *types.NoSuchKey
		if errors.As(err, &t) {
			return nil, ErrNotFound
		}
		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) ObjectExists(ctx context.Context, bucket, key string) (bool, error) {
	_, err := s.client.GetObjectAttributes(ctx, &s3.GetObjectAttributesInput{
		Bucket: &bucket,
		Key:    &key,
		ObjectAttributes: []types.ObjectAttributes{
			types.ObjectAttributesObjectSize,
		},
	})
	if err != nil {
		var t *types.NoSuchKey
		if errors.As(err, &t) {
			return false, nil
		}
		return false, err
	}
	return true, nil
}

func (s *awsStorage) ObjectSize(ctx context.Context, bucket, key string) (int64, error) {
	output, err := s.client.GetObjectAttributes(ctx, &s3.GetObjectAttributesInput{
		Bucket: &bucket,
		Key:    &key,
		ObjectAttributes: []types.ObjectAttributes{
			types.ObjectAttributesObjectSize,
		},
	})
	if err != nil {
		var t *types.NoSuchKey
		if errors.As(err, &t) {
			return 0, nil
		}
		return 0, err
	}
	return *output.ObjectSize, 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) ObjectsDelimiter(ctx context.Context, bucket, prefix string) ([]string, error) {
	input := &s3.ListObjectsV2Input{
		Bucket:    &bucket,
		Prefix:    &prefix,
		Delimiter: aws.String("/"),
	}
	output, err := s.client.ListObjectsV2(ctx, input)
	if err != nil {
		return nil, err
	}
	if len(output.CommonPrefixes) == 0 {
		return nil, nil
	}
	prefixes := make([]string, len(output.CommonPrefixes))
	for n, p := range output.CommonPrefixes {
		prefixes[n] = *p.Prefix
	}
	return prefixes, 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
}

func (s *awsStorage) SaveObject(ctx context.Context, bucket, key string, body io.Reader) error {
	_, err := s.client.PutObject(ctx, &s3.PutObjectInput{
		Bucket: &bucket,
		Key:    &key,
		Body:   body,
	})
	return err
}