package v1beta1

import (
	"context"
	"fmt"
	"strconv"
	"strings"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/aws/aws-sdk-go-v2/service/s3"
	"github.com/tidbcloud/pkgv2/bizerr"
	"github.com/tidbcloud/pkgv2/import/cloud_storage"
	"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/config"
	"github.com/tidbcloud/serverless-service/internal/domain/constant"
	"github.com/tidbcloud/serverless-service/internal/domain/entity"
	"github.com/tidbcloud/serverless-service/internal/domain/facade"
	"github.com/tidbcloud/serverless-service/internal/domain/repo"
	"github.com/tidbcloud/serverless-service/internal/domain/service"
	"github.com/tidbcloud/serverless-service/internal/domain/vo"
	"github.com/tidbcloud/serverless-service/internal/errorcode"
	stashing "github.com/tidbcloud/serverless-service/internal/infra/stashing_area"
	"github.com/tidbcloud/serverless-service/internal/service/import_svc"
	"github.com/tidbcloud/serverless-service/internal/utils"
	globalserverlessv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/serverless/v1"
	tidbmgmtservicev1 "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/trace"
	"go.uber.org/zap"
	"google.golang.org/protobuf/types/known/emptypb"
)

var Tracer = otel.Tracer("import.v1beta1")

var (
	ListSupportedOrderByField = []string{"create_time"}
	ListOrderByFieldMap       = map[string]string{
		"create_time": "created_at",
	}
)

type importService struct {
	uploadRepo           repo.UploadRepo
	stashingAreaProvider *stashing.StashingAreaProvider
	importHelper         *import_svc.ImportHelperService
	importSvc            service.ImportService
	mgmtSvc              facade.MgmtSvcFacade
	v1beta1.UnimplementedImportServiceServer
}

func NewImportService(
	uploadRepo repo.UploadRepo,
	stashingAreaProvider *stashing.StashingAreaProvider,
	importHelper *import_svc.ImportHelperService,
	importSvc service.ImportService,
	mgmtSvc facade.MgmtSvcFacade,
) v1beta1.ImportServiceServer {
	return &importService{
		uploadRepo:           uploadRepo,
		stashingAreaProvider: stashingAreaProvider,
		importSvc:            importSvc,
		importHelper:         importHelper,
		mgmtSvc:              mgmtSvc,
	}
}

func (s *importService) ListImports(
	ctx context.Context,
	req *v1beta1.ListImportsRequest,
) (_ *v1beta1.ListImportsResponse, err error) {
	ctx, span := Tracer.Start(ctx, "v1beta1.import.List")
	defer func() {
		tracing.EndSpan(span, err)
		err = errorcode.PrepareImportErrorMessage(ctx, err)
	}()

	pager := utils.NewAipPager()
	err = pager.Normalize(req)
	if err != nil {
		return nil, err
	}

	// check orderBy
	err = utils.CheckOrderBy(req, ListSupportedOrderByField)
	if err != nil {
		return nil, bizerr.New(
			errorcode.ErrCodeInvalidOrderby,
			"%s",
			errorcode.ExportErrorMessages[errorcode.ErrCodeInvalidOrderby]+" Error: "+err.Error(),
		)
	}
	orderBy := req.OrderBy
	for k, v := range ListOrderByFieldMap {
		orderBy = strings.ReplaceAll(orderBy, k, v)
	}

	span.AddEvent("call ListImports")
	total, imports, bizErr := s.importSvc.ListImports(ctx, &vo.ListImportOpt{
		ClusterID: req.ClusterId,
		Offset:    int64(pager.Offset()),
		PageSize:  int32(pager.PageSize()),
		OrderBy:   orderBy,
	})
	if bizErr != nil {
		return nil, bizErr
	}
	nextPageToken := pager.NextPageToken(uint32(total))

	return vo.ListImport2oa(imports, total, nextPageToken)
}

func (s *importService) CreateImport(
	ctx context.Context,
	req *v1beta1.CreateImportRequest,
) (_ *v1beta1.Import, err error) {
	ctx, span := Tracer.Start(ctx, "v1beta1.import.Create")
	defer func() {
		tracing.EndSpan(span, err)
		err = errorcode.PrepareImportErrorMessage(ctx, err)
	}()

	bizErr := validateImportReq(ctx, req)
	if bizErr != nil {
		return nil, bizErr
	}

	switch req.GetSource().GetType() {
	case v1beta1.ImportSourceType_LOCAL:
		return s.createImportFromLocalFile(ctx, req, span)
	case v1beta1.ImportSourceType_S3, v1beta1.ImportSourceType_S3_COMPATIBLE, v1beta1.ImportSourceType_OSS:
		return s.createImportFromS3(ctx, req, span)
	case v1beta1.ImportSourceType_GCS:
		return s.createImportFromGCS(ctx, req, span)
	case v1beta1.ImportSourceType_AZURE_BLOB:
		return s.createImportFromAzBlob(ctx, req, span)
	default:
		return nil, bizerr.New(bizerr.ErrCodeParams, "unrecognized import type")
	}
}

