package imports

import (
	"math"
	"math/rand"
	"sort"
	"time"

	"github.com/docker/go-units"
	"github.com/go-logr/logr"
	"github.com/pingcap/errors"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
)

const (
	estimatedImportSizePerSecond      = 500 * units.KiB
	estimatedPostProcessSizePerSecond = 500 * units.KiB
	maxPreProcessedPercent            = 50
)

func getTableStatusForTiDBBackend(tableName string, tblProgress *TableProgress) (*cloudv1.CurrentTableStatus, error) {
	tableImportPhase, err := getTableImportPhaseForTiDBBackend(tblProgress)
	if err != nil {
		return nil, errors.Annotatef(err, "Getting table import phase for the table '%s' failed", tableName)
	}
	totalSize := tblProgress.TotalSize
	processedSize := tblProgress.TotalWritten
	completedPercent := uint32(float64(processedSize) * 100 / float64(totalSize))
	if completedPercent > 100 {
		completedPercent = 100
	}
	var postProgressPercent uint32 = 0
	if completedPercent == 100 {
		postProgressPercent = 100
	}

	return &cloudv1.CurrentTableStatus{
		Name:                       tableName,
		Phase:                      tableImportPhase,
		Size:                       uint64(totalSize),
		ProcessedSize:              uint64(processedSize),
		CompletedPercent:           completedPercent,
		PostImportCompletedPercent: postProgressPercent,
		Message:                    tblProgress.Message,
	}, nil
}

func getTableImportPhaseForTiDBBackend(tblProgress *TableProgress) (cloudv1.TableImportPhaseType, error) {
	switch tblProgress.Status {
	case taskStatusNotStarted:
		return cloudv1.TableImportPhasePending, nil
	case taskStatusRunning:
		return cloudv1.TableImportPhaseImporting, nil
	case taskStatusCompleted:
		if len(tblProgress.Message) > 0 {
			return cloudv1.TableImportPhaseFailed, nil
		} else {
			return cloudv1.TableImportPhaseComplete, nil
		}
	default:
		return cloudv1.TableImportPhaseUnknown, errors.Errorf("unknown table progress: %v", tblProgress.Status)
	}
}

func calculateTableCompletePercentForLocalBackend(tblStatus *cloudv1.CurrentTableStatus, prevCollectTS, currentCollectTS int64) float64 {
	switch tblStatus.Phase {
	case cloudv1.TableImportPhaseUnknown, cloudv1.TableImportPhasePending:
		return 0.0
	case cloudv1.TableImportPhasePreProcessing:
		if tblStatus.Size > 0 {
			return float64(tblStatus.ProcessedSize*maxPreProcessedPercent) / float64(tblStatus.Size)
		} else {
			return float64(maxPreProcessedPercent)
		}
	case cloudv1.TableImportPhaseImporting, cloudv1.TableImportPhaseIndexBuilding:
		if tblStatus.CompletedPercent == 0 {
			return float64(maxPreProcessedPercent + 1)
		}
		percentDelta := 0.0
		if prevCollectTS > 0 && currentCollectTS > 0 && currentCollectTS > prevCollectTS && tblStatus.Size > 0 {
			percentDelta = float64(100-maxPreProcessedPercent) * // rest of the percent
				(float64(currentCollectTS-prevCollectTS) / // elapsed time since last collection
					float64(tblStatus.Size/estimatedImportSizePerSecond)) // estimated total time
		}
		result := float64(tblStatus.CompletedPercent) + percentDelta
		if result >= 100.0 {
			return 99.0
		}
		return result
	case cloudv1.TableImportPhaseFailed:
		// For failed state, set to a small value,
		// so that if there is an existing progress value,
		// it will be kept.
		return 0.0
	default:
		return 100.0
	}
}

