package impl

import (
	"context"
	"encoding/csv"
	"errors"
	"fmt"
	"io"
	"strings"
	"time"

	"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"
	s3manager "github.com/aws/aws-sdk-go-v2/feature/s3/manager"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/aws/aws-sdk-go-v2/service/s3control"
	s3ctypes "github.com/aws/aws-sdk-go-v2/service/s3control/types"
	"github.com/aws/aws-sdk-go-v2/service/sts"
	"github.com/google/uuid"
	"github.com/tidbcloud/dataflow-service/app/internal/aws_facade"
	dfaws "github.com/tidbcloud/dataflow-service/pkg/aws"
	"github.com/tidbcloud/dataflow-service/pkg/log"
	"go.uber.org/zap"
)

// batchOpsJobPriority defines priority used for submitted S3 Batch Ops jobs.
const batchOpsJobPriority int32 = 10

func GenerateDPConfig(ctx context.Context, accountID, region string) (aws.Config, error) {
	baseCfg, err := awsconfig.LoadDefaultConfig(ctx, awsconfig.WithRegion(region))
	if err != nil {
		return aws.Config{}, fmt.Errorf("load base aws config: %w", err)
	}
	roleArn := (&dfaws.GetAWSClientReq{AccountID: accountID, Region: region}).GenerateDataflowRoleArn()
	stsClient := sts.NewFromConfig(baseCfg)
	baseCfg.Credentials = aws.NewCredentialsCache(stscreds.NewAssumeRoleProvider(stsClient, roleArn))
	return baseCfg, nil
}

var _ aws_facade.BatchOpsClient = (*BatchOpsClientImpl)(nil)

type BatchOpsClientImpl struct{}

func NewBatchOpsClient() *BatchOpsClientImpl {
	return &BatchOpsClientImpl{}
}

// GenerateAndUploadManifest streams a CSV manifest for all objects under the source prefix directly to S3.
func (c *BatchOpsClientImpl) GenerateAndUploadManifest(
	ctx context.Context, region, accountID, sourceBucket, sourcePrefix, manifestBucket, manifestKey string,
) (string, error) {
	logger := log.GetLogger(ctx).Named("aws-batch-ops")
	startAt := time.Now()
	logger.Info("generate and upload manifest start",
		zap.String("region", region),
		zap.String("source_bucket", sourceBucket),
		zap.String("source_prefix", sourcePrefix),
		zap.String("manifest_bucket", manifestBucket),
		zap.String("manifest_key", manifestKey))

	cfg, err := GenerateDPConfig(ctx, accountID, region)
	if err != nil {
		return "", fmt.Errorf("load aws config: %w", err)
	}

	s3Client := s3.NewFromConfig(cfg)
	reader, writer := io.Pipe()
	resultCh := make(chan manifestResult, 1)
	go func() {
		stats, err := writeManifest(ctx, writer, s3Client, sourceBucket, sourcePrefix)
		if err == nil && stats.objects == 0 {
			err = fmt.Errorf("manifest is empty for source %s/%s", sourceBucket, sourcePrefix)
		}
		if err != nil {
			_ = writer.CloseWithError(err)
			resultCh <- manifestResult{err: err}
			return
		}
		_ = writer.Close()
		resultCh <- manifestResult{stats: stats}
	}()

	uploader := s3manager.NewUploader(s3Client)
	output, uploadErr := uploader.Upload(ctx, &s3.PutObjectInput{
		Bucket:      aws.String(manifestBucket),
		Key:         aws.String(manifestKey),
		Body:        reader,
		ContentType: aws.String("text/csv"),
	})
	if uploadErr != nil {
		_ = reader.CloseWithError(uploadErr)
	}
	result := <-resultCh
	if err := errors.Join(uploadErr, result.err); err != nil {
		return "", fmt.Errorf("generate and upload manifest: %w", err)
	}

	etag := strings.Trim(aws.ToString(output.ETag), "\"")
	if etag == "" {
		return "", fmt.Errorf("manifest etag is empty")
	}
	logger.Info("generate and upload manifest finished",
		zap.Int64("manifest_objects", result.stats.objects),
		zap.Int64("manifest_bytes", result.stats.bytes),
		zap.Duration("elapsed", time.Since(startAt)))
	return etag, nil
}

