package precheck

import (
	"context"
	"errors"
	"fmt"
	"sync"

	"github.com/tidbcloud/pkgv2/bizerr"
	"github.com/tidbcloud/pkgv2/import/cloud_storage"
	"github.com/tidbcloud/pkgv2/log"
	"go.uber.org/zap"
	"golang.org/x/sync/errgroup"
)

const (
	displayLimit         = 5
	concurrentCheckLimit = 5
)

// PrecheckStatus represents the status of precheck operation
type PrecheckStatus string

const (
	// PrecheckStatusCompleted indicates the precheck completed successfully
	PrecheckStatusCompleted PrecheckStatus = "completed"
	// PrecheckStatusTruncated indicates the precheck was truncated due to file listing limits
	PrecheckStatusTruncated PrecheckStatus = "truncated"
	// PrecheckStatusTimeout indicates the precheck was cancelled due to timeout
	PrecheckStatusTimeout PrecheckStatus = "timeout"
)

type PrecheckResult struct {
	TableResults []TableResult
	Status       PrecheckStatus
}

type TableResult struct {
	DataFiles []string
	// SchemaFile is the path of the schema file, if table exists, it should be empty.
	SchemaFile string
	// UseSchemaFile indicates whether the schema file is used for import.
	UseSchemaFile bool
	Table         Table
	ErrorMsg      string
}

type DbMeta struct {
	Name       string
	SchemaFile *RouteResult
	TableMap   map[string]*TableMeta
}

type TableMeta struct {
	Name string
	// use slice to detect duplicated schema files error
	SchemaFiles []RouteResult
	DataFiles   []RouteResult
}

type Checker struct {
	store                cloud_storage.CloudStorage
	fileRouter           FileRouter
	schemaChecker        SchemaChecker
	useCustomFileMapping bool
}

// NewChecker returns a new Checker.
//
// If useCustomFileMapping is true, the file mapping rules are defined by the user,
// and the target tables must already exist.
func NewChecker(store cloud_storage.CloudStorage,
	fileRouter FileRouter,
	schemaChecker SchemaChecker,
	useCustomFileMapping bool) *Checker {
	return &Checker{
		store:                store,
		fileRouter:           fileRouter,
		schemaChecker:        schemaChecker,
		useCustomFileMapping: useCustomFileMapping,
	}
}

// CheckSchemasResult represents the result of checking schemas with cancellation status
type CheckSchemasResult struct {
	Results     []TableResult
	IsCancelled bool
}

// isTimeoutError checks if the error is a timeout-related error
func isTimeoutError(err error) bool {
	if err == nil {
		return false
	}

	// Check for context timeout/cancellation
	if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
		return true
	}

	// Check for business timeout error
	var bizErr bizerr.BizError
	if errors.As(err, &bizErr) && bizErr.RawErrCode() == bizerr.ErrCodeTimeout {
		return true
	}

	return false
}

func (c *Checker) Run(ctx context.Context) (*PrecheckResult, bizerr.BizError) {
	files, isTruncated, err := c.store.ListLimited(ctx, nil)
	if err != nil {
		log.WithContextL(ctx).Error("failed to listFiles", zap.Error(err))
		return nil, bizerr.New(bizerr.ErrCodeUnknown, "failed to list files")
	}

	dbMap, bizErr := c.buildDbMetaMap(ctx, files)
	if bizErr != nil {
		return nil, bizErr
	}

	result, bizErr := c.checkSchemas(ctx, dbMap, isTruncated)
	if bizErr != nil {
		return nil, bizErr
	}

	// Determine the final status
	var status PrecheckStatus

	if result.IsCancelled {
		status = PrecheckStatusTimeout
	} else if isTruncated {
		status = PrecheckStatusTruncated
	} else {
		status = PrecheckStatusCompleted
	}

	return &PrecheckResult{
		TableResults: result.Results,
		Status:       status,
	}, nil
}