func estimateTableImportPhaseForLocalBackend(overallPhase cloudv1.ImportConditionType, taskPhase cloudv1.ImportTaskPhaseType, tblProgress *TableProgress) cloudv1.TableImportPhaseType {
	if len(tblProgress.Message) > 0 {
		return cloudv1.TableImportPhaseFailed
	}
	switch tblProgress.Status {
	case taskStatusRunning:
		if tblProgress.TotalWritten >= tblProgress.TotalSize {
			return cloudv1.TableImportPhaseImporting
		} else {
			return cloudv1.TableImportPhasePreProcessing
		}
	case taskStatusCompleted:
		switch {
		case taskPhase == cloudv1.ImportTaskPhasePostProcessing:
			return cloudv1.TableImportPhasePostProcessing
		case overallPhase == cloudv1.ImportCompleted:
			return cloudv1.TableImportPhaseComplete
		default:
			return cloudv1.TableImportPhaseImported
		}
	default:
		return cloudv1.TableImportPhasePending
	}
}

func saveImportStatus(importStatus *cloudv1.ImportStatus, task *TaskProgress, lightningBackend string, logger logr.Logger) {
	collectTime := time.Now().Unix()
	prevCollectTime := importStatus.CollectTime

	// update phase
	if task.Status == taskStatusRunning {
		importStatus.Phase = cloudv1.ImportRunning
		importStatus.IsTaskCreated = true
	}

	if task.Status == taskStatusCompleted {
		if task.Message == "" {
			importStatus.Phase = cloudv1.ImportCompleted
		} else {
			importStatus.Phase = cloudv1.ImportFailed
		}
	}

	// collect some basic statistics
	var totalSize uint64
	var totalProcessedSize int64
	var completedTables int
	var pendingTables int
	totalTables := len(task.Tables)
	for _, tableInfo := range task.Tables {
		totalSize += uint64(tableInfo.TotalSize)     // sum of total size
		totalProcessedSize += tableInfo.TotalWritten // sum of processed size
		switch tableInfo.Status {
		case taskStatusNotStarted:
			pendingTables++
		case taskStatusCompleted:
			completedTables++
		}
	}

	// get the running phase
	if importStatus.Phase == cloudv1.ImportRunning {
		if totalTables > 0 {
			if totalTables == completedTables {
				importStatus.TaskPhase = cloudv1.ImportTaskPhasePostProcessing
			} else {
				importStatus.TaskPhase = cloudv1.ImportTaskPhaseImporting
			}
		} else {
			importStatus.TaskPhase = cloudv1.ImportTaskPhaseUnknown
		}
	} else {
		importStatus.TaskPhase = cloudv1.ImportTaskPhaseUnknown
	}

	var files uint32
	var weightedPercentSum float64 = 0.0
	tblDetailStatusMap := map[string]cloudv1.CurrentTableStatus{}
	for _, tblStatus := range importStatus.AllTableDetails {
		tblDetailStatusMap[tblStatus.Name] = tblStatus
	}
	if len(tblDetailStatusMap) == 0 {
		for _, tblStatus := range importStatus.CurrentTables {
			tblDetailStatusMap[tblStatus.Name] = tblStatus
		}
	}
	// construct current table status list and percent
	for name, tableInfo := range task.Tables {
		var (
			current cloudv1.CurrentTableStatus
			ok      bool
		)
		// TODO: use interface and implementations to replace if/else for tidb backend and local backend
		if lightningBackend == "tidb" {
			currentTblStatus, err := getTableStatusForTiDBBackend(name, tableInfo)
			if err != nil {
				logger.Error(err, "get table status for TiDB backend error")
				continue
			}
			current = *currentTblStatus
		} else {
			current, ok = tblDetailStatusMap[name]
			if !ok {
				current = cloudv1.CurrentTableStatus{
					Name:             name,
					Phase:            cloudv1.TableImportPhaseUnknown,
					Size:             uint64(tableInfo.TotalSize),
					ProcessedSize:    uint64(tableInfo.TotalWritten),
					CompletedPercent: 0,
					Message:          tableInfo.Message,
				}
			} else {
				current.Size = uint64(tableInfo.TotalSize)
				current.ProcessedSize = uint64(tableInfo.TotalWritten)
				current.Message = tableInfo.Message
			}

			// determine the current table phase
			curTblPhase := estimateTableImportPhaseForLocalBackend(importStatus.Phase, importStatus.TaskPhase, tableInfo)
			if curTblPhase.StepGte(current.Phase) {
				current.Phase = curTblPhase
			}

			// calculate the percent of a table
			var percent uint32 = randomRound(calculateTableCompletePercentForLocalBackend(&current, prevCollectTime, collectTime))
			if percent > current.CompletedPercent {
				current.CompletedPercent = percent
			}
		}
		tblDetailStatusMap[name] = current
		if totalSize > 0 {
			weightedPercentSum += float64(current.CompletedPercent) * float64(current.Size) / float64(totalSize)
		} else {
			weightedPercentSum += float64(current.CompletedPercent) / float64(totalTables)
		}
	}
	importStatus.CollectTime = collectTime
	importStatus.TotalTables = uint32(totalTables)
	importStatus.CompletedTables = uint32(completedTables)
	importStatus.PendingTables = uint32(pendingTables)
	importStatus.TotalSize = totalSize
	importStatus.ProcessedDataSize = uint64(totalProcessedSize)
	importStatus.Message = importErrorMessageDigest(task.Message)

	currentTables := []cloudv1.CurrentTableStatus{}
	allTableStatuses := []cloudv1.CurrentTableStatus{}
	for _, tblStatus := range tblDetailStatusMap {
		allTableStatuses = append(allTableStatuses, tblStatus)
		switch tblStatus.Phase {
		case cloudv1.TableImportPhasePreProcessing,
			cloudv1.TableImportPhaseImporting,
			cloudv1.TableImportPhaseIndexBuilding:
			currentTables = append(currentTables, tblStatus)
		}
	}
	sort.Slice(currentTables, func(i, j int) bool { return currentTables[i].Name < currentTables[j].Name })
	importStatus.CurrentTables = currentTables
	if len(tblDetailStatusMap) < 1000 {
		sort.Slice(allTableStatuses, func(i, j int) bool { return allTableStatuses[i].Name < allTableStatuses[j].Name })
		importStatus.AllTableDetails = allTableStatuses
	} else {
		importStatus.AllTableDetails = nil
	}

	// set overall percent
	switch {
	case importStatus.TaskPhase == cloudv1.ImportTaskPhasePostProcessing:
		importStatus.CompletedPercent = 100
	case importStatus.Phase == cloudv1.ImportCompleted:
		importStatus.CompletedPercent = 100
	default:
		if uint32(weightedPercentSum) > importStatus.CompletedPercent {
			importStatus.CompletedPercent = uint32(weightedPercentSum)
		}
	}

	// get the overall post-processing percent
	postProcessProgress := randomRound(estimateOverallPostProcessingPercent(importStatus, prevCollectTime, collectTime))
	if postProcessProgress > importStatus.PostImportCompletedPercent {
		importStatus.PostImportCompletedPercent = postProcessProgress
	}
	if importStatus.TotalFiles == nil || *importStatus.TotalFiles < files {
		importStatus.TotalFiles = &files
	}
}

