package pimport

import (
	"context"
	"encoding/json"
	"fmt"
	"net/url"
	"strings"

	"github.com/pingcap/errcode"
	"github.com/pingcap/errors"
	agentpb "github.com/tidbcloud/dataflow-service/idl/pbgen/proto/agent"
	commonpb "github.com/tidbcloud/dataflow-service/idl/pbgen/proto/common"
	dataflowpb "github.com/tidbcloud/dataflow-service/idl/pbgen/proto/dataflow"
	"github.com/tidbcloud/dataflow-service/internal/db/entity"
	pimport "github.com/tidbcloud/dataflow-service/internal/model/import"
	"github.com/tidbcloud/dataflow-service/internal/utils"
	errhandler "github.com/tidbcloud/dataflow-service/pkg/error-handler"
	"github.com/tidbcloud/dataflow-service/pkg/s3"
	"github.com/tidbcloud/idl/pbgen/proto/openapi/dp"
	"go.uber.org/zap"
)

func ImportEntity2ImportItem(logger *zap.Logger, imp *pimport.ImportTask) *dp.ImportItem {
	plainURL, _, err := imp.ParseSourceURL()
	if err != nil {
		logger.Error("cannot parse the source URL of the import entity", zap.Error(err))
	}
	resp := &dp.ImportItem{
		Metadata: &dp.ImportMetadata{
			Id:              uint64(imp.ID),
			Name:            imp.Name,
			CreateTimestamp: imp.CreatedAt.Unix(),
		},
		Spec: &dp.ImportSpec{
			Source: &dp.ImportSource{
				Uri: plainURL,
			},
			Target: &dp.ImportTarget{},
		},
		Status: getImportStatusFromEntity(imp),
	}
	resp.Spec.Target.Tables = getImportTargetTablesFromEntity(imp, logger)

	switch imp.DataSourceType {
	case entity.ImportDataSourceS3:
		resp.Spec.Source.Type = dp.ImportSource_S3
		// get access config
		switch {
		case len(imp.AwsRoleArn) > 0:
			resp.Spec.Source.Access = &dp.ImportSource_AwsAssumeRoleAccess{
				AwsAssumeRoleAccess: &dp.AwsAssumeRoleAccess{
					AssumeRole: imp.AwsRoleArn,
				},
			}
		case imp.LocalImportCfg != nil:
			// local file import, no need to show access config
		default:
			resp.Spec.Source.Access = &dp.ImportSource_AwsKeyAccess{
				AwsKeyAccess: &dp.AwsKeyAccess{
					AccessKeyId:     "REDACTED",
					SecretAccessKey: "REDACTED",
				},
			}
		}
	case entity.ImportDataSourceGcs:
		resp.Spec.Source.Type = dp.ImportSource_GCS
	}
	switch imp.DataFormat {
	case entity.ImportDataFormatCsv:
		resp.Spec.Source.Format = &dp.ImportSourceFormat{
			Type: dp.ImportSourceFormat_CSV,
		}
		csvConfig, err := entityCSVFormatStrToConfig(imp.CustomCSVFormat)
		if err != nil {
			logger.Error("convert CSV format string to CSV config error", zap.Error(err))
		} else {
			resp.Spec.Source.Format.FormatConfig = &dp.ImportSourceFormat_CsvConfig{
				CsvConfig: csvConfig,
			}
		}
	case entity.ImportDataFormatAuroraSnapshot:
		resp.Spec.Source.Format = &dp.ImportSourceFormat{
			Type: dp.ImportSourceFormat_AURORA_SNAPSHOT,
		}
	case entity.ImportDataFormatParquet:
		resp.Spec.Source.Format = &dp.ImportSourceFormat{
			Type: dp.ImportSourceFormat_PARQUET,
		}
	case entity.ImportDataFormatSql:
		resp.Spec.Source.Format = &dp.ImportSourceFormat{
			Type: dp.ImportSourceFormat_SQL,
		}
	}
	if imp.LocalImportCfg != nil {
		// don't show the sourceUrl when import local file
		resp.Spec.Source.Uri = imp.LocalImportCfg.OriginalFileName
		resp.Spec.Source.Type = dp.ImportSource_LOCAL_FILE
	}
	return resp
}

