package vo

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"path/filepath"
	"strconv"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/juju/errors"
	"github.com/pingcap/tidb/pkg/lightning/config"
	"github.com/tidbcloud/pkgv2/bizerr"
	"github.com/tidbcloud/pkgv2/import/cloud_storage"
	"github.com/tidbcloud/pkgv2/log"
	"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"
	"github.com/tidbcloud/serverless-service/internal/utils"
	tidbmgmtservicev1 "github.com/tidbcloud/tidb-management-service/api/spec/tidb_mgmt_service/v1"
	"go.uber.org/zap"
	"google.golang.org/protobuf/types/known/timestamppb"
)

func ListImport2oa(imports []*entity.Import, total int64, token string) (*v1beta1.ListImportsResponse, error) {
	v := new(v1beta1.ListImportsResponse)
	v.TotalSize = uint32(total)
	v.NextPageToken = token

	for _, i := range imports {
		df2oa, err := GetImport2oa(i)
		if err != nil {
			return nil, err
		}
		v.Imports = append(v.Imports, df2oa)
	}
	return v, nil
}

func GetImport2oa(i *entity.Import) (*v1beta1.Import, error) {
	if i == nil {
		return nil, bizerr.New(bizerr.ErrCodeUnknown, "import is nil")
	}
	req, err := ImportEntity2CreationDetails(i)
	if err != nil {
		return nil, err
	}

	resp := &v1beta1.Import{
		Id:              i.ID,
		Name:            fmt.Sprintf("clusters/%s/imports/%s", i.ClusterID, i.ID),
		ClusterId:       i.ClusterID,
		CreateTime:      timestamppb.New(i.CreatedAt),
		State:           entityImportStatus2PBState(i.Status),
		CompletePercent: i.CompletedPercent,
		Message:         PrettyMessage(i.Message, i.DataFormat),
		CreatedBy:       i.Creator,
		CreationDetails: req,
		ImportId:        i.ID,
	}

	if i.TotalSize != nil {
		resp.TotalSize = *i.TotalSize
	}

	if i.CompletedAt != nil {
		resp.CompleteTime = timestamppb.New(*i.CompletedAt)
	}

	return resp, nil
}

func BuildV1beta1ImportWithS3UriLocal(
	ctx context.Context,
	req *v1beta1.CreateImportRequest,
	cluster *tidbmgmtservicev1.Cluster,
	upload *entity.Upload,
	area *stashing.StashingArea,
) (*entity.ImportTask, bizerr.BizError) {
	imp := &entity.ImportTask{
		Import: &entity.Import{
			ID:             entity.GenerateIDWithPrefix(constant.ImportIDPrefix),
			ProjectID:      strconv.FormatUint(cluster.GetOriginalProjectId(), 10),
			TenantID:       strconv.FormatUint(cluster.OrganizationId, 10),
			ClusterID:      req.ClusterId,
			SourceURL:      fmt.Sprintf("s3://%s/%s", area.GetBucketName(), upload.BucketKey),
			BucketRegion:   area.Region,
			DataSourceType: pbDataSource2Entity(req.GetSource().GetType()),
			DataFormat:     PbFileType2Entity(req.GetImportOptions().GetFileType()),
			NeedDropTables: false,
		},
		Cluster: cluster,
	}

	if req.GetImportOptions().GetFileType() == v1beta1.ImportFileType_CSV {
		csvBytes, err := json.Marshal(req.GetImportOptions().GetCsvFormat())
		if err != nil {
			log.Error(ctx, "failed to marshal CSV format", zap.Error(err))
			return nil, bizerr.New(errorcode.ErrCodeInternalError, errorcode.InternalErrorMessage)
		}
		imp.CSVFormat = string(csvBytes)
	}

	if req.GetSource().GetLocal() != nil {
		imp.LocalImportCfg = &entity.LocalImportCfg{
			FileName:    filepath.Base(upload.BucketKey),
			TargetDB:    req.GetSource().GetLocal().GetTargetDatabase(),
			TargetTable: req.GetSource().GetLocal().GetTargetTable(),
			UploadID:    upload.ID,
		}
	}

	if area.Endpoint != "" {
		imp.S3Endpoint = area.Endpoint
		imp.S3AccessKeyID = area.AccessKeyID
		imp.S3SecretAccessKey = area.SecretAccessKey
	}

	imp.Creator = utils.BuildCreatorID(ctx)

	return imp, nil
}

