package s3

import (
	"context"
	"errors"
	"fmt"
	"net/url"
	"os"
	"path"

	awsv2 "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/feature/s3/manager"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	types "github.com/aws/aws-sdk-go-v2/service/s3/types"
	"github.com/aws/aws-sdk-go-v2/service/sts"
	"github.com/sirupsen/logrus"

	"github.com/tidbcloud/observability/backup-worker/pkg/common"
	"github.com/tidbcloud/observability/backup-worker/pkg/configs"
)

// Client wraps the AWS S3 SDK with backup-specific operations: existence checks,
// file uploads, object listing, and cross-bucket copying with multipart support.
type Client struct {
	cli      *s3.Client
	cfg      configs.S3Sink
	rp       configs.RetryPolicy
	uploader *manager.Uploader
}

// New creates an S3 client with optional IAM role assumption.
// Credential chain:
// 1. If RoleARN is configured, assume that role using the default credential chain
// 2. Otherwise, use the default credential chain (env vars, EC2 IAM role, etc.)
//
// ExternalID provides an additional security layer for cross-account role assumption,
// preventing the "confused deputy" attack in multi-tenant scenarios.
func New(ctx context.Context, cfg configs.S3Sink, retryPolicy configs.RetryPolicy) (*Client, error) {
	baseCfg, err := config.LoadDefaultConfig(ctx, config.WithRegion(cfg.Region))
	if err != nil {
		return nil, fmt.Errorf("load aws config: %w", err)
	}

	s3Cfg := baseCfg

	// Set up role assumption only when RoleARN is configured.
	// This enables cross-account backup scenarios where the backup service
	// needs to write to buckets in different AWS accounts.
	if cfg.RoleARN != "" {
		stsClient := sts.NewFromConfig(baseCfg)
		provider := stscreds.NewAssumeRoleProvider(stsClient, cfg.RoleARN, func(o *stscreds.AssumeRoleOptions) {
			if cfg.ExternalID != "" {
				o.ExternalID = &cfg.ExternalID
			}
		})
		s3Cfg.Credentials = awsv2.NewCredentialsCache(provider)
	}

	cli := s3.NewFromConfig(s3Cfg)
	logrus.Infof("S3 client created: bucket=%s, region=%s", cfg.Bucket, cfg.Region)
	return &Client{
		cli:      cli,
		cfg:      cfg,
		rp:       retryPolicy,
		uploader: manager.NewUploader(cli),
	}, nil
}

func (c *Client) Type() string {
	return common.StorageS3
}

// CheckObject checks whether an object already exists in the bucket.
//
// Existence checks enable resumable, idempotent backups: each shard can be
// skipped if its target object was already written in a previous run,
// avoiding redundant uploads after interruptions.
func (c *Client) CheckObject(ctx context.Context, key string) (bool, error) {
	normalizeKey := c.normalizePath(key)
	_, err := c.cli.HeadObject(ctx, &s3.HeadObjectInput{
		Bucket: &c.cfg.Bucket,
		Key:    normalizeKey,
	})
	if err == nil {
		return true, nil
	}
	var notFound *types.NotFound
	if errors.As(err, &notFound) {
		return false, nil
	}
	return false, err
}

// SinkObject uploads a local file to S3 with automatic retries.
//
// Retries handle transient failures (network timeouts, throttling) that are
// common in distributed storage systems, without which a single blip would
// fail the entire shard.
func (c *Client) SinkObject(ctx context.Context, key, localPath string) error {
	normalizeKey := c.normalizePath(key)
	return c.rp.Do(ctx, func() error {
		return c.sinkObject(ctx, normalizeKey, localPath)
	})
}

func (c *Client) sinkObject(ctx context.Context, key *string, localPath string) error {
	f, err := os.Open(localPath)
	if err != nil {
		return fmt.Errorf("open local file: %w", err)
	}
	defer f.Close()

	_, err = c.uploader.Upload(ctx, &s3.PutObjectInput{
		Bucket: &c.cfg.Bucket,
		Key:    key,
		Body:   f,
	})
	if err != nil {
		return fmt.Errorf("upload s3 object: %w", err)
	}
	return nil
}