func getImportStatusFromEntity(imp *pimport.ImportTask) *dp.ImportStatus {
	resultStatus := &dp.ImportStatus{
		Phase:        entityStatusToPBStatus(imp.Status),
		ErrorMessage: errhandler.PrettyMessage(imp.Message, uint32(imp.DataFormat)),
		Progress: &dp.ImportProgress{
			ImportProgress:     float64(imp.CompletedPercent),
			ValidationProgress: float64(imp.PostImportCompletedPercent),
		},
	}
	if imp.TotalSize != nil {
		resultStatus.SourceTotalSizeBytes = *imp.TotalSize
	}
	switch resultStatus.Phase {
	case dp.ImportStatus_PREPARING:
		resultStatus.Progress.ImportProgress = 0
		resultStatus.Progress.ValidationProgress = 0
	case dp.ImportStatus_COMPLETED:
		resultStatus.Progress.ImportProgress = 100
		resultStatus.Progress.ValidationProgress = 100
	}
	if imp.StartedAt != nil {
		resultStatus.StartTimestamp = imp.StartedAt.Unix()
	}
	if imp.FinishedAt != nil {
		resultStatus.EndTimestamp = imp.FinishedAt.Unix()
	}
	// TODO: convertSyntaxErr(resp, imp, logger)
	return resultStatus
}

func entityStatusToPBStatus(status entity.ImportStatus) dp.ImportStatus_ImportTaskPhase {
	switch status {
	case entity.ImportStatusPreparing:
		return dp.ImportStatus_PREPARING
	case entity.ImportStatusImporting:
		return dp.ImportStatus_IMPORTING
	case entity.ImportStatusCompleted:
		return dp.ImportStatus_COMPLETED
	case entity.ImportStatusFailed:
		return dp.ImportStatus_FAILED
	case entity.ImportStatusCanceled:
		return dp.ImportStatus_CANCELED
	case entity.ImportStatusCanceling:
		return dp.ImportStatus_CANCELING
	default:
		return dp.ImportStatus_UNSPECIFIED
	}
}

func entityCSVFormatStrToConfig(customCSVFormatStr string) (*dp.ImportSourceCSVConfig, error) {
	customCSVFormat := new(dataflowpb.CustomCSVFormat)
	if err := json.Unmarshal([]byte(customCSVFormatStr), customCSVFormat); err != nil {
		return nil, errors.Annotate(err, "parse custom CSV format JSON string error")
	}
	return &dp.ImportSourceCSVConfig{
		Delimiter:       customCSVFormat.Separator,
		Quote:           customCSVFormat.Delimiter,
		BackslashEscape: customCSVFormat.BackslashEscape,
		HasHeaderRow:    customCSVFormat.Header,
	}, nil
}

func getImportTargetTablesFromEntity(imp *pimport.ImportTask, logger *zap.Logger) []*dp.ImportTargetTable {
	var targetTableInfos []*dp.ImportTargetTable
	for _, tableMapping := range imp.TargetTableInfos.Infos {
		targetTableInfo := &dp.ImportTargetTable{
			DatabaseName:    tableMapping.Schema,
			TableName:       tableMapping.Table,
			FileNamePattern: tableMapping.CustomFile,
		}

		targetTableInfos = append(targetTableInfos, targetTableInfo)
	}
	if len(targetTableInfos) == 0 && imp.LocalImportCfg != nil {
		targetTableInfos = append(targetTableInfos,
			&dp.ImportTargetTable{
				DatabaseName:    imp.LocalImportCfg.TargetDb,
				TableName:       imp.LocalImportCfg.TargetTable,
				FileNamePattern: imp.LocalImportCfg.OriginalFileName,
			},
		)
	}

	return targetTableInfos
}

func ImportSpec2Entity(ctx context.Context, spec *dp.ImportSpec, stub *entity.ImportUploadStub, logger *zap.Logger) (*pimport.ImportTask, error) {
	// preprocess some string input value
	sourceURL := strings.TrimSpace(spec.GetSource().GetUri())
	sourceType := spec.GetSource().GetType()
	imp := &pimport.ImportTask{
		Import: entity.Import{
			SourceURL:      sourceURL,
			DataSourceType: specSourceType2EntityDataSource(sourceType),
			DataFormat:     specFormatType2EntityDataFormat(spec.GetSource().GetFormat().GetType()),
		},
	}
	if sourceType == dp.ImportSource_LOCAL_FILE {
		if stub == nil {
			return nil, errors.New("stub is empty when creating a local file import task")
		}

		imp.SourceURL = stub.UploadedURL
		imp.BucketRegion = stub.UploadedRegion
		target := spec.GetTarget()
		if len(target.GetTables()) != 1 {
			return nil, errcode.NewInvalidInputErr(errors.New("Only one target table is supported for local file import"))
		}
		targetTable := target.GetTables()[0]

		imp.LocalImportCfg = &entity.LocalImportCfg{
			FileName:         stub.UploadedFileName,
			TargetDb:         targetTable.GetDatabaseName(),
			TargetTable:      targetTable.GetTableName(),
			OriginalFileName: stub.LocalFileName,
		}
	} else {
		if sourceType == dp.ImportSource_S3 {
			bucketRegion, err := getBucketRegion(ctx, sourceURL)
			if err != nil {
				logger.Error("get bucket region error", zap.Error(err))
				return nil, errcode.NewInvalidInputErr(errors.Errorf("Failed to get bucket region of the S3 URI: '%s'", sourceURL))
			}
			imp.BucketRegion = bucketRegion
			if err := applyS3AccessInfoToEntity(spec.GetSource(), imp); err != nil {
				logger.Error("apply S3 access info to entity failed", zap.Error(err))
				return nil, errcode.NewInvalidInputErr(errors.New("Failed to parse the S3 access information"))
			}
		}
		applyTargetSpecToEntity(spec.GetTarget(), imp)
	}
	if spec.GetSource().GetFormat().GetType() == dp.ImportSourceFormat_CSV {
		entityCSVConfig, err := buildCSVConfigInEntity(spec.GetSource().GetFormat().GetCsvConfig(), logger)
		if err != nil {
			logger.Error("build CSV config error", zap.Error(err))
			return nil, utils.WrapAsErrorCodeErr(err, "parse CSV config error")
		}

		csvConfigBytes, err := json.Marshal(entityCSVConfig)
		if err != nil {
			return nil, errors.Annotate(err, "failed to marshal csv config")
		}

		imp.CustomCSVFormat = string(csvConfigBytes)
	}
	return imp, nil
}