func (c *Checker) buildDbMetaMap(ctx context.Context, files []string) (map[string]*DbMeta, bizerr.BizError) {
	dbMap := make(map[string]*DbMeta)
	for _, file := range files {
		logger := log.WithContextL(ctx).With(zap.String("file", file))
		routeResult, err := c.fileRouter.Route(file)
		if err != nil {
			logger.Warn("failed to route file", zap.String("file", file), zap.Error(err))
			return nil, bizerr.New(bizerr.ErrCodeUnknown, "failed to route file")
		}
		if routeResult == nil {
			logger.Debug("skip route file", zap.String("file", file))
			continue
		}

		dbMeta := ensureDbMeta(dbMap, routeResult.Schema)

		switch routeResult.Type {
		case SchemaSchema:
			dbMeta.SchemaFile = routeResult
		case TableSchema:
			tableMeta := ensureTableMeta(dbMeta, routeResult.Name)
			tableMeta.SchemaFiles = append(tableMeta.SchemaFiles, *routeResult)
		default:
			tableMeta := ensureTableMeta(dbMeta, routeResult.Name)
			tableMeta.DataFiles = append(tableMeta.DataFiles, *routeResult)
		}
	}

	return dbMap, nil
}

func (c *Checker) checkSchemas(ctx context.Context, dbMap map[string]*DbMeta, isTruncated bool) (*CheckSchemasResult, bizerr.BizError) {
	var (
		results []TableResult
		mu      sync.Mutex
	)

	logger := log.WithContextL(ctx)
	g, ctx := errgroup.WithContext(ctx)
	g.SetLimit(concurrentCheckLimit)

	for schemaName, dbMeta := range dbMap {
		for tableName, tableMeta := range dbMeta.TableMap {
			schema := schemaName
			table := tableName
			meta := tableMeta

			g.Go(func() error {
				result, err := c.checkSingleTable(ctx, schema, table, meta, isTruncated)
				if err != nil {
					return err
				}
				if result != nil {
					mu.Lock()
					results = append(results, *result)
					mu.Unlock()
				}
				return nil
			})
		}
	}

	if err := g.Wait(); err != nil {
		if isTimeoutError(err) {
			logger.Warn("context canceled while waiting for table checks", zap.Error(err))
			return &CheckSchemasResult{
				Results:     results,
				IsCancelled: true,
			}, nil
		}
		logger.Error("failed to check schemas", zap.Error(err))
		return nil, bizerr.New(bizerr.ErrCodeUnknown, "%s", err.Error())
	}

	return &CheckSchemasResult{
		Results:     results,
		IsCancelled: false,
	}, nil
}

func (c *Checker) checkSingleTable(ctx context.Context, schema, table string, tableMeta *TableMeta, isTruncated bool) (*TableResult, bizerr.BizError) {
	logger := log.WithContextL(ctx).With(zap.String("schema", schema), zap.String("table", table))

	// Prepare the basic result with table info and data files
	result := c.prepareTableResult(schema, table, tableMeta)

	// Check table schema existence
	checkResult, err := c.schemaChecker.CheckSchema(ctx, schema, table)
	if err != nil {
		logger.Error("failed to check schema", zap.Error(err))
		return nil, err
	}

	// If using custom file mapping, we expect the table to exist and schema file validation is skipped.
	if c.useCustomFileMapping {
		if !checkResult.Exists {
			result.ErrorMsg = fmt.Sprintf("table %s.%s does not exist", schema, table)
		}
		// If table already exists and not empty, return error
		if checkResult.IsNotEmpty() {
			result.ErrorMsg = fmt.Sprintf("table %s.%s is not empty", schema, table)
		}
		return result, nil
	}

	// If table does not exist, we will use the schema file for import.
	if !checkResult.Exists {
		result.UseSchemaFile = true
	}

	// Handle schema file logic
	result, err = c.handleSchemaFile(ctx, schema, table, tableMeta, result, isTruncated)
	if err != nil {
		logger.Error("failed to handle schema file", zap.Error(err))
		return nil, err
	}

	// If table already exists and not empty, return error
	if checkResult.IsNotEmpty() {
		result.ErrorMsg = fmt.Sprintf("table %s.%s is not empty", schema, table)
	}

	return result, nil
}