func BuildV1beta1ImportWithS3Uri(
	ctx context.Context,
	req *v1beta1.CreateImportRequest,
	cluster *tidbmgmtservicev1.Cluster,
	region string,
) (*entity.ImportTask, bizerr.BizError) {
	imp := &entity.ImportTask{
		Import: &entity.Import{
			ID:             entity.GenerateIDWithPrefix(constant.ImportIDPrefix),
			ProjectID:      strconv.FormatUint(cluster.GetOriginalProjectId(), 10),
			TenantID:       strconv.FormatUint(cluster.OrganizationId, 10),
			ClusterID:      req.ClusterId,
			BucketRegion:   region,
			DataSourceType: pbDataSource2Entity(req.GetSource().GetType()),
			DataFormat:     PbFileType2Entity(req.GetImportOptions().GetFileType()),
			NeedDropTables: false,
		},
		Cluster: cluster,
	}

	if req.GetImportOptions().GetFileType() == v1beta1.ImportFileType_CSV {
		csvBytes, err := json.Marshal(req.GetImportOptions().GetCsvFormat())
		if err != nil {
			log.Error(ctx, "failed to marshal CSV format", zap.Error(err))
			return nil, bizerr.New(errorcode.ErrCodeInternalError, errorcode.InternalErrorMessage)
		}
		imp.CSVFormat = string(csvBytes)
	}

	if req.GetSource().GetS3() != nil {
		imp.SourceURL = req.GetSource().GetS3().GetUri()
		switch req.GetSource().GetS3().GetAuthType() {
		case v1beta1.ImportS3AuthType_ROLE_ARN:
			imp.S3RoleArn = req.GetSource().GetS3().GetRoleArn()
		case v1beta1.ImportS3AuthType_ACCESS_KEY:
			imp.S3AccessKeyID = req.GetSource().GetS3().GetAccessKey().Id
			imp.S3SecretAccessKey = req.GetSource().GetS3().GetAccessKey().Secret
		}
	}

	if req.GetSource().GetS3Compatible() != nil {
		imp.SourceURL = req.GetSource().GetS3Compatible().GetUri()
		switch req.GetSource().GetS3Compatible().GetAuthType() {
		case v1beta1.ImportS3CompatibleAuthType_ACCESS_KEY:
			imp.S3AccessKeyID = req.GetSource().GetS3Compatible().GetAccessKey().Id
			imp.S3SecretAccessKey = req.GetSource().GetS3Compatible().GetAccessKey().Secret
		}
		if req.GetSource().GetS3Compatible().GetEndpoint() != "" {
			imp.S3Endpoint = req.GetSource().GetS3Compatible().GetEndpoint()
		}
	}

	if req.GetSource().GetOss() != nil {
		imp.SourceURL = req.GetSource().GetOss().GetUri()
		imp.S3Endpoint = cloud_storage.GenerateOssEndpoint(region)
		switch req.GetSource().GetOss().GetAuthType() {
		case v1beta1.ImportOSSAuthType_ACCESS_KEY:
			imp.S3AccessKeyID = req.GetSource().GetOss().GetAccessKey().Id
			imp.S3SecretAccessKey = req.GetSource().GetOss().GetAccessKey().Secret
		}
	}

	imp.Creator = utils.BuildCreatorID(ctx)

	return imp, nil
}

func PbFileType2Entity(fileType v1beta1.ImportFileType_Enum) constant.DataFormat {
	switch fileType {
	case v1beta1.ImportFileType_CSV:
		return constant.ImportDataFormatCsv
	case v1beta1.ImportFileType_SQL:
		return constant.ImportDataFormatSql
	case v1beta1.ImportFileType_AURORA_SNAPSHOT:
		return constant.ImportDataFormatAuroraSnapshot
	case v1beta1.ImportFileType_PARQUET:
		return constant.ImportDataFormatParquet
	}
	return constant.ImportDataFormatUnknown
}

