package objectstore

import (
	"context"
	"errors"
	"io"
	"path"
	"strings"

	"github.com/aws/aws-sdk-go-v2/aws"
	awsconfig "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"
	"github.com/aws/smithy-go"
)

type S3V2Client interface {
	PutObject(context.Context, *s3.PutObjectInput, ...func(*s3.Options)) (*s3.PutObjectOutput, error)
	GetObject(context.Context, *s3.GetObjectInput, ...func(*s3.Options)) (*s3.GetObjectOutput, error)
	HeadObject(context.Context, *s3.HeadObjectInput, ...func(*s3.Options)) (*s3.HeadObjectOutput, error)
	ListObjectsV2(context.Context, *s3.ListObjectsV2Input, ...func(*s3.Options)) (*s3.ListObjectsV2Output, error)
}

type S3V2Store struct {
	client S3V2Client
	bucket string
	prefix string
}

func NewS3V2Client(ctx context.Context, region string, roleARN string) (*s3.Client, error) {
	cfg, err := awsconfig.LoadDefaultConfig(ctx, awsconfig.WithRegion(region))
	if err != nil {
		return nil, err
	}
	if roleARN != "" {
		stsClient := sts.NewFromConfig(cfg)
		cfg.Credentials = aws.NewCredentialsCache(stscreds.NewAssumeRoleProvider(stsClient, roleARN))
	}
	return s3.NewFromConfig(cfg), nil
}

func NewS3V2Store(client S3V2Client, bucket string, prefix string) *S3V2Store {
	return &S3V2Store{
		client: client,
		bucket: bucket,
		prefix: strings.Trim(prefix, "/"),
	}
}

func (s *S3V2Store) Put(ctx context.Context, key string, body io.Reader, opts PutOptions) (ObjectInfo, error) {
	fullKey := s.fullKey(key)
	input := &s3.PutObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(fullKey),
		Body:   body,
	}
	if opts.ContentType != "" {
		input.ContentType = aws.String(opts.ContentType)
	}
	if opts.ContentEncoding != "" {
		input.ContentEncoding = aws.String(opts.ContentEncoding)
	}
	if len(opts.Metadata) != 0 {
		input.Metadata = opts.Metadata
	}
	out, err := s.client.PutObject(ctx, input)
	if err != nil {
		return ObjectInfo{}, err
	}
	return ObjectInfo{
		Key:      fullKey,
		ETag:     aws.ToString(out.ETag),
		Metadata: opts.Metadata,
	}, nil
}

func (s *S3V2Store) Get(ctx context.Context, key string) (io.ReadCloser, ObjectInfo, error) {
	fullKey := s.fullKey(key)
	out, err := s.client.GetObject(ctx, &s3.GetObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(fullKey),
	})
	if err != nil {
		return nil, ObjectInfo{}, err
	}
	return out.Body, ObjectInfo{
		Key:          fullKey,
		ETag:         aws.ToString(out.ETag),
		Size:         aws.ToInt64(out.ContentLength),
		LastModified: out.LastModified,
		Metadata:     out.Metadata,
	}, nil
}

func (s *S3V2Store) Head(ctx context.Context, key string) (ObjectInfo, bool, error) {
	fullKey := s.fullKey(key)
	out, err := s.client.HeadObject(ctx, &s3.HeadObjectInput{
		Bucket: aws.String(s.bucket),
		Key:    aws.String(fullKey),
	})
	if err != nil {
		if isNotFound(err) {
			return ObjectInfo{}, false, nil
		}
		return ObjectInfo{}, false, err
	}
	return ObjectInfo{
		Key:          fullKey,
		ETag:         aws.ToString(out.ETag),
		Size:         aws.ToInt64(out.ContentLength),
		LastModified: out.LastModified,
		Metadata:     out.Metadata,
	}, true, nil
}

func (s *S3V2Store) List(ctx context.Context, prefix string) ([]ObjectInfo, error) {
	fullPrefix := s.fullKey(prefix)
	ret := make([]ObjectInfo, 0)
	var continuationToken *string
	for {
		out, err := s.client.ListObjectsV2(ctx, &s3.ListObjectsV2Input{
			Bucket:            aws.String(s.bucket),
			Prefix:            aws.String(fullPrefix),
			ContinuationToken: continuationToken,
		})
		if err != nil {
			return nil, err
		}
		for _, item := range out.Contents {
			ret = append(ret, ObjectInfo{
				Key:          aws.ToString(item.Key),
				Size:         aws.ToInt64(item.Size),
				LastModified: item.LastModified,
				ETag:         aws.ToString(item.ETag),
			})
		}
		if !aws.ToBool(out.IsTruncated) {
			break
		}
		continuationToken = out.NextContinuationToken
		if continuationToken == nil {
			break
		}
	}
	return ret, nil
}

func (s *S3V2Store) fullKey(key string) string {
	key = strings.TrimLeft(key, "/")
	if s.prefix == "" {
		return key
	}
	if key == "" {
		return s.prefix
	}
	if key == s.prefix || strings.HasPrefix(key, s.prefix+"/") {
		return key
	}
	return path.Join(s.prefix, key)
}

func isNotFound(err error) bool {
	var apiErr smithy.APIError
	if errors.As(err, &apiErr) {
		return apiErr.ErrorCode() == "NotFound" || apiErr.ErrorCode() == "NoSuchKey"
	}
	return false
}