func (s *importService) createImportFromS3(
	ctx context.Context,
	req *v1beta1.CreateImportRequest,
	span trace.Span,
) (*v1beta1.Import, error) {
	logger := log.WithContextL(ctx).With(zap.String("cluster_id", req.ClusterId))

	cluster, bizErr := s.validateCluster(ctx, req, logger)
	if bizErr != nil {
		return nil, bizErr
	}
	if cluster.CloudProvider != config.ProviderAWS &&
		req.GetSource().GetS3().GetAuthType() == v1beta1.ImportS3AuthType_ROLE_ARN {
		return nil, bizerr.New(bizerr.ErrCodeParams, "RoleArn is only supported for AWS Clusters")
	}

	region, bizErr := s.getS3bucketRegion(ctx, req.GetSource(), cluster)
	if bizErr != nil {
		return nil, bizErr
	}
	imp, bizErr := vo.BuildV1beta1ImportWithS3Uri(ctx, req, cluster, region)
	if bizErr != nil {
		logger.Error("build import error", zap.Error(bizErr))
		return nil, bizErr
	}

	span.AddEvent("call CreateImport")
	resp, err := s.create(ctx, imp, nil)
	if err != nil {
		logger.Error("CreateImport failed", zap.Error(err))
		return nil, err
	}

	r, err := vo.GetImport2oa(resp)
	if err != nil {
		logger.Error("GetImport2oa failed", zap.Error(err))
		return nil, err
	}
	return r, nil
}

func (s *importService) createImportFromGCS(
	ctx context.Context,
	req *v1beta1.CreateImportRequest,
	span trace.Span,
) (*v1beta1.Import, error) {
	logger := log.WithContextL(ctx).With(zap.String("cluster_id", req.ClusterId))

	cluster, bizErr := s.validateCluster(ctx, req, logger)
	if bizErr != nil {
		return nil, bizErr
	}

	imp, bizErr := vo.BuildGCSImport(ctx, req, cluster)
	if bizErr != nil {
		logger.Error("build import error", zap.Error(bizErr))
		return nil, bizErr
	}

	span.AddEvent("call CreateImport")
	resp, err := s.create(ctx, imp, nil)
	if err != nil {
		logger.Error("CreateImport failed", zap.Error(err))
		return nil, err
	}

	r, err := vo.GetImport2oa(resp)
	if err != nil {
		logger.Error("GetImport2oa failed", zap.Error(err))
		return nil, err
	}
	return r, nil
}

func (s *importService) createImportFromAzBlob(
	ctx context.Context,
	req *v1beta1.CreateImportRequest,
	span trace.Span,
) (*v1beta1.Import, error) {
	logger := log.WithContextL(ctx).With(zap.String("cluster_id", req.ClusterId))

	cluster, bizErr := s.validateCluster(ctx, req, logger)
	if bizErr != nil {
		return nil, bizErr
	}

	imp, bizErr := vo.BuildAzureImport(ctx, req, cluster)
	if bizErr != nil {
		logger.Error("build import error", zap.Error(bizErr))
		return nil, bizErr
	}

	span.AddEvent("call CreateImport")
	resp, err := s.create(ctx, imp, nil)
	if err != nil {
		logger.Error("CreateImport failed", zap.Error(err))
		return nil, err
	}

	r, err := vo.GetImport2oa(resp)
	if err != nil {
		logger.Error("GetImport2oa failed", zap.Error(err))
		return nil, err
	}
	return r, nil
}