type manifestStats struct {
	objects int64
	bytes   int64
}

type manifestResult struct {
	stats manifestStats
	err   error
}

type countingWriter struct {
	io.Writer
	bytes int64
}

func (w *countingWriter) Write(p []byte) (int, error) {
	n, err := w.Writer.Write(p)
	w.bytes += int64(n)
	return n, err
}

func writeManifest(ctx context.Context, output io.Writer, s3Client *s3.Client, sourceBucket, sourcePrefix string) (manifestStats, error) {
	input := &s3.ListObjectsV2Input{
		Bucket: aws.String(sourceBucket),
		Prefix: aws.String(sourcePrefix),
	}
	return writeManifestPages(ctx, output, sourceBucket, &s3ManifestPaginator{
		paginator: s3.NewListObjectsV2Paginator(s3Client, input),
	})
}

type manifestPaginator interface {
	HasMorePages() bool
	NextPage(context.Context) (*s3.ListObjectsV2Output, error)
}

type s3ManifestPaginator struct {
	paginator *s3.ListObjectsV2Paginator
}

func (p *s3ManifestPaginator) HasMorePages() bool {
	return p.paginator.HasMorePages()
}

func (p *s3ManifestPaginator) NextPage(ctx context.Context) (*s3.ListObjectsV2Output, error) {
	return p.paginator.NextPage(ctx)
}

func writeManifestPages(ctx context.Context, output io.Writer, sourceBucket string, paginator manifestPaginator) (manifestStats, error) {
	countingOutput := &countingWriter{Writer: output}
	writer := csv.NewWriter(countingOutput)
	stats := manifestStats{}
	for paginator.HasMorePages() {
		page, err := paginator.NextPage(ctx)
		if err != nil {
			return manifestStats{}, fmt.Errorf("list objects: %w", err)
		}
		for _, obj := range page.Contents {
			key := aws.ToString(obj.Key)
			if strings.HasSuffix(key, "/") {
				continue
			}
			if err := writer.Write([]string{sourceBucket, key}); err != nil {
				return manifestStats{}, fmt.Errorf("write csv record: %w", err)
			}
			stats.objects++
		}
	}
	writer.Flush()
	if err := writer.Error(); err != nil {
		return manifestStats{}, fmt.Errorf("flush manifest: %w", err)
	}
	stats.bytes = countingOutput.bytes
	return stats, nil
}

// UploadManifest uploads the generated manifest content to the specified S3 bucket and key, returning the ETag of the uploaded object.
func (c *BatchOpsClientImpl) UploadManifest(
	ctx context.Context, region, accountID, bucket, key, content string,
) (string, error) {
	logger := log.GetLogger(ctx).Named("aws-batch-ops")
	startAt := time.Now()
	logger.Info("upload manifest start",
		zap.String("region", region),
		zap.String("bucket", bucket),
		zap.String("key", key),
		zap.Int("manifest_bytes", len(content)))

	cfg, err := GenerateDPConfig(ctx, accountID, region)
	if err != nil {
		return "", fmt.Errorf("load aws config: %w", err)
	}
	s3Client := s3.NewFromConfig(cfg)

	_, err = s3Client.PutObject(ctx, &s3.PutObjectInput{
		Bucket:      aws.String(bucket),
		Key:         aws.String(key),
		Body:        strings.NewReader(content),
		ContentType: aws.String("text/csv"),
	})
	if err != nil {
		return "", fmt.Errorf("put manifest: %w", err)
	}

	head, err := s3Client.HeadObject(ctx, &s3.HeadObjectInput{
		Bucket: aws.String(bucket),
		Key:    aws.String(key),
	})
	if err != nil {
		return "", fmt.Errorf("head manifest: %w", err)
	}
	etag := strings.Trim(aws.ToString(head.ETag), "\"")
	if etag == "" {
		return "", fmt.Errorf("manifest etag is empty")
	}
	logger.Info("upload manifest finished",
		zap.String("bucket", bucket),
		zap.String("key", key),
		zap.String("etag", etag),
		zap.Duration("elapsed", time.Since(startAt)))
	return etag, nil
}