func estimateOverallPostProcessingPercent(importStatus *cloudv1.ImportStatus, prevCollectTS, currentCollectTS int64) float64 {
	switch {
	case importStatus.Phase == cloudv1.ImportCompleted:
		return 100.0
	case importStatus.Phase == cloudv1.ImportRunning && importStatus.TaskPhase == cloudv1.ImportTaskPhasePostProcessing:
		if importStatus.PostImportCompletedPercent == 0 {
			return 1.0
		}
		percentDelta := 0.0
		if prevCollectTS > 0 && currentCollectTS > 0 && currentCollectTS > prevCollectTS && importStatus.TotalSize > 0 {
			percentDelta = float64(100) *
				(float64(currentCollectTS-prevCollectTS) / // elapsed time since last collection
					float64(importStatus.TotalSize/estimatedPostProcessSizePerSecond)) // estimated total time
		}
		result := float64(importStatus.PostImportCompletedPercent) + percentDelta
		if result >= 100.0 {
			return 99.0
		}
		return result
	default:
		// set to a small value,
		// so that if there is an existing progress value,
		// it will be kept.
		return 0.0
	}
}

// assumption: the float number is greater than 0
func randomRound(num float64) uint32 {
	// Calculate the fractional part of the float64
	fraction := num - math.Floor(num)

	// Generate a random number between 0 and 1
	randomNum := rand.Float64()

	// Calculate the probability for choosing the ceiling integer
	ceilingProbability := 1.0 - fraction

	// Return the floor or ceiling integer based on the random number
	if randomNum <= ceilingProbability {
		return uint32(math.Floor(num))
	} else {
		return uint32(math.Ceil(num))
	}
}