package lightning

import (
	"context"
	"path/filepath"

	"github.com/pingcap/errors"
	"github.com/pingcap/tidb/pkg/lightning/mydump"
	filter "github.com/pingcap/tidb/pkg/util/table-filter"
)

// FileScan is a helper to scan files from a given path and route them to different types, to generate a list of MDDatabaseMeta.
type FileScan struct {
	fileRouter mydump.FileRouter
	fileIter   mydump.FileIterator
	filter     filter.Filter
	charSet    string

	dbFileInfoMap map[string]*DBFileInfo
}

// TableFileInfo is a helper struct to store table schema and data files.
type TableFileInfo struct {
	// TableName is the name of the table.
	TableName string
	// SchemaFile is the schema file of the table.
	SchemaFile *mydump.FileInfo
	// DataFiles is the list of data files of the table.
	DataFiles []mydump.FileInfo
}

// DBFileInfo is a helper struct to store database schema and tables.
type DBFileInfo struct {
	// DBName is the name of the database.
	DBName string
	// SchemaFile is the schema file of the database.
	SchemaFile *mydump.FileInfo
	// Tables is the list of tables in the database organized in a map.
	Tables map[string]*TableFileInfo
}

// NewFileScan creates a new FileScan.
func NewFileScan(router mydump.FileRouter, iter mydump.FileIterator, filter filter.Filter, charSet string) *FileScan {
	return &FileScan{
		fileRouter:    router,
		fileIter:      iter,
		filter:        filter,
		charSet:       charSet,
		dbFileInfoMap: make(map[string]*DBFileInfo),
	}
}

// GetMDDatabaseMetas returns a list of MDDatabaseMeta.
func (f *FileScan) GetMDDatabaseMetas(ctx context.Context) ([]*mydump.MDDatabaseMeta, error) {
	if err := f.fileIter.IterateFiles(ctx, f.addToFileInfos); err != nil {
		return nil, err
	}

	dbMetas := make([]*mydump.MDDatabaseMeta, 0)
	for _, dbFileInfo := range f.dbFileInfoMap {
		dbMeta := mydump.NewMDDatabaseMeta(f.charSet)
		dbMeta.Name = dbFileInfo.DBName
		if dbFileInfo.SchemaFile != nil {
			dbMeta.SchemaFile = *dbFileInfo.SchemaFile
		}
		dbMeta.Tables = make([]*mydump.MDTableMeta, 0)
		for _, tblFileInfo := range dbFileInfo.Tables {
			tblMeta := mydump.NewMDTableMeta(f.charSet)
			tblMeta.DB = dbFileInfo.DBName
			tblMeta.Name = tblFileInfo.TableName
			if tblFileInfo.SchemaFile != nil {
				tblMeta.SchemaFile = *tblFileInfo.SchemaFile
			}
			tblMeta.DataFiles = tblFileInfo.DataFiles
			dbMeta.Tables = append(dbMeta.Tables, tblMeta)
		}
		dbMetas = append(dbMetas, dbMeta)
	}

	return dbMetas, nil
}

// addToFileInfos adds a file to the internal file info map.
func (f *FileScan) addToFileInfos(ctx context.Context, path string, size int64) error {
	res, err := f.fileRouter.Route(filepath.ToSlash(path))
	if err != nil {
		return errors.Annotate(err, "route file failed")
	}
	if res == nil {
		return nil
	}
	info := mydump.FileInfo{
		TableName: filter.Table{Schema: res.Schema, Name: res.Name},
		FileMeta:  mydump.SourceFileMeta{Path: path, Type: res.Type, Compression: res.Compression, SortKey: res.Key, FileSize: size},
	}
	if shouldSkip(f.filter, res) {
		return nil
	}

	switch res.Type {
	case mydump.SourceTypeSchemaSchema:
		dbFileInfo, ok := f.dbFileInfoMap[res.Schema]
		if !ok {
			dbFileInfo = &DBFileInfo{
				DBName:     res.Schema,
				SchemaFile: &info,
				Tables:     make(map[string]*TableFileInfo),
			}
			f.dbFileInfoMap[res.Schema] = dbFileInfo
		} else {
			if dbFileInfo.SchemaFile != nil {
				return errors.Errorf("duplicate schema file for database %s", res.Schema)
			}
			dbFileInfo.SchemaFile = &info
		}
	case mydump.SourceTypeTableSchema:
		dbFileInfo, ok := f.dbFileInfoMap[res.Schema]
		if !ok {
			dbFileInfo = &DBFileInfo{
				DBName:     res.Schema,
				SchemaFile: nil,
				Tables:     make(map[string]*TableFileInfo),
			}
			f.dbFileInfoMap[res.Schema] = dbFileInfo
		}
		tableFileInfo, ok := dbFileInfo.Tables[res.Name]
		if !ok {
			tableFileInfo = &TableFileInfo{
				TableName:  res.Name,
				SchemaFile: &info,
				DataFiles:  make([]mydump.FileInfo, 0),
			}
			dbFileInfo.Tables[res.Name] = tableFileInfo
		} else {
			if tableFileInfo.SchemaFile != nil {
				return errors.Errorf("duplicate schema file for table %s.%s", res.Schema, res.Name)
			}
			tableFileInfo.SchemaFile = &info
		}
	case mydump.SourceTypeSQL, mydump.SourceTypeCSV, mydump.SourceTypeParquet:
		dbFileInfo, ok := f.dbFileInfoMap[res.Schema]
		if !ok {
			dbFileInfo = &DBFileInfo{
				DBName:     res.Schema,
				SchemaFile: nil,
				Tables:     make(map[string]*TableFileInfo),
			}
			f.dbFileInfoMap[res.Schema] = dbFileInfo
		}
		tableFileInfo, ok := dbFileInfo.Tables[res.Name]
		if !ok {
			tableFileInfo = &TableFileInfo{
				TableName:  res.Name,
				SchemaFile: nil,
				DataFiles:  make([]mydump.FileInfo, 0),
			}
			dbFileInfo.Tables[res.Name] = tableFileInfo
		}
		tableFileInfo.DataFiles = append(tableFileInfo.DataFiles, info)
	}
	return nil
}