// CreateCopyJob creates an S3 Batch Operations copy job using the provided parameters and returns the job ID.
func (c *BatchOpsClientImpl) CreateCopyJob(
	ctx context.Context, region string, params *aws_facade.CreateCopyJobParams,
) (string, error) {
	logger := log.GetLogger(ctx).Named("aws-batch-ops")
	startAt := time.Now()
	logger.Info("create batch ops job start",
		zap.String("region", region),
		zap.String("account_id", params.AccountID),
		zap.String("manifest_bucket", params.ManifestBucket),
		zap.String("manifest_key", params.ManifestKey),
		zap.String("report_bucket", params.ReportBucket),
		zap.String("report_prefix", params.ReportPrefix),
		zap.String("target_bucket_arn", params.TargetBucketARN),
		zap.String("target_key_prefix", params.TargetKeyPrefix))

	cfg, err := GenerateDPConfig(ctx, params.AccountID, region)
	if err != nil {
		return "", fmt.Errorf("load aws config: %w", err)
	}
	s3ControlClient := s3control.NewFromConfig(cfg)

	copyOp := &s3ctypes.S3CopyObjectOperation{
		TargetResource:  aws.String(params.TargetBucketARN),
		TargetKeyPrefix: aws.String(normalizeS3KeyPrefix(params.TargetKeyPrefix)),
	}
	if params.CannedACL != "" {
		acl, err := parseCannedACL(params.CannedACL)
		if err != nil {
			return "", err
		}
		copyOp.CannedAccessControlList = acl
	}

	token := uuid.NewString()
	resp, err := s3ControlClient.CreateJob(ctx, &s3control.CreateJobInput{
		AccountId:          aws.String(params.AccountID),
		ClientRequestToken: aws.String(token),
		Operation: &s3ctypes.JobOperation{
			S3PutObjectCopy: copyOp,
		},
		Priority:             aws.Int32(batchOpsJobPriority),
		RoleArn:              aws.String(params.RoleARN),
		ConfirmationRequired: aws.Bool(false),
		Manifest: &s3ctypes.JobManifest{
			Spec: &s3ctypes.JobManifestSpec{
				Format: s3ctypes.JobManifestFormatS3BatchOperationsCsv20180820,
				Fields: []s3ctypes.JobManifestFieldName{
					s3ctypes.JobManifestFieldNameBucket,
					s3ctypes.JobManifestFieldNameKey,
				},
			},
			Location: &s3ctypes.JobManifestLocation{
				ObjectArn: aws.String(fmt.Sprintf("arn:aws:s3:::%s/%s", params.ManifestBucket, params.ManifestKey)),
				ETag:      aws.String(params.ManifestETag),
			},
		},
		Report: &s3ctypes.JobReport{
			Enabled:     true,
			Bucket:      aws.String(fmt.Sprintf("arn:aws:s3:::%s", params.ReportBucket)),
			Prefix:      aws.String(normalizeS3KeyPrefix(params.ReportPrefix)),
			Format:      s3ctypes.JobReportFormatReportCsv20180820,
			ReportScope: s3ctypes.JobReportScopeAllTasks,
		},
	})
	if err != nil {
		return "", fmt.Errorf("create job: %w", err)
	}
	jobID := aws.ToString(resp.JobId)
	logger.Info("create batch ops job finished",
		zap.String("job_id", jobID),
		zap.Duration("elapsed", time.Since(startAt)))
	return jobID, nil
}