func specSourceType2EntityDataSource(specSourceType dp.ImportSource_ImportSourceType) entity.DataSource {
	switch specSourceType {
	case dp.ImportSource_S3:
		return entity.ImportDataSourceS3
	case dp.ImportSource_GCS:
		return entity.ImportDataSourceGcs
	case dp.ImportSource_LOCAL_FILE:
		// TODO: if later serverless support GCS, maybe this will change
		return entity.ImportDataSourceS3
	}
	return entity.ImportDataSourceUnknown
}

func specFormatType2EntityDataFormat(specFormatType dp.ImportSourceFormat_ImportSourceFormatType) entity.DataFormat {
	switch specFormatType {
	case dp.ImportSourceFormat_CSV:
		return entity.ImportDataFormatCsv
	case dp.ImportSourceFormat_PARQUET:
		return entity.ImportDataFormatParquet
	case dp.ImportSourceFormat_SQL:
		return entity.ImportDataFormatSql
	case dp.ImportSourceFormat_AURORA_SNAPSHOT:
		return entity.ImportDataFormatAuroraSnapshot
	}

	return entity.ImportDataFormatUnknown
}

func applyTargetSpecToEntity(target *dp.ImportTarget, imp *pimport.ImportTask) {
	imp.TargetTableInfos = entity.TargetTableInfos{
		Infos: make([]entity.TargetTableInfo, 0, len(target.GetTables())),
	}

	for _, tbl := range target.GetTables() {
		imp.TargetTableInfos.Infos = append(imp.TargetTableInfos.Infos, entity.TargetTableInfo{
			Schema:     tbl.DatabaseName,
			Table:      tbl.TableName,
			CustomFile: tbl.FileNamePattern,
		})
	}
}

func applyS3AccessInfoToEntity(sourceSpec *dp.ImportSource, imp *pimport.ImportTask) error {
	accessConfig := new(entity.ImportSourceAccessConfig)
	switch {
	case sourceSpec.GetAwsAssumeRoleAccess() != nil:
		roleARN := sourceSpec.GetAwsAssumeRoleAccess().GetAssumeRole()
		imp.AwsRoleArn = roleARN
		accessConfig.AWSRoleARN = roleARN
	case sourceSpec.GetAwsKeyAccess() != nil:
		access := sourceSpec.GetAwsKeyAccess()
		accessConfig.AWSAccessKeyID = access.GetAccessKeyId()
		accessConfig.AWSSecretAccessKey = access.GetSecretAccessKey()
	}
	if err := imp.SetAccessConfigInSourceURL(accessConfig); err != nil {
		return errors.Annotate(err, "set access config in source URL failed")
	}
	return nil
}

func getBucketRegion(ctx context.Context, s3URI string) (string, error) {
	u, err := s3.NewS3URL(s3URI)
	if err != nil {
		return "", errors.Annotatef(err, "cannot parse S3 URL '%s'", s3URI)
	}
	region, err := u.GetBucketRegion(ctx)
	if err != nil {
		return "", errors.Annotate(err, "get bucket region error")
	}
	return region, nil
}

