package aws

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"strings"

	"github.com/tidbcloud/observability/pkg/common/message"
	"github.com/tidbcloud/observability/pkg/common/utils"
	"github.com/tidbcloud/observability/pkg/resourceworker/providers/lifecyclecompare"
	"github.com/tidbcloud/observability/pkg/resourceworker/providers/types"

	"github.com/tidbcloud/observability/pkg/common/metrics"

	"github.com/aws/aws-sdk-go-v2/aws"
	s3api "github.com/aws/aws-sdk-go-v2/service/s3"
	s3types "github.com/aws/aws-sdk-go-v2/service/s3/types"
)

const defaultRegion = "us-east-1"

const (
	s3ClientLabel         = "s3-client"
	errCreateBucketFailed = "create bucket failed"
	errHeadBucketFailed   = "determine if a bucket exists failed"
	errGetObjectFailed    = "get object failed"
	errPutObjectFailed    = "put object failed"
	errGetLifecycleFailed = "get lifecycle failed"
	errPutLifecycleFailed = "put lifecycle failed"
	errDeleteObjectFailed = "delete object failed"
	s3DeletionBatchSize   = 1000
)

type S3Client struct {
	cli s3Client
	cfg aws.Config
}

func NewS3(cfg aws.Config) S3Client {
	c := s3api.NewFromConfig(cfg)
	return S3Client{
		cli: c,
		cfg: cfg,
	}
}

// CreateBucket used to create a new bucket in the S3
func (s3 S3Client) CreateBucket(ctx context.Context, name string) error {
	input := &s3api.CreateBucketInput{
		Bucket: aws.String(name),
		ACL:    s3types.BucketCannedACLPrivate,
	}
	if s3.cfg.Region != defaultRegion {
		input.CreateBucketConfiguration = &s3types.CreateBucketConfiguration{LocationConstraint: s3types.BucketLocationConstraint(s3.cfg.Region)}
	}
	_, err := s3.cli.CreateBucket(ctx, input)
	if err != nil {
		metrics.HandleAWSAPIRequestMetrics(s3.cfg.Region, s3ClientLabel, "CreateBucket", name, errCreateBucketFailed)
	}
	return err
}

// BucketExists returns if a bucket exists in the S3
func (s3 S3Client) BucketExists(ctx context.Context, name string) (bool, error) {
	_, err := s3.cli.HeadBucket(ctx, &s3api.HeadBucketInput{
		Bucket: aws.String(name),
	})
	if err != nil {
		var nf *s3types.NotFound
		if errors.As(err, &nf) {
			return false, nil
		}
		metrics.HandleAWSAPIRequestMetrics(s3.cfg.Region, s3ClientLabel, "HeadBucket", name, errHeadBucketFailed)
		return false, err
	}
	return true, nil
}

func (s3 S3Client) Get(ctx context.Context, bucketName string, objectName string) (*io.ReadCloser, error) {
	out, err := s3.cli.GetObject(ctx, &s3api.GetObjectInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(objectName),
	})
	if err != nil {
		var nf *s3types.NoSuchKey
		if errors.As(err, &nf) {
			return nil, types.ObjectNotFoundError
		}
		metrics.HandleAWSAPIRequestMetrics(s3.cfg.Region, s3ClientLabel, "GetObject", objectName, errGetObjectFailed)
		return nil, err
	}
	return &out.Body, nil
}

// ListObjects returns a list of objects in the bucket with the given prefix.
func (s3 S3Client) ListObjects(ctx context.Context, bucketName string, prefix string, options types.ListOptions) ([]string, error) {
	paginator := s3api.NewListObjectsV2Paginator(s3.cli, &s3api.ListObjectsV2Input{
		Bucket:  aws.String(bucketName),
		Prefix:  aws.String(prefix),
		MaxKeys: options.MaxKeys,
	})
	objects := make([]string, 0)
	for paginator.HasMorePages() {
		output, err := paginator.NextPage(ctx)
		if err != nil {
			return nil, err
		}
		for _, object := range output.Contents {
			objects = append(objects, *object.Key)
		}
	}
	return objects, nil
}

func (s3 S3Client) Put(ctx context.Context, bucketName string, objectName string, reader *bytes.Reader) error {
	_, err := s3.cli.PutObject(ctx, &s3api.PutObjectInput{
		Bucket: aws.String(bucketName),
		Key:    aws.String(objectName),
		Body:   reader,
	})
	if err != nil {
		metrics.HandleAWSAPIRequestMetrics(s3.cfg.Region, s3ClientLabel, "PutObject", objectName, errPutObjectFailed)
		return err
	}

	return nil
}