// DescribeJob retrieves the status and progress of the specified S3 Batch Operations job.
func (c *BatchOpsClientImpl) DescribeJob(
	ctx context.Context, region, accountID, jobID string,
) (*aws_facade.JobStatus, error) {
	cfg, err := GenerateDPConfig(ctx, accountID, region)
	if err != nil {
		return nil, fmt.Errorf("load aws config: %w", err)
	}
	s3ControlClient := s3control.NewFromConfig(cfg)

	out, err := s3ControlClient.DescribeJob(ctx, &s3control.DescribeJobInput{
		AccountId: aws.String(accountID),
		JobId:     aws.String(jobID),
	})
	if err != nil {
		return nil, fmt.Errorf("describe job: %w", err)
	}
	if out.Job == nil {
		return nil, fmt.Errorf("describe job: empty job descriptor")
	}

	res := &aws_facade.JobStatus{
		Status: string(out.Job.Status),
	}
	if ps := out.Job.ProgressSummary; ps != nil {
		res.Total = aws.ToInt64(ps.TotalNumberOfTasks)
		res.Succeeded = aws.ToInt64(ps.NumberOfTasksSucceeded)
		res.Failed = aws.ToInt64(ps.NumberOfTasksFailed)
	}

	if out.Job.StatusUpdateReason != nil && *out.Job.StatusUpdateReason != "" {
		res.Message = *out.Job.StatusUpdateReason
	} else if len(out.Job.FailureReasons) > 0 {
		// keep it short for status_msg/error_msg
		res.Message = aws.ToString(out.Job.FailureReasons[0].FailureReason)
	}
	return res, nil
}

// CancelJob cancels the specified S3 Batch Operations job.
func (c *BatchOpsClientImpl) CancelJob(ctx context.Context, region, accountID, jobID string) error {
	cfg, err := GenerateDPConfig(ctx, accountID, region)
	if err != nil {
		return fmt.Errorf("load aws config: %w", err)
	}
	s3ControlClient := s3control.NewFromConfig(cfg)

	_, err = s3ControlClient.UpdateJobStatus(ctx, &s3control.UpdateJobStatusInput{
		AccountId:          aws.String(accountID),
		JobId:              aws.String(jobID),
		RequestedJobStatus: s3ctypes.RequestedJobStatusCancelled,
		StatusUpdateReason: aws.String("Canceled by user"),
	})
	if err != nil {
		if strings.Contains(strings.ToLower(err.Error()), aws_facade.JobStatusForbiddenMsg) {
			// Job status cannot be changed, e.g. already completed
			status := "unknown"
			if describe, dErr := c.DescribeJob(ctx, region, accountID, jobID); dErr == nil && describe != nil {
				status = describe.Status
			}
			return fmt.Errorf("cancel job failed, %s, current state=%s", aws_facade.JobStatusForbiddenMsg, status)
		}
		return fmt.Errorf("cancel job: %w", err)
	}
	return nil
}

// normalizeS3KeyPrefix ensures the S3 key prefix is in a standard format without leading and trailing slash
func normalizeS3KeyPrefix(prefix string) string {
	prefix = strings.TrimPrefix(prefix, "/")
	prefix = strings.TrimSuffix(prefix, "/")
	if prefix == "" {
		return ""
	}
	return prefix
}

// parseCannedACL converts a string representation of a canned ACL to the corresponding S3CannedAccessControlList type.
func parseCannedACL(v string) (s3ctypes.S3CannedAccessControlList, error) {
	switch strings.ToLower(strings.TrimSpace(v)) {
	case "bucketownerfullcontrol", "bucket-owner-full-control", "bucket_owner_full_control":
		return s3ctypes.S3CannedAccessControlListBucketOwnerFullControl, nil
	default:
		return "", fmt.Errorf("unsupported canned acl: %s", v)
	}
}

func (c *BatchOpsClientImpl) GenerateServiceRole(
	ctx context.Context, accountID string, projectId, clusterId uint64,
) string {
	return fmt.Sprintf("arn:aws:iam::%s:role/batchops-role-%d-%d", accountID, projectId, clusterId)
}