func buildCSVConfigInEntity(csvConfig *dp.ImportSourceCSVConfig, logger *zap.Logger) (*dataflowpb.CustomCSVFormat, error) {
	var err error
	var entityCSVConfig *dataflowpb.CustomCSVFormat
	if csvConfig == nil {
		entityCSVConfig = &dataflowpb.CustomCSVFormat{
			Separator:       ",",
			Delimiter:       `"`,
			Header:          true,
			BackslashEscape: false,
		}
	} else {
		quote := strings.TrimSpace(csvConfig.GetQuote())
		if len(quote) > 1 {
			originalQuote := quote
			quote, err = utils.RestoreSpecialChars(quote)
			if err != nil {
				logger.Error("restore special characters for quote chars failed", zap.Error(err))
				return nil, errcode.NewInvalidInputErr(errors.Errorf("invalid quote character in CSV config: %s", originalQuote))
			}
		}
		delimiter := csvConfig.GetDelimiter()
		if len(delimiter) > 1 {
			originalDelimiter := delimiter
			delimiter, err = utils.RestoreSpecialChars(delimiter)
			if err != nil {
				logger.Error("restore special characters for delimiter chars failed", zap.Error(err))
				return nil, errcode.NewInvalidInputErr(errors.Errorf("invalid delimiter character in CSV config: %s", originalDelimiter))
			}
		}
		entityCSVConfig = &dataflowpb.CustomCSVFormat{
			Separator:       delimiter,
			Delimiter:       quote,
			Header:          csvConfig.GetHasHeaderRow(),
			BackslashEscape: csvConfig.GetBackslashEscape(),
		}
	}
	return entityCSVConfig, nil
}

func previewResult2TableSchema(previewResult *agentpb.Table, logger *zap.Logger) (*dp.TableSchema, error) {
	colDefs := []*dp.ColumnDefinition{}
	if len(previewResult.GetRowHeaders()) != len(previewResult.GetRowDataTypes()) {
		logger.Error("row header count and row data type count mismatch",
			zap.Int("row_header_count", len(previewResult.GetRowHeaders())),
			zap.Int("row_type_count", len(previewResult.GetRowDataTypes())),
		)
		return nil, errcode.NewInternalErr(errors.New("The returned preview data is abnormal for extracting table definition"))
	}
	for i, rowHeader := range previewResult.GetRowHeaders() {
		colName := rowHeader
		if len(colName) == 0 {
			colName = fmt.Sprintf("column_%d", i+1)
		}
		rowType := previewResult.GetRowDataTypes()[i]
		colDefs = append(colDefs, &dp.ColumnDefinition{
			ColumnName: colName,
			ColumnType: rowType,
		})
	}
	return &dp.TableSchema{
		ColumnDefinitions: colDefs,
	}, nil
}

func previewResult2TableData(previewResult *agentpb.Table, logger *zap.Logger) (*dp.TableData, error) {
	rows := []*dp.TableDataRow{}
	colNames := []string{}
	for i, rowHeader := range previewResult.GetRowHeaders() {
		colName := rowHeader
		if len(colName) == 0 {
			colName = fmt.Sprintf("column_%d", i+1)
		}
		colNames = append(colNames, colName)
	}
	for _, row := range previewResult.GetRowValues() {
		rows = append(rows, &dp.TableDataRow{
			Columns: append([]string{}, row.GetColumnValues()...),
		})
	}
	return &dp.TableData{
		ColumnNames: colNames,
		Rows:        rows,
	}, nil
}

func tableDefinition2LocalImportConfig(tableDef *dp.TableDefinition) *dataflowpb.LocalImportCfg {
	colDefs := tableDef.GetSchema().GetColumnDefinitions()
	if len(colDefs) <= 0 {
		return nil
	}
	targetColDefs := []*dataflowpb.ColumnInfo{}
	for _, colDef := range colDefs {
		targetColDefs = append(targetColDefs, &dataflowpb.ColumnInfo{
			ColumnName: colDef.GetColumnName(),
			ColumnType: colDef.GetColumnType(),
		})
	}
	return &dataflowpb.LocalImportCfg{
		TargetTable: &commonpb.Table{
			Schema: tableDef.GetDatabaseName(),
			Table:  tableDef.GetTableName(),
		},
		Columns:   targetColDefs,
		PkColumns: append([]string{}, tableDef.GetSchema().GetPrimaryKeyColumns()...),
	}
}

func tweakImportEntityForCreate(imp *pimport.ImportTask) error {
	if imp.LocalImportCfg != nil {
		// same as GenCompletedUrl in internal/model/import/naming.go
		fullSourceURL, err := url.JoinPath(imp.SourceURL, imp.LocalImportCfg.FileName)
		if err != nil {
			return errors.Annotate(err, "cannot generate the new source URL for the stub")
		}
		imp.SourceURL = fullSourceURL
	}
	return nil
}