package v1beta1

import (
	"context"
	"fmt"
	"path/filepath"
	"strconv"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3/types"
	"github.com/tidbcloud/pkgv2/bizerr"
	"github.com/tidbcloud/pkgv2/log"
	"github.com/tidbcloud/pkgv2/tracing"
	"github.com/tidbcloud/serverless-service/idl/pbgen/proto/v1beta1"
	"github.com/tidbcloud/serverless-service/internal/domain/constant"
	"github.com/tidbcloud/serverless-service/internal/domain/entity"
	"github.com/tidbcloud/serverless-service/internal/errorcode"
	stashing "github.com/tidbcloud/serverless-service/internal/infra/stashing_area"
	tidbmgmtservicev1 "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
	"go.opentelemetry.io/otel/attribute"
	"go.uber.org/zap"
	"google.golang.org/protobuf/types/known/emptypb"
)

func (s *importService) StartUpload(
	ctx context.Context,
	req *v1beta1.StartUploadRequest,
) (_ *v1beta1.StartUploadResponse, err error) {
	ctx, span := Tracer.Start(ctx, "importService.StartUpload")
	defer func() {
		tracing.EndSpan(span, err)
		err = errorcode.PrepareImportErrorMessage(ctx, err)
	}()
	span.SetAttributes(
		attribute.String("cluster_id", req.ClusterId),
		attribute.Int("part_number", int(req.PartNumber)),
		attribute.String("file_name", req.FileName),
	)
	logger := log.WithContextL(ctx).
		With(zap.String("cluster_id", req.ClusterId), zap.Int32("part_number", req.PartNumber), zap.String("file_name", req.FileName))

	clusterID, err := strconv.ParseUint(req.ClusterId, 10, 64)
	if err != nil {
		return nil, bizerr.New(bizerr.ErrCodeParams, "invalid cluster id")
	}
	cluster, err := s.mgmtSvc.GetCluster(ctx, clusterID, tidbmgmtservicev1.ClusterFilter_CLUSTER_VIEW_LOCAL)
	if err != nil {
		logger.Error("failed to get cluster by id", zap.Error(err))
		return nil, err
	}
	stashingArea, err := s.stashingAreaProvider.GetStashingArea(cluster.CloudProvider, cluster.Region)
	if err != nil {
		logger.Error("failed to get stashing area", zap.Error(err))
		return nil, bizerr.New(bizerr.ErrCodeUnknown, "failed to get stashing area")
	}

	id := entity.GenerateIDWithPrefix(constant.UploadIDPrefix)
	key := fmt.Sprintf(
		"%s/%s",
		stashingArea.GetObjectPath(strconv.FormatUint(cluster.OrganizationId, 10), stashing.Import, id),
		filepath.Base(req.FileName),
	)

	// check storage limit
	exceeded, err := stashingArea.CheckExceedSize(ctx, strconv.FormatUint(cluster.OrganizationId, 10))
	if err != nil {
		logger.Error("failed to check storage limit", zap.Error(err))
		return nil, bizerr.New(bizerr.ErrCodeUnknown, "failed to check storage limit")
	}
	if exceeded {
		return nil, bizerr.New(errorcode.ErrCodeStorageLimitExceed, "")
	}

	if req.PartNumber < 1 {
		return nil, bizerr.New(bizerr.ErrCodeParams, "part number should be greater than 0")
	} else if req.PartNumber == 1 {
		request, err := stashingArea.PresignPutObject(ctx, key, 6*time.Hour)
		if err != nil {
			logger.Error("failed to presign put object", zap.Error(err))
			return nil, bizerr.New(bizerr.ErrCodeUnknown, "failed to presign put object")
		}
		bizErr := s.uploadRepo.Create(ctx, &entity.Upload{
			ID:            id,
			Status:        constant.UploadStatusCreated,
			Multipart:     false,
			AwsUploadID:   "",
			BucketKey:     key,
			TenantID:      cluster.OrganizationId,
			ClusterID:     clusterID,
			CloudProvider: cluster.GetCloudProvider(),
			RegionName:    cluster.Region,
		})
		if bizErr != nil {
			return nil, bizErr
		}
		return &v1beta1.StartUploadResponse{
			UploadUrl: []string{request.URL},
			UploadId:  id,
		}, nil
	} else if req.PartNumber > 10000 {
		return nil, bizerr.New(bizerr.ErrCodeParams, "part number should be less than 10000")
	}

	uploadId, urls, err := stashingArea.MultipartUpload(ctx, key, req.PartNumber, 30*time.Minute)
	if err != nil {
		logger.Error("failed to multipart upload", zap.Error(err))
		return nil, bizerr.New(bizerr.ErrCodeUnknown, "failed to multipart upload")
	}

	bizErr := s.uploadRepo.Create(ctx, &entity.Upload{
		ID:            id,
		Status:        constant.UploadStatusCreated,
		Multipart:     true,
		AwsUploadID:   *uploadId,
		BucketKey:     key,
		TenantID:      cluster.OrganizationId,
		ClusterID:     clusterID,
		CloudProvider: cluster.GetCloudProvider(),
		RegionName:    cluster.Region,
	})
	if bizErr != nil {
		return nil, bizErr
	}

	logger.Info("generate upload url successfully", zap.String("upload_id", *uploadId))
	return &v1beta1.StartUploadResponse{
		UploadUrl: urls,
		UploadId:  id,
	}, nil
}

