package pimport

import (
	"context"
	"fmt"
	"io"
	"strings"

	"github.com/dustin/go-humanize"
	"github.com/pingcap/errors"
	"github.com/pingcap/tidb/br/pkg/storage"
	"github.com/pingcap/tidb/pkg/lightning/config"
	"github.com/pingcap/tidb/pkg/lightning/mydump"
	"github.com/pingcap/tidb/pkg/lightning/worker"
	"github.com/pingcap/tidb/pkg/parser/mysql"
	ptypes "github.com/pingcap/tidb/pkg/parser/types"
	"github.com/pingcap/tidb/pkg/types"
	agentpb "github.com/tidbcloud/dataflow-service/idl/pbgen/proto/agent"
	"github.com/tidbcloud/dataflow-service/idl/pbgen/proto/dataflow"
	"github.com/tidbcloud/dataflow-service/internal/db/entity"
	str_utils "github.com/tidbcloud/dataflow-service/pkg/string-utils"
	"github.com/tidbcloud/pkgv2/log"
	"go.uber.org/zap"
)

type ErrCSVSyntax struct{}

func (e *ErrCSVSyntax) Error() string {
	return "CSV syntax error"
}

type ErrCSVRowOverflow struct {
	RowNum int
}

func (e *ErrCSVRowOverflow) Error() string {
	return fmt.Sprintf("the data num of one row is greater than header on %s row", humanize.Ordinal(e.RowNum+1))
}

type ErrHeaderDataColumnMismatch struct {
	HeaderColumnsCount int
	DataColumnsCount   int
}

func (e *ErrHeaderDataColumnMismatch) Error() string {
	return fmt.Sprintf("the header column count (%d) does not equal to data column count (%d).", e.HeaderColumnsCount, e.DataColumnsCount)
}

type ErrPreviewHasWarning struct {
	WarningMessage string
}

func (e *ErrPreviewHasWarning) Error() string {
	return fmt.Sprintf("the preview has warning: %s", e.WarningMessage)
}

type ImportPreviewer interface {
	Preview(ctx context.Context, limitCount uint) ([]*agentpb.Schema, error)
}

// PreviewFileMapping defines the source file to target table mapping.
type PreviewFileMapping struct {
	DatabaseName string
	TableName    string
	FileName     string
}

// LocalImportPreviewer is the implementation that the preview can be run on the server locally without any problems.
// Usually this means:
//   - The server can access the source URL without any problems. (no assuming role issues)
//   - The list of files to preview is clear from the import entity:
//     no need to scan the source to get how many tables are there,
//     and no need to find the first data file name for each target table to do the preview.
//
// As a result, currently the previewer only applies to local file import situations, which meets all the pre-conditions.
// TODO: If later, the import supports importing a single S3 file object right into the TiDB, with provided AK/SK,
// it can also use this implementation
type LocalImportPreviewer struct {
	PreviewFileMappings []*PreviewFileMapping
	CSVFormat           *dataflow.CustomCSVFormat
	SourceURL           string
	DataFormat          entity.DataFormat
}

func (p *LocalImportPreviewer) Preview(ctx context.Context, limitCount uint) ([]*agentpb.Schema, error) {
	storageBackend, err := storage.ParseBackend(p.SourceURL, nil)
	if err != nil {
		return nil, errors.Annotate(err, "failed to parse external storage backend")
	}
	store, err := storage.New(ctx, storageBackend, &storage.ExternalStorageOptions{})
	if err != nil {
		return nil, errors.Annotate(err, "failed to new storage backend")
	}

	schemaPreviewMap := map[string]*agentpb.Schema{}
	for _, fileMapping := range p.PreviewFileMappings {
		previewTable, err := p.doSingleTablePreview(ctx, store, fileMapping, limitCount)
		if err != nil {
			return nil, err
		}
		schemaName := fileMapping.DatabaseName
		schemaPreview, ok := schemaPreviewMap[schemaName]
		if !ok {
			schemaPreview = &agentpb.Schema{
				SchemaName: schemaName,
			}
			schemaPreviewMap[schemaName] = schemaPreview
		}
		schemaPreview.Tables = append(schemaPreview.GetTables(), previewTable)
	}
	result := []*agentpb.Schema{}
	for _, schemaPreview := range schemaPreviewMap {
		result = append(result, schemaPreview)
	}
	return result, nil
}