// ListObjects paginates through all objects under the given prefix.
func (c *Client) ListObjects(ctx context.Context, prefix string) ([]common.ObjectInfo, error) {
	var objects []common.ObjectInfo
	paginator := s3.NewListObjectsV2Paginator(c.cli, &s3.ListObjectsV2Input{
		Bucket: &c.cfg.Bucket,
		Prefix: &prefix,
	})
	for paginator.HasMorePages() {
		page, err := paginator.NextPage(ctx)
		if err != nil {
			return nil, fmt.Errorf("list '%s' objects: %w", prefix, err)
		}
		for _, content := range page.Contents {
			objects = append(objects, common.ObjectInfo{
				Name: *content.Key,
				Size: *content.Size,
			})
		}
	}
	return objects, nil
}

// CopyObject copies an object from another bucket into this client's bucket.
// Retries on failure and aborts any in-progress multipart upload to avoid
// orphaned parts that would consume storage and incur costs.
func (c *Client) CopyObject(ctx context.Context, srcBucket, srcKey, destKey string, size int64) error {
	return c.rp.Do(ctx, func() error {
		return c.copyObject(ctx, srcBucket, srcKey, destKey, size)
	})
}

func (c *Client) copyObject(ctx context.Context, srcBucket, srcKey, destKey string, size int64) error {
	copySource := srcBucket + "/" + url.PathEscape(srcKey)

	// Small files use CopyObject API (server-side copy), large files use multipart copy
	const partSize int64 = 5 * 1024 * 1024 * 1024 // 5 GB
	if size <= partSize {
		_, err := c.cli.CopyObject(ctx, &s3.CopyObjectInput{
			Bucket:     &c.cfg.Bucket,
			Key:        &destKey,
			CopySource: &copySource,
		})
		if err != nil {
			return fmt.Errorf("copy s3 object: %w", err)
		}
		return nil
	}

	// Large files use multipart copy
	mpu, err := c.cli.CreateMultipartUpload(ctx, &s3.CreateMultipartUploadInput{
		Bucket: &c.cfg.Bucket,
		Key:    &destKey,
	})
	if err != nil {
		return fmt.Errorf("create multipart upload: %w", err)
	}

	cleanup := func() {
		// Abort on any failure to avoid orphaned multipart uploads.
		_, _ = c.cli.AbortMultipartUpload(ctx, &s3.AbortMultipartUploadInput{
			Bucket:   &c.cfg.Bucket,
			Key:      &destKey,
			UploadId: mpu.UploadId,
		})
	}

	var completedParts []types.CompletedPart
	var partNumber int32 = 1

	for offset := int64(0); offset < size; offset += partSize {
		pn := partNumber // capture loop variable for CompletedPart
		end := offset + partSize - 1
		if end >= size {
			end = size - 1
		}
		rangeStr := fmt.Sprintf("bytes=%d-%d", offset, end)

		partResult, err := c.cli.UploadPartCopy(ctx, &s3.UploadPartCopyInput{
			Bucket:          &c.cfg.Bucket,
			Key:             &destKey,
			UploadId:        mpu.UploadId,
			PartNumber:      &pn,
			CopySource:      &copySource,
			CopySourceRange: &rangeStr,
		})
		if err != nil {
			cleanup()
			return fmt.Errorf("upload part %d: %w", pn, err)
		}

		completedParts = append(completedParts, types.CompletedPart{
			PartNumber: &pn,
			ETag:       partResult.CopyPartResult.ETag,
		})
		partNumber++
	}

	_, err = c.cli.CompleteMultipartUpload(ctx, &s3.CompleteMultipartUploadInput{
		Bucket:          &c.cfg.Bucket,
		Key:             &destKey,
		UploadId:        mpu.UploadId,
		MultipartUpload: &types.CompletedMultipartUpload{Parts: completedParts},
	})
	if err != nil {
		cleanup()
	}
	return err
}

func (c *Client) Close() error {
	return nil
}

func (c *Client) normalizePath(key string) *string {
	p := path.Join(c.cfg.PathPrefix, key)
	return &p
}