func (s3 S3Client) AddLifecycle(ctx context.Context, bucketName string, rules []message.LifecycleRuleParam, unusedRulePrefixes []string) error {
	o, err := s3.cli.GetBucketLifecycleConfiguration(ctx, &s3api.GetBucketLifecycleConfigurationInput{
		Bucket: &bucketName,
	})
	if err != nil {
		if !strings.Contains(err.Error(), "NoSuchLifecycleConfiguration") {
			metrics.HandleAWSAPIRequestMetrics(s3.cfg.Region, s3ClientLabel, "GetLifecycle", bucketName, errGetLifecycleFailed)
			return err
		} else {
			o = &s3api.GetBucketLifecycleConfigurationOutput{
				Rules: []s3types.LifecycleRule{},
			}
		}
	}

	newRules := make([]s3types.LifecycleRule, 0)
	for _, rule := range rules {
		newRules = append(newRules, buildLifecycle(rule.RuleName, rule.Prefix, rule.ColdDays, rule.ExpirationDays))
	}
	for _, oldRule := range o.Rules {
		drop := false
		for _, prefix := range unusedRulePrefixes {
			if strings.HasPrefix(*oldRule.ID, prefix) {
				drop = true
			}
		}
		if drop {
			continue
		}
		for _, newRule := range newRules {
			if *newRule.ID == *oldRule.ID {
				drop = true
			}
		}
		if drop {
			continue
		}
		newRules = append(newRules, oldRule)
	}

	if lifecyclecompare.AWSLifecycleRulesEqual(newRules, o.Rules) {
		fmt.Printf("lifecycle rules unchanged for bucket %s, skip update\n", bucketName)
		return nil
	}

	_, err = s3.cli.PutBucketLifecycleConfiguration(ctx, &s3api.PutBucketLifecycleConfigurationInput{
		Bucket: aws.String(bucketName),
		LifecycleConfiguration: &s3types.BucketLifecycleConfiguration{
			Rules: newRules,
		},
	})
	if err != nil {
		metrics.HandleAWSAPIRequestMetrics(s3.cfg.Region, s3ClientLabel, "PutLifecycle", bucketName, errPutLifecycleFailed)
		return err
	}

	return nil
}

func buildLifecycle(ruleName string, prefix string, coldDays int32, expirationDays int32) s3types.LifecycleRule {
	rule := s3types.LifecycleRule{
		ID: &ruleName,
		Filter: &s3types.LifecycleRuleFilter{
			Prefix: &prefix,
		},
		Status: s3types.ExpirationStatusEnabled,
	}
	if expirationDays > 0 {
		rule.Expiration = &s3types.LifecycleExpiration{
			Days: utils.Int32Ptr(expirationDays),
		}
	}

	if coldDays > 0 {
		rule.Transitions = []s3types.Transition{
			{
				Days:         utils.Int32Ptr(coldDays),
				StorageClass: s3types.TransitionStorageClassGlacierIr,
			},
		}
	}

	return rule
}

func (s3 S3Client) RemoveLifecycle(ctx context.Context, bucketName string, name string) error {
	return nil
}

// Delete removes objects from the bucket.
func (s3 S3Client) Delete(ctx context.Context, bucketName string, objectNames ...string) error {
	for i := 0; i < len(objectNames); i += s3DeletionBatchSize {
		batch := objectNames[i:min(i+s3DeletionBatchSize, len(objectNames))]
		err := s3.deleteObjects(ctx, bucketName, batch)
		if err != nil {
			return err
		}
	}
	return nil
}

func (s3 S3Client) deleteObjects(ctx context.Context, bucketName string, objectNames []string) error {
	objects := make([]s3types.ObjectIdentifier, 0)
	for _, objectName := range objectNames {
		objects = append(objects, s3types.ObjectIdentifier{
			Key: aws.String(objectName),
		})
	}
	_, err := s3.cli.DeleteObjects(ctx, &s3api.DeleteObjectsInput{
		Bucket: aws.String(bucketName),
		Delete: &s3types.Delete{
			Objects: objects,
		},
	})

	if err != nil {
		metrics.HandleAWSAPIRequestMetrics(s3.cfg.Region, s3ClientLabel, "DeleteObjects", bucketName, errDeleteObjectFailed)
		return err
	}
	return nil
}

func MockS3(cli s3Client) S3Client {
	return S3Client{
		cli: cli,
	}
}