func (p *LocalImportPreviewer) newParser(ctx context.Context, store storage.ExternalStorage, fileMapping *PreviewFileMapping) (mydump.Parser, error) {
	var parser mydump.Parser
	switch p.DataFormat {
	case entity.ImportDataFormatAuroraSnapshot, entity.ImportDataFormatParquet:
		// MaxUploadSize is less than smallParquetFileThreshold
		reader, err := mydump.OpenParquetReader(ctx, store, fileMapping.FileName, MaxUploadSize)
		if err != nil {
			return nil, errors.Annotate(err, "failed to open parquet reader")
		}
		parser, err = mydump.NewParquetParser(ctx, store, reader, fileMapping.FileName)
		if err != nil {
			return nil, errors.Annotate(err, "failed to create the parquet parser")
		}
	case entity.ImportDataFormatCsv:
		if p.CSVFormat == nil {
			return nil, errors.New("CSV format configuration is missing")
		}
		reader, err := store.Open(ctx, fileMapping.FileName, nil)
		if err != nil {
			return nil, errors.Annotatef(err, "failed to open the path: %s", fileMapping.FileName)
		}
		ioWorkers := worker.NewPool(ctx, 5, "preview_local")

		cfg := config.MydumperRuntime{
			CSV: config.CSVConfig{
				FieldsTerminatedBy: p.CSVFormat.Separator,
				FieldsEnclosedBy:   p.CSVFormat.Delimiter,
				TrimLastEmptyField: p.CSVFormat.GetTrimLastSeparator(),
				Header:             p.CSVFormat.Header,
				NotNull:            p.CSVFormat.GetNotNull(),
				FieldNullDefinedBy: config.StringOrStringSlice{p.CSVFormat.GetNull()},
				HeaderSchemaMatch:  p.CSVFormat.Header,
			},
		}
		if p.CSVFormat.GetBackslashEscape() {
			cfg.CSV.FieldsEscapedBy = `\`
		}

		charsetConvertor, err := mydump.NewCharsetConvertor(cfg.DataCharacterSet, cfg.DataInvalidCharReplace)
		if err != nil {
			return nil, errors.Annotate(err, "failed to new charset convertor")
		}
		parser, err = mydump.NewCSVParser(ctx, &cfg.CSV, reader, int64(config.ReadBlockSize), ioWorkers, cfg.CSV.Header, charsetConvertor)
		if err != nil {
			return nil, errors.Annotate(err, "failed to create csv parser")
		}
	case entity.ImportDataFormatSql:
		reader, err := store.Open(ctx, fileMapping.FileName, nil)
		if err != nil {
			return nil, errors.Annotatef(err, "failed to open the path: %s", fileMapping.FileName)
		}
		ioWorkers := worker.NewPool(ctx, 5, "preview_local")
		defaultSQLMode, _ := mysql.GetSQLMode(mysql.DefaultSQLMode)
		parser = mydump.NewChunkParser(ctx, defaultSQLMode, reader, int64(config.ReadBlockSize), ioWorkers)
	default:
		return nil, errors.Errorf("unsupported data format: %v", p.DataFormat)
	}
	return parser, nil
}

func (p *LocalImportPreviewer) doSingleTablePreview(ctx context.Context, store storage.ExternalStorage, fileMapping *PreviewFileMapping, limitCount uint) (*agentpb.Table, error) {
	logger := log.WithContextL(ctx)
	parser, err := p.newParser(ctx, store, fileMapping)
	if err != nil {
		return nil, errors.Annotate(err, "create source data parser error")
	}
	//nolint: errcheck
	defer parser.Close()

	result := &agentpb.Table{
		SchemaName: fileMapping.DatabaseName,
		TableName:  fileMapping.TableName,
		RowValues:  make([]*agentpb.Row, limitCount),
	}

	var rowDataTypes []*ptypes.FieldType
	for i := 0; i < int(limitCount); i++ {
		if err := parser.ReadRow(); err != nil {
			if errors.Cause(err) == io.EOF {
				result.RowValues = result.RowValues[:i]
				break
			}
			if strings.Contains(err.Error(), "syntax error") {
				return nil, new(ErrCSVSyntax)
			}

			logger.Error("failed to read row in parser", zap.Error(err))
			return nil, errors.Annotate(err, "failed to read row in parser")
		}

		readRow := parser.LastRow().Row
		writeRow := &agentpb.Row{ColumnValues: make([]string, 0, len(readRow))}
		if len(rowDataTypes) == 0 {
			rowDataTypes = make([]*ptypes.FieldType, len(readRow))
		}
		for j, datum := range readRow {
			columnVal, err := datum.ToString()
			if err != nil {
				return nil, errors.Annotate(err, "failed to get string from a datum")
			}
			writeRow.ColumnValues = append(writeRow.ColumnValues, str_utils.ForceValidUTF8(columnVal))
			fp := inferDataType(datum, writeRow.ColumnValues[j])

			if len(rowDataTypes) <= j {
				logger.Warn("the data num of one row is greater than previous row", zap.String("row", humanize.Ordinal(i+1)))
				return nil, &ErrCSVRowOverflow{RowNum: i}
			}
			if rowDataTypes[j] == nil {
				rowDataTypes[j] = fp
			} else {
				rowDataTypes[j] = types.AggFieldType([]*types.FieldType{rowDataTypes[j], fp})
			}
		}

		result.RowValues[i] = writeRow
	}
	for _, tp := range rowDataTypes {
		result.RowDataTypes = append(result.RowDataTypes, strings.ToUpper(tp.CompactStr()))
	}

	if len(parser.Columns()) != 0 {
		result.RowHeaders = parser.Columns()
		if len(parser.Columns()) != len(rowDataTypes) {
			logger.Error("the data header num does not equal to data type.", zap.Int("data type num", len(rowDataTypes)))
			return nil, &ErrHeaderDataColumnMismatch{
				HeaderColumnsCount: len(parser.Columns()),
				DataColumnsCount:   len(rowDataTypes),
			}
		}
	} else {
		result.RowHeaders = make([]string, len(rowDataTypes))
	}

	return result, nil
}