func (s *importService) createImportFromLocalFile(
	ctx context.Context,
	req *v1beta1.CreateImportRequest,
	span trace.Span,
) (_ *v1beta1.Import, err error) {
	logger := log.WithContextL(ctx).With(zap.String("cluster_id", req.ClusterId))

	cluster, bizErr := s.validateCluster(ctx, req, logger)
	if bizErr != nil {
		return nil, bizErr
	}

	span.AddEvent("call GetUploadByID")
	upload, bizErr := s.uploadRepo.GetUploadByID(ctx, req.GetSource().GetLocal().GetUploadId())
	if bizErr != nil {
		return nil, bizErr
	}
	bizErr = import_svc.ValidateUpload(ctx, upload, cluster.ClusterId)
	if bizErr != nil {
		return nil, bizErr
	}

	span.AddEvent("get import file from stashing area")
	area, err := s.stashingAreaProvider.GetStashingArea(cluster.CloudProvider, cluster.Region)
	if err != nil {
		logger.Error("GetStashingArea failed", zap.Error(err))
		return nil, bizerr.New(errorcode.ErrCodeInternalError, errorcode.InternalErrorMessage)
	}

	object, err := area.GetObject(ctx, &s3.GetObjectInput{
		Bucket: aws.String(area.GetBucketName()),
		Key:    aws.String(upload.BucketKey),
		Range:  aws.String(fmt.Sprintf("bytes=0-%d", constant.PreviewReadSize)),
	})
	if err != nil {
		logger.Error(
			"GetObject failed",
			zap.Error(err),
			zap.String("bucket", area.GetBucketName()),
			zap.String("key", upload.BucketKey),
		)
		return nil, bizerr.New(errorcode.ErrCodeInternalError, errorcode.InternalErrorMessage)
	}
	defer object.Body.Close()

	dbPreviewResult, err := import_svc.PreviewLocalImport(
		ctx,
		object.Body,
		vo.Pb2CsvConfig(req.GetImportOptions().GetCsvFormat()),
		vo.PbFileType2Entity(req.GetImportOptions().GetFileType()),
	)
	if err != nil {
		logger.Error("PreviewLocalImport failed", zap.Error(err))
		return nil, err
	}
	if len(dbPreviewResult.GetTables()) < 1 {
		logger.Error("no preview result for the table")
		return nil, bizerr.New(
			bizerr.ErrCodeParams,
			"no preview result for the table, please make sure the CSV file is correct",
		)
	}
	fileColumns := import_svc.ParseColumns(dbPreviewResult.GetTables()[0])

	span.AddEvent("call GetTableColumns")
	tableColumns, bizErr := s.importHelper.GetTableColumns(
		ctx,
		cluster,
		req.GetSource().GetLocal().GetTargetDatabase(),
		req.GetSource().GetLocal().GetTargetTable(),
	)
	var targetColumns []*entity.ColumnInfo
	if bizErr != nil {
		if bizErr.RawErrCode() != errorcode.ErrCodeTableNotExist {
			logger.Error("GetTableColumns failed", zap.Error(bizErr))
			return nil, bizErr
		}
	} else {
		targetColumns = import_svc.ParseColumns(tableColumns)
	}

	// if the target table exists, validate the columns
	span.AddEvent("validate table columns")
	if err = import_svc.ValidateColumns(fileColumns, targetColumns); err != nil {
		logger.Error("ValidateColumns failed", zap.Error(err))
		return nil, err
	}

	var columns []*entity.ColumnInfo
	// if targetTable does not exist, add fileColumns in order to create target table later
	if len(targetColumns) == 0 {
		columns = fileColumns
	}

	imp, err := vo.BuildV1beta1ImportWithS3UriLocal(ctx, req, cluster, upload, area)
	if err != nil {
		logger.Error("build import error", zap.Error(err))
		return nil, err
	}

	span.AddEvent("call CreateImport")
	resp, err := s.create(ctx, imp, columns)
	if err != nil {
		logger.Error("CreateImport failed", zap.Error(err))
		return nil, err
	}

	r, err := vo.GetImport2oa(resp)
	if err != nil {
		logger.Error("GetImport2oa failed", zap.Error(err))
		return nil, err
	}
	return r, nil
}

func (s *importService) create(
	ctx context.Context,
	imp *entity.ImportTask,
	columns []*entity.ColumnInfo,
) (_ *entity.Import, err error) {
	logger := log.WithContextL(ctx).
		With(zap.String("project_id", imp.ProjectID)).
		With(zap.String("cluster_id", imp.ClusterID))

	if imp.LocalImportCfg != nil {
		err = s.importHelper.CreateTargetTable(ctx, imp.Cluster, imp.LocalImportCfg, columns, nil)
		if err != nil {
			logger.Error("CreateTargetTable failed", zap.Error(err))
			return nil, err
		}
	}

	id, bizErr := s.importSvc.CreateImport(ctx, imp)
	if bizErr != nil {
		return nil, bizErr
	}

	logger.Info("create import success", zap.String("import_id", id))
	im, bizErr := s.importSvc.GetImport(ctx, imp.ClusterID, id)
	if bizErr != nil {
		return nil, bizErr
	}

	return im, nil
}

func (s *importService) GetImport(ctx context.Context, req *v1beta1.GetImportRequest) (_ *v1beta1.Import, err error) {
	ctx, span := Tracer.Start(ctx, "v1beta1.import.Get")
	defer func() {
		tracing.EndSpan(span, err)
		err = errorcode.PrepareImportErrorMessage(ctx, err)
	}()

	span.AddEvent("call GetImport")
	i, err := s.importSvc.GetImport(ctx, req.ClusterId, req.Id)
	if err != nil {
		log.WithContextL(ctx).Error("GetImport failed",
			zap.String("clusterId", req.ClusterId), zap.Error(err))
		return nil, err
	}

	return vo.GetImport2oa(i)
}