// prepareTableResult creates a new TableResult with basic information
func (c *Checker) prepareTableResult(schema, table string, tableMeta *TableMeta) *TableResult {
	dataFiles := make([]string, 0, len(tableMeta.DataFiles))
	for _, file := range tableMeta.DataFiles {
		dataFiles = append(dataFiles, file.FilePath)
	}

	return &TableResult{
		DataFiles: dataFiles,
		Table: Table{
			Schema: schema,
			Name:   table,
		},
	}
}

// handleSchemaFile processes schema file logic based on different scenarios
func (c *Checker) handleSchemaFile(
	ctx context.Context,
	schema,
	table string,
	tableMeta *TableMeta,
	result *TableResult,
	isTruncated bool,
) (*TableResult, bizerr.BizError) {
	logger := log.WithContextL(ctx).With(zap.String("schema", schema), zap.String("table", table))
	// Case 1: Multiple schema files found (error case)
	if len(tableMeta.SchemaFiles) > 1 {
		result.ErrorMsg = fmt.Sprintf("table %s.%s has duplicated schema files: %v",
			schema, table, truncatePaths(tableMeta.SchemaFiles))
		return result, nil
	}

	// Case 2: No schema files found but response might be truncated
	if len(tableMeta.SchemaFiles) == 0 && isTruncated && result.UseSchemaFile {
		return c.retrieveSchemaFromStore(ctx, schema, table, result, logger)
	}

	// Case 3: Exactly one schema file found
	if len(tableMeta.SchemaFiles) == 1 {
		result.SchemaFile = tableMeta.SchemaFiles[0].FilePath
		return result, nil
	}

	// Case 4: No schema file found
	if len(tableMeta.SchemaFiles) == 0 && result.UseSchemaFile {
		result.ErrorMsg = fmt.Sprintf("table `%s`.`%s` schema not found", schema, table)
	}
	return result, nil
}

// retrieveSchemaFromStore attempts to retrieve the schema file from storage
func (c *Checker) retrieveSchemaFromStore(
	ctx context.Context,
	schema,
	table string,
	result *TableResult,
	logger *zap.Logger,
) (*TableResult, bizerr.BizError) {
	subDir := cloud_storage.GetBaseDir(result.DataFiles[0])
	logger.Info("retrieving schema file from store", zap.String("subDir", subDir))

	schemaFile, err := c.store.RetrieveSchemaFile(ctx, schema, table, subDir)
	if err != nil {
		logger.Error("failed to retrieve schema file", zap.Error(err))
		return nil, err
	}

	if schemaFile != nil && *schemaFile != "" {
		result.SchemaFile = *schemaFile
		return result, nil
	}

	result.ErrorMsg = fmt.Sprintf("table `%s`.`%s` schema not found", schema, table)
	return result, nil
}

func truncatePaths(files []RouteResult) []string {
	result := make([]string, 0, displayLimit+1)
	for i, f := range files {
		if i >= displayLimit {
			result = append(result, "...")
			break
		}
		result = append(result, f.FilePath)
	}
	return result
}

func ensureDbMeta(dbMap map[string]*DbMeta, schema string) *DbMeta {
	dbMeta, ok := dbMap[schema]
	if !ok {
		dbMeta = &DbMeta{
			Name:     schema,
			TableMap: make(map[string]*TableMeta),
		}
		dbMap[schema] = dbMeta
	}
	return dbMeta
}

func ensureTableMeta(dbMeta *DbMeta, tableName string) *TableMeta {
	tableMeta, ok := dbMeta.TableMap[tableName]
	if !ok {
		tableMeta = &TableMeta{
			Name: tableName,
		}
		dbMeta.TableMap[tableName] = tableMeta
	}
	return tableMeta
}