func (s *importService) CompleteUpload(
	ctx context.Context,
	req *v1beta1.CompleteUploadRequest,
) (_ *emptypb.Empty, err error) {
	ctx, span := Tracer.Start(ctx, "importService.CompleteMultipartUpload")
	defer func() {
		tracing.EndSpan(span, err)
		err = errorcode.PrepareImportErrorMessage(ctx, err)
	}()
	span.SetAttributes(
		attribute.String("cluster_id", req.ClusterId),
		attribute.String("upload_id", req.UploadId),
	)
	logger := log.WithContextL(ctx).With(zap.String("upload_id", req.UploadId), zap.String("cluster_id", req.ClusterId))

	clusterID, err := strconv.ParseUint(req.ClusterId, 10, 64)
	if err != nil {
		return nil, bizerr.New(bizerr.ErrCodeParams, "invalid cluster id")
	}
	upload, bizErr := s.uploadRepo.GetUploadByID(ctx, req.UploadId)
	if bizErr != nil {
		return nil, bizErr
	}
	if clusterID != upload.ClusterID {
		logger.Error("inconsistent cluster id or file name", zap.Uint64("upload_cluster_id", upload.ClusterID))
		return nil, bizerr.New(bizerr.ErrCodeParams, "inconsistent cluster id or file name")
	}

	// complete multipart upload
	if upload.Multipart {
		stashingArea, err := s.stashingAreaProvider.GetStashingArea(upload.GetCloudProvider(), upload.RegionName)
		if err != nil {
			logger.Error("failed to get stashing area", zap.Error(err))
			return nil, bizerr.New(bizerr.ErrCodeUnknown, "failed to get stashing area")
		}

		parts := make([]types.CompletedPart, len(req.Parts))
		for i := range req.Parts {
			parts[i] = types.CompletedPart{
				ETag:       aws.String(req.Parts[i].Etag),
				PartNumber: aws.Int32(req.Parts[i].PartNumber),
			}
		}
		_, err = stashingArea.CompleteMultipartUpload(ctx, upload.BucketKey, upload.AwsUploadID, parts)
		if err != nil {
			logger.Error("failed to complete multipart upload", zap.Error(err))
			return nil, err
		}
	}

	bizErr = s.uploadRepo.UpdateStatus(ctx, req.UploadId, constant.UploadStatusCompleted)
	if bizErr != nil {
		return nil, bizErr
	}
	logger.Info("complete multipart upload successfully")
	return &emptypb.Empty{}, nil
}

func (s *importService) CancelUpload(
	ctx context.Context,
	req *v1beta1.CancelUploadRequest,
) (_ *emptypb.Empty, err error) {
	ctx, span := Tracer.Start(ctx, "importService.CancelMultipartUpload")
	defer func() {
		tracing.EndSpan(span, err)
		err = errorcode.PrepareImportErrorMessage(ctx, err)
	}()
	span.SetAttributes(
		attribute.String("cluster_id", req.ClusterId),
		attribute.String("upload_id", req.UploadId),
	)
	logger := log.WithContextL(ctx).With(zap.String("upload_id", req.UploadId), zap.String("cluster_id", req.ClusterId))

	clusterID, err := strconv.ParseUint(req.ClusterId, 10, 64)
	if err != nil {
		return nil, bizerr.New(bizerr.ErrCodeParams, "invalid cluster id")
	}
	upload, bizErr := s.uploadRepo.GetUploadByID(ctx, req.UploadId)
	if bizErr != nil {
		return nil, bizErr
	}
	if clusterID != upload.ClusterID {
		logger.Error("inconsistent cluster id", zap.Uint64("upload_cluster_id", upload.ClusterID))
		return nil, bizerr.New(bizerr.ErrCodeParams, "inconsistent cluster id")
	}

	if upload.Multipart {
		stashingArea, err := s.stashingAreaProvider.GetStashingArea(upload.GetCloudProvider(), upload.RegionName)
		if err != nil {
			logger.Error("failed to get stashing area", zap.Error(err))
			return nil, bizerr.New(bizerr.ErrCodeUnknown, "failed to get stashing area")
		}
		_, err = stashingArea.AbortMultipartUpload(ctx, upload.BucketKey, upload.AwsUploadID)
		if err != nil {
			logger.Error("failed to abort multipart upload", zap.Error(err))
			return nil, bizerr.New(bizerr.ErrCodeUnknown, "failed to cancel upload")
		}
	}

	bizErr = s.uploadRepo.UpdateStatus(ctx, req.UploadId, constant.UploadStatusCanceled)
	if bizErr != nil {
		return nil, bizErr
	}
	logger.Info("abort multipart upload successfully")
	return &emptypb.Empty{}, nil
}