func pbDataSource2Entity(dataSource v1beta1.ImportSourceType_Enum) constant.DataSource {
	switch dataSource {
	case v1beta1.ImportSourceType_LOCAL:
		return constant.ImportDataSourceLocal
	case v1beta1.ImportSourceType_S3:
		return constant.ImportDataSourceS3
	case v1beta1.ImportSourceType_GCS:
		return constant.ImportDataSourceGCS
	case v1beta1.ImportSourceType_AZURE_BLOB:
		return constant.ImportDataSourceAzBlob
	case v1beta1.ImportSourceType_S3_COMPATIBLE:
		return constant.ImportDataSourceS3Compatible
	case v1beta1.ImportSourceType_OSS:
		return constant.ImportDataSourceOSS
	}
	return constant.ImportDataSourceUnknown
}

func entityImportStatus2PBState(entityResult constant.ImportStatus) v1beta1.ImportState_Enum {
	switch entityResult {
	case constant.ImportStatusUnknown:
		return v1beta1.ImportState_UNSPECIFIED
	case constant.ImportStatusPreparing:
		return v1beta1.ImportState_PREPARING
	case constant.ImportStatusImporting:
		return v1beta1.ImportState_IMPORTING
	case constant.ImportStatusCompleted:
		return v1beta1.ImportState_COMPLETED
	case constant.ImportStatusFailed:
		return v1beta1.ImportState_FAILED
	case constant.ImportStatusCanceling:
		return v1beta1.ImportState_CANCELING
	case constant.ImportStatusCanceled:
		return v1beta1.ImportState_CANCELED
	default:
		return v1beta1.ImportState_UNSPECIFIED
	}
}