func (s *importService) CancelImport(
	ctx context.Context,
	req *v1beta1.CancelImportRequest,
) (_ *emptypb.Empty, err error) {
	ctx, span := Tracer.Start(ctx, "v1beta1.import.Cancel")
	defer func() {
		tracing.EndSpan(span, err)
		err = errorcode.PrepareImportErrorMessage(ctx, err)
	}()

	span.AddEvent("call CancelImport")
	bizErr := s.importSvc.CancelImport(ctx, req.ClusterId, req.Id)
	if bizErr != nil {
		return nil, bizErr
	}

	return &emptypb.Empty{}, nil
}

func (s *importService) validateCluster(
	ctx context.Context,
	req *v1beta1.CreateImportRequest,
	logger *zap.Logger,
) (*tidbmgmtservicev1.Cluster, bizerr.BizError) {
	cID, err := strconv.ParseUint(req.ClusterId, 10, 64)
	if err != nil {
		return nil, bizerr.New(bizerr.ErrCodeParams, "Invalid cluster ID: %s", req.ClusterId)
	}
	cluster, bizErr := s.mgmtSvc.GetCluster(ctx, cID, tidbmgmtservicev1.ClusterFilter_CLUSTER_VIEW_BASIC)
	if bizErr != nil {
		logger.Error("GetClusterByID failed", zap.String("cluster_id", req.ClusterId), zap.Error(bizErr))
		return nil, bizErr
	}

	if cluster.GetServerlessCluster().
		GetUserStatus() !=
		globalserverlessv1.UserClusterStatus_USER_CLUSTER_STATUS_NORMAL {
		return nil, bizerr.New(
			errorcode.ErrCodeClusterNotRunning,
			"Cluster ID %s is not in normal state, please try again later.",
			req.ClusterId,
		)
	}
	return cluster, nil
}

func (s *importService) getS3bucketRegion(
	ctx context.Context,
	source *v1beta1.ImportSource,
	cluster *tidbmgmtservicev1.Cluster,
) (string, bizerr.BizError) {
	var cs cloud_storage.CloudStorage
	var bizErr bizerr.BizError
	switch source.GetType() {
	case v1beta1.ImportSourceType_S3_COMPATIBLE:
		accessKey := source.GetS3Compatible().GetAccessKey()
		cs, bizErr = cloud_storage.NewS3Compatible(ctx, source.GetS3Compatible().GetUri(),
			cloud_storage.WithS3CompatibleAccessKey(accessKey.GetId(), accessKey.GetSecret()),
			cloud_storage.WithS3CompatibleEndpoint(source.GetS3Compatible().GetEndpoint()),
		)
		if bizErr != nil {
			return "", bizErr
		}
		return cs.(*cloud_storage.S3).Region, nil
	case v1beta1.ImportSourceType_OSS:
		oss := source.GetOss()
		options := []cloud_storage.OSSOption{cloud_storage.WithOSSAccessKey(
			oss.GetAccessKey().GetId(), oss.GetAccessKey().GetSecret())}
		var region string
		region, bizErr = cloud_storage.GetOssBucketRegion(ctx, oss.GetUri(), options...)
		if bizErr != nil {
			return "", bizErr
		}
		return region, nil
	case v1beta1.ImportSourceType_S3:
		switch source.GetS3().GetAuthType() {
		case v1beta1.ImportS3AuthType_ROLE_ARN:
			var profile *entity.DataIngestionProfile
			profile, bizErr = s.importSvc.GetDataIngestionProfile(ctx, cluster)
			if bizErr != nil {
				return "", bizErr
			}
			cs, bizErr = cloud_storage.NewS3(ctx, source.GetS3().GetUri(),
				cloud_storage.WithAssumeRole(source.GetS3().GetRoleArn(), profile.ExternalId),
			)
			if bizErr != nil {
				return "", bizErr
			}
		case v1beta1.ImportS3AuthType_ACCESS_KEY:
			accessKey := source.GetS3().GetAccessKey()
			cs, bizErr = cloud_storage.NewS3(ctx, source.GetS3().GetUri(),
				cloud_storage.WithAccessKey(accessKey.GetId(), accessKey.GetSecret()),
			)
			if bizErr != nil {
				return "", bizErr
			}
		}
		return cs.(*cloud_storage.S3).Region, nil
	}

	return "", bizerr.New(bizerr.ErrCodeParams, "unsupported import source type for S3 bucket: %s", source.GetType())
}