func ImportEntity2CreationDetails(imp *entity.Import) (*v1beta1.CreationDetails, error) {
	var csvFormat v1beta1.CSVFormat
	if imp.DataFormat == constant.ImportDataFormatCsv && imp.CSVFormat != "" {
		if err := json.Unmarshal([]byte(imp.CSVFormat), &csvFormat); err != nil {
			return nil, errors.Annotate(err, "parse custom CSV format failed")
		}
	}
	resp := &v1beta1.CreationDetails{
		ImportOptions: &v1beta1.ImportOptions{
			FileType: dataFormat2FileType(imp.DataFormat),
			Format: &v1beta1.ImportOptions_CsvFormat{
				CsvFormat: &csvFormat,
			},
		},
	}
	if imp.LocalImportCfg != nil {
		resp.Source = &v1beta1.ImportSource{
			Type: v1beta1.ImportSourceType_LOCAL,
			Source: &v1beta1.ImportSource_Local{
				Local: &v1beta1.LocalSource{
					UploadId:       imp.LocalImportCfg.UploadID,
					TargetTable:    imp.LocalImportCfg.TargetTable,
					TargetDatabase: imp.LocalImportCfg.TargetDB,
					FileName:       imp.LocalImportCfg.FileName,
				},
			},
		}
	}
	switch imp.DataSourceType {
	case constant.ImportDataSourceS3:
		resp.Source = &v1beta1.ImportSource{
			Type: v1beta1.ImportSourceType_S3,
			Source: &v1beta1.ImportSource_S3{
				S3: &v1beta1.S3Source{
					Uri: imp.SourceURL,
				},
			},
		}
		if imp.S3RoleArn != "" {
			resp.GetSource().GetS3().AuthType = v1beta1.ImportS3AuthType_ROLE_ARN
			resp.Source.GetS3().Auth = &v1beta1.S3Source_RoleArn{
				RoleArn: imp.S3RoleArn,
			}
		} else if imp.S3AccessKeyID != "" {
			resp.GetSource().GetS3().AuthType = v1beta1.ImportS3AuthType_ACCESS_KEY
			resp.Source.GetS3().Auth = &v1beta1.S3Source_AccessKey_{
				AccessKey: &v1beta1.S3Source_AccessKey{
					Id: imp.S3AccessKeyID,
				},
			}
		}
	case constant.ImportDataSourceGCS:
		resp.Source = &v1beta1.ImportSource{
			Type: v1beta1.ImportSourceType_GCS,
			Source: &v1beta1.ImportSource_Gcs{
				Gcs: &v1beta1.GCSSource{
					Uri:      imp.SourceURL,
					AuthType: v1beta1.ImportGcsAuthType_SERVICE_ACCOUNT_KEY,
				},
			},
		}
	case constant.ImportDataSourceAzBlob:
		resp.Source = &v1beta1.ImportSource{
			Type: v1beta1.ImportSourceType_AZURE_BLOB,
			Source: &v1beta1.ImportSource_AzureBlob{
				AzureBlob: &v1beta1.AzureBlobSource{
					Uri:      imp.SourceURL,
					AuthType: v1beta1.ImportAzureBlobAuthType_SAS_TOKEN,
				},
			},
		}
	case constant.ImportDataSourceS3Compatible:
		resp.Source = &v1beta1.ImportSource{
			Type: v1beta1.ImportSourceType_S3_COMPATIBLE,
			Source: &v1beta1.ImportSource_S3Compatible{
				S3Compatible: &v1beta1.S3CompatibleSource{
					Uri: imp.SourceURL,
				},
			},
		}
		if imp.S3AccessKeyID != "" {
			resp.GetSource().GetS3Compatible().AuthType = v1beta1.ImportS3CompatibleAuthType_ACCESS_KEY
			resp.Source.GetS3Compatible().Auth = &v1beta1.S3CompatibleSource_AccessKey_{
				AccessKey: &v1beta1.S3CompatibleSource_AccessKey{
					Id: imp.S3AccessKeyID,
				},
			}
		}
		if imp.S3Endpoint != "" {
			resp.GetSource().GetS3Compatible().Endpoint = aws.String(imp.S3Endpoint)
		}
	case constant.ImportDataSourceOSS:
		resp.Source = &v1beta1.ImportSource{
			Type: v1beta1.ImportSourceType_OSS,
			Source: &v1beta1.ImportSource_Oss{
				Oss: &v1beta1.OSSSource{
					Uri: imp.SourceURL,
				},
			},
		}
		if imp.S3AccessKeyID != "" {
			resp.GetSource().GetOss().AuthType = v1beta1.ImportOSSAuthType_ACCESS_KEY
			resp.Source.GetOss().Auth = &v1beta1.OSSSource_AccessKey_{
				AccessKey: &v1beta1.OSSSource_AccessKey{
					Id: imp.S3AccessKeyID,
				},
			}
		}
	}

	return resp, nil
}

func dataFormat2FileType(format constant.DataFormat) v1beta1.ImportFileType_Enum {
	switch format {
	case constant.ImportDataFormatCsv:
		return v1beta1.ImportFileType_CSV
	case constant.ImportDataFormatSql:
		return v1beta1.ImportFileType_SQL
	case constant.ImportDataFormatAuroraSnapshot:
		return v1beta1.ImportFileType_AURORA_SNAPSHOT
	case constant.ImportDataFormatParquet:
		return v1beta1.ImportFileType_PARQUET
	default:
		return v1beta1.ImportFileType_UNSPECIFIED
	}
}

type ListImportOpt struct {
	ClusterID string
	PageSize  int32
	Offset    int64
	Statuses  []constant.ImportStatus
	OrderBy   string
}

func BuildAzureImport(
	ctx context.Context,
	req *v1beta1.CreateImportRequest,
	cluster *tidbmgmtservicev1.Cluster,
) (*entity.ImportTask, bizerr.BizError) {
	_, _, bizErr := utils.ParseAzureURI(req.GetSource().GetAzureBlob().GetUri())
	if bizErr != nil {
		return nil, bizErr
	}

	imp := &entity.ImportTask{
		Import: &entity.Import{
			ID:             entity.GenerateIDWithPrefix(constant.ImportIDPrefix),
			ProjectID:      strconv.FormatUint(cluster.GetOriginalProjectId(), 10),
			TenantID:       strconv.FormatUint(cluster.OrganizationId, 10),
			ClusterID:      req.ClusterId,
			SourceURL:      req.GetSource().GetAzureBlob().GetUri(),
			DataSourceType: pbDataSource2Entity(req.GetSource().GetType()),
			DataFormat:     PbFileType2Entity(req.GetImportOptions().GetFileType()),
			NeedDropTables: false,
			AzureSASToken:  req.GetSource().GetAzureBlob().GetSasToken(),
		},
		Cluster: cluster,
	}

	if req.GetImportOptions().GetFileType() == v1beta1.ImportFileType_CSV {
		csvBytes, err := json.Marshal(req.GetImportOptions().GetCsvFormat())
		if err != nil {
			log.Error(ctx, "failed to marshal CSV format", zap.Error(err))
			return nil, bizerr.New(errorcode.ErrCodeInternalError, errorcode.InternalErrorMessage)
		}
		imp.CSVFormat = string(csvBytes)
	}

	imp.Creator = utils.BuildCreatorID(ctx)

	return imp, nil
}

func BuildGCSImport(
	ctx context.Context,
	req *v1beta1.CreateImportRequest,
	cluster *tidbmgmtservicev1.Cluster,
) (*entity.ImportTask, bizerr.BizError) {
	accountKey := req.Source.GetGcs().GetServiceAccountKey()
	decodeString, err := DecodeBase64(accountKey)
	if err != nil {
		return nil, bizerr.New(bizerr.ErrCodeParams, "invalid service account key")
	}
	imp := &entity.ImportTask{
		Import: &entity.Import{
			ID:             entity.GenerateIDWithPrefix(constant.ImportIDPrefix),
			ProjectID:      strconv.FormatUint(cluster.GetOriginalProjectId(), 10),
			TenantID:       strconv.FormatUint(cluster.OrganizationId, 10),
			ClusterID:      req.ClusterId,
			SourceURL:      req.Source.GetGcs().GetUri(),
			DataSourceType: pbDataSource2Entity(req.GetSource().GetType()),
			DataFormat:     PbFileType2Entity(req.GetImportOptions().GetFileType()),
			NeedDropTables: false,
			GCSCredentials: base64.URLEncoding.EncodeToString(decodeString),
		},
		Cluster: cluster,
	}

	if req.GetImportOptions().GetFileType() == v1beta1.ImportFileType_CSV {
		csvBytes, err := json.Marshal(req.GetImportOptions().GetCsvFormat())
		if err != nil {
			log.Error(ctx, "failed to marshal CSV format", zap.Error(err))
			return nil, bizerr.New(errorcode.ErrCodeInternalError, errorcode.InternalErrorMessage)
		}
		imp.CSVFormat = string(csvBytes)
	}

	imp.Creator = utils.BuildCreatorID(ctx)

	return imp, nil
}

// DecodeBase64 decodes a Base64 URL encoded string.
// If the input string's length is a multiple of 4, it uses the standard Base64 decoding
// with padding. Otherwise, it uses the raw Base64 decoding without padding.
func DecodeBase64(blob string) ([]byte, error) {
	if len(blob)%4 == 0 {
		// decode with padding
		return base64.StdEncoding.DecodeString(blob)
	}
	// decode without padding
	return base64.RawStdEncoding.DecodeString(blob)
}

func Pb2CsvConfig(format *v1beta1.CSVFormat) *config.CSVConfig {
	if format == nil {
		return nil
	}

	c := &config.CSVConfig{
		Separator:         format.GetSeparator(),
		Delimiter:         format.GetDelimiter(),
		Header:            format.GetHeader(),
		TrimLastSep:       format.GetTrimLastSeparator(),
		NotNull:           format.GetNotNull(),
		Null:              config.StringOrStringSlice{format.GetNull()},
		HeaderSchemaMatch: true,
	}
	if format.GetBackslashEscape() {
		c.EscapedBy = `\`
	}

	return c
}