package imports

import (
	"encoding/json"
	"fmt"
	"testing"
	"time"

	"github.com/docker/go-units"
	"github.com/go-logr/logr"
	"github.com/stretchr/testify/suite"

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

type importProgressSuite struct {
	suite.Suite
}

func TestImportProgressSuite(t *testing.T) {
	suite.Run(t, new(importProgressSuite))
}

func (s *importProgressSuite) TestGenImportStatus() {
	taskProgressRespStr := "{\"t\":{\"`edstats`.`country`\":{\"w\":125660,\"z\":125660,\"s\":2},\"`edstats`.`countryseries`\":{\"w\":44906,\"z\":44906,\"s\":2},\"`edstats`.`data`\":{\"w\":0,\"z\":215136009,\"s\":0},\"`edstats`.`footnote`\":{\"w\":33060496,\"z\":33060496,\"s\":1},\"`edstats`.`series`\":{\"w\":3568676,\"z\":3568676,\"s\":1}},\"s\":1}"
	var task TaskProgress
	s.Require().NoError(json.Unmarshal([]byte(taskProgressRespStr), &task))
	var importStatus cloudv1.ImportStatus
	// to simulate `footnote` is in pre-processing
	task.Tables["`edstats`.`footnote`"].TotalWritten = 0
	saveImportStatus(&importStatus, &task, "local", logr.Discard())
	for _, tblStatus := range importStatus.CurrentTables {
		if tblStatus.Name != "`edstats`.`footnote`" {
			continue
		}
		s.Require().Equal(uint64(0), tblStatus.ProcessedSize)
		s.Require().Equal(cloudv1.TableImportPhasePreProcessing, tblStatus.Phase)
	}

	saveImportStatus(&importStatus, &task, "local", logr.Discard())
	for _, tblStatus := range importStatus.CurrentTables {
		if tblStatus.Name != "`edstats`.`footnote`" {
			continue
		}
		s.Require().Equal(uint64(0), tblStatus.ProcessedSize)
		s.Require().Equal(cloudv1.TableImportPhasePreProcessing, tblStatus.Phase)
		s.Require().Equal(uint32(0), tblStatus.CompletedPercent)
	}

	importStatusJson, err := json.MarshalIndent(&importStatus, "", "  ")
	s.Require().NoError(err)
	s.T().Logf("import status: %+v\n", string(importStatusJson))
	task.Tables["`edstats`.`footnote`"].TotalWritten = 33060496
	saveImportStatus(&importStatus, &task, "local", logr.Discard())
	importStatusJson, err = json.MarshalIndent(&importStatus, "", "  ")
	s.Require().NoError(err)
	s.T().Logf("import status 2: %+v\n", string(importStatusJson))
	for _, tblStatus := range importStatus.CurrentTables {
		if tblStatus.Name != "`edstats`.`footnote`" {
			continue
		}
		s.Require().Equal(uint64(33060496), tblStatus.ProcessedSize)
		s.Require().Equal(cloudv1.TableImportPhaseImporting, tblStatus.Phase)
		s.Require().Equal(51, int(tblStatus.CompletedPercent))
	}

	// simulate the importing phase
	time.Sleep(2 * time.Second)
	saveImportStatus(&importStatus, &task, "local", logr.Discard())
	importStatusJson, err = json.MarshalIndent(&importStatus, "", "  ")
	s.Require().NoError(err)
	s.T().Logf("import status 3: %+v\n", string(importStatusJson))
	for _, tblStatus := range importStatus.CurrentTables {
		if tblStatus.Name != "`edstats`.`footnote`" {
			continue
		}
		s.Require().Equal(cloudv1.TableImportPhaseImporting, tblStatus.Phase)
		// 51 + randomRound((100 -50) * 2 / (33060496 / 500KiB)) == 52 or 53
		s.Require().Contains([]int{52, 53}, int(tblStatus.CompletedPercent))
	}

	// simulate the importing phase 2
	time.Sleep(4 * time.Second)
	saveImportStatus(&importStatus, &task, "local", logr.Discard())
	importStatusJson, err = json.MarshalIndent(&importStatus, "", "  ")
	s.Require().NoError(err)
	s.T().Logf("import status 4: %+v\n", string(importStatusJson))
	for _, tblStatus := range importStatus.CurrentTables {
		if tblStatus.Name != "`edstats`.`footnote`" {
			continue
		}
		s.Require().Equal(cloudv1.TableImportPhaseImporting, tblStatus.Phase)
		// 52/53 + randomRound((100 -50) * 4 / (33060496 / 500KiB)) == 55 ~ 57
		s.Require().Contains([]int{55, 56, 57}, int(tblStatus.CompletedPercent))
	}

	// simulate the table complete
	task.Tables["`edstats`.`footnote`"].Status = taskStatusCompleted
	saveImportStatus(&importStatus, &task, "local", logr.Discard())
	var (
		prevCompletedPercent            uint32
		prevPostProcessCompletedPercent uint32
	)
	for _, tblStatus := range importStatus.CurrentTables {
		if tblStatus.Name != "`edstats`.`footnote`" {
			continue
		}
		s.Require().Equal(cloudv1.TableImportPhaseImported, tblStatus.Phase)
		s.Require().Equal(100, tblStatus.CompletedPercent)
		prevCompletedPercent = tblStatus.CompletedPercent
		prevPostProcessCompletedPercent = tblStatus.PostImportCompletedPercent
	}

	// simulate some other table failed, causing the task failed.
	task.Status = taskStatusCompleted
	task.Message = "some other table failed"
	saveImportStatus(&importStatus, &task, "local", logr.Discard())
	s.Require().Equal(cloudv1.ImportFailed, importStatus.Phase)
	for _, tblStatus := range importStatus.CurrentTables {
		if tblStatus.Name != "`edstats`.`footnote`" {
			continue
		}
		s.Require().Equal(cloudv1.TableImportPhaseImported, tblStatus.Phase)
		s.Require().Equal(prevCompletedPercent, tblStatus.CompletedPercent)
		s.Require().Equal(prevPostProcessCompletedPercent, tblStatus.PostImportCompletedPercent)
	}
}

func (s *importProgressSuite) TestSaveHugeTablesImportStatus() {
	var importStatus cloudv1.ImportStatus
	totalFiles := uint32(0)
	// just start the task
	task := TaskProgress{
		Tables: make(map[string]*TableProgress),
		Status: taskStatusRunning,
	}
	expectedImportStatus := cloudv1.ImportStatus{
		Phase:                      cloudv1.ImportRunning,
		IsTaskCreated:              true,
		TaskPhase:                  cloudv1.ImportTaskPhaseUnknown,
		TotalTables:                0,
		CompletedTables:            0,
		PendingTables:              0,
		TotalSize:                  0,
		ProcessedDataSize:          0,
		CurrentTables:              []cloudv1.CurrentTableStatus{},
		AllTableDetails:            []cloudv1.CurrentTableStatus{},
		CompletedPercent:           0,
		PostImportCompletedPercent: 0,
		TotalFiles:                 &totalFiles,
	}
	saveImportStatus(&importStatus, &task, "local", logr.Discard())
	importStatusJson, err := json.MarshalIndent(&importStatus, "", "  ")
	s.Require().NoError(err)
	s.T().Logf("import status: %+v\n", string(importStatusJson))
	expectedImportStatus.CollectTime = importStatus.CollectTime
	s.Require().Equal(expectedImportStatus, importStatus)

	// simulate the importing
	time.Sleep(1 * time.Second)
	expectedImportStatus = cloudv1.ImportStatus{
		Phase:                      cloudv1.ImportRunning,
		IsTaskCreated:              true,
		TaskPhase:                  cloudv1.ImportTaskPhaseImporting,
		CollectTime:                importStatus.CollectTime,
		TotalTables:                5000,
		CompletedTables:            200,
		PendingTables:              5000 - 232,
		TotalSize:                  5000 * 100 * units.KiB,
		ProcessedDataSize:          216*100*units.KiB + 16*50*units.KiB,
		CurrentTables:              nil,
		AllTableDetails:            nil,
		CompletedPercent:           4, // 200 * 100 * 100 / (5000 * 100) + 16 * 50 * 100 / (5000 * 100) + 16 * 25 * 100 / (5000 * 100) = 4.24
		PostImportCompletedPercent: 0,
		TotalFiles:                 &totalFiles,
	}
	curTables := make([]cloudv1.CurrentTableStatus, 0)
	for i := 0; i < 5000; i++ {
		tblName := fmt.Sprintf("`testdb`.`tbl_%04d`", i)
		tblProgress := &TableProgress{
			TotalSize: 100 * units.KiB,
			Status:    taskStatusNotStarted,
		}
		switch {
		case i < 200:
			tblProgress.TotalWritten = 100 * units.KiB
			tblProgress.Status = taskStatusCompleted
		case i < 216:
			tblProgress.TotalWritten = 100 * units.KiB
			tblProgress.Status = taskStatusRunning
			curTables = append(curTables, cloudv1.CurrentTableStatus{
				Name:                       tblName,
				Phase:                      cloudv1.TableImportPhaseImporting,
				Size:                       100 * units.KiB,
				ProcessedSize:              100 * units.KiB,
				CompletedPercent:           51,
				PostImportCompletedPercent: 0,
			})
		case i < 232:
			tblProgress.TotalWritten = 50 * units.KiB
			tblProgress.Status = taskStatusRunning
			curTables = append(curTables, cloudv1.CurrentTableStatus{
				Name:                       tblName,
				Phase:                      cloudv1.TableImportPhasePreProcessing,
				Size:                       100 * units.KiB,
				ProcessedSize:              50 * units.KiB,
				CompletedPercent:           25,
				PostImportCompletedPercent: 0,
			})
		}
		task.Tables[tblName] = tblProgress
		expectedImportStatus.CurrentTables = curTables
	}
	saveImportStatus(&importStatus, &task, "local", logr.Discard())
	importStatusJson, err = json.MarshalIndent(&importStatus, "", "  ")
	s.Require().NoError(err)
	s.T().Logf("import status: %+v\n", string(importStatusJson))
	expectedImportStatus.CollectTime = importStatus.CollectTime
	s.Require().Equal(expectedImportStatus, importStatus)

	// simulate the post-processing phase
	time.Sleep(1 * time.Second)
	fmt.Println("sleep after import status")
	task = TaskProgress{
		Tables: make(map[string]*TableProgress),
		Status: taskStatusRunning,
	}
	for i := 0; i < 5000; i++ {
		tblName := fmt.Sprintf("`testdb`.`tbl_%04d`", i)
		tblProgress := &TableProgress{
			TotalWritten: 100 * units.KiB,
			TotalSize:    100 * units.KiB,
			Status:       taskStatusCompleted,
		}
		task.Tables[tblName] = tblProgress
	}
	saveImportStatus(&importStatus, &task, "local", logr.Discard())
	importStatusJson, err = json.MarshalIndent(&importStatus, "", "  ")
	s.Require().NoError(err)
	s.T().Logf("import status 2: %+v\n", string(importStatusJson))
	expectedImportStatus = cloudv1.ImportStatus{
		Phase:                      cloudv1.ImportRunning,
		IsTaskCreated:              true,
		TaskPhase:                  cloudv1.ImportTaskPhasePostProcessing,
		CollectTime:                importStatus.CollectTime,
		TotalTables:                5000,
		CompletedTables:            5000,
		PendingTables:              0,
		TotalSize:                  5000 * 100 * units.KiB,
		ProcessedDataSize:          5000 * 100 * units.KiB,
		CurrentTables:              []cloudv1.CurrentTableStatus{},
		AllTableDetails:            nil,
		CompletedPercent:           100,
		PostImportCompletedPercent: 1,
		TotalFiles:                 &totalFiles,
	}
	s.Require().Equal(expectedImportStatus, importStatus)

	time.Sleep(1 * time.Second)
	fmt.Println("sleep after import status 2")
	saveImportStatus(&importStatus, &task, "local", logr.Discard())
	postImportCompletedPercent := importStatus.PostImportCompletedPercent
	importStatusJson, err = json.MarshalIndent(&importStatus, "", "  ")
	s.Require().NoError(err)
	s.T().Logf("import status 3: %+v\n", string(importStatusJson))
	expectedImportStatus.CollectTime = importStatus.CollectTime
	importStatus.PostImportCompletedPercent = 1
	s.Require().Equal(expectedImportStatus, importStatus)
	s.Require().Contains([]uint32{1, 2}, postImportCompletedPercent)
	importStatus.PostImportCompletedPercent = postImportCompletedPercent
}

func (s *importProgressSuite) TestCalculateTableCompletePercent() {
	tblStatus := &cloudv1.CurrentTableStatus{
		Name:          "`edstats`.`footnote`",
		Phase:         cloudv1.TableImportPhaseImporting,
		Size:          50000 * units.KiB,
		ProcessedSize: 50000 * units.KiB,
	}
	prevTS := int64(0)
	currentTS := int64(1000)
	percent := calculateTableCompletePercentForLocalBackend(tblStatus, prevTS, currentTS)
	s.Require().Equal(51.0, percent)
	tblStatus.CompletedPercent = uint32(percent)

	prevTS = int64(1000)
	currentTS = int64(1010)
	percent = calculateTableCompletePercentForLocalBackend(tblStatus, prevTS, currentTS)
	// 51 + round((100 -50) * 10 / (50000KiB / 500KiB)) == 56
	s.Require().Equal(56.0, percent)
	tblStatus.CompletedPercent = uint32(percent)

	prevTS = int64(1010)
	currentTS = int64(1020)
	percent = calculateTableCompletePercentForLocalBackend(tblStatus, prevTS, currentTS)
	// 56 + round((100 -50) * 10 / (50000KiB / 500KiB)) == 61
	s.Require().Equal(61.0, percent)
	tblStatus.CompletedPercent = uint32(percent)

	tblStatus.Phase = cloudv1.TableImportPhaseComplete
	percent = calculateTableCompletePercentForLocalBackend(tblStatus, prevTS, currentTS)
	s.Require().Equal(100.0, percent)
	tblStatus.CompletedPercent = uint32(percent)

	tblStatus.Phase = cloudv1.TableImportPhaseFailed
	percent = calculateTableCompletePercentForLocalBackend(tblStatus, prevTS, currentTS)
	s.Require().Equal(0.0, percent)
}

func (s *importProgressSuite) TestEstimatePostProcessPercent() {
	tblStatus := &cloudv1.ImportStatus{
		Phase:     cloudv1.ImportRunning,
		TaskPhase: cloudv1.ImportTaskPhaseImporting,
		TotalSize: 500 * 1024 * 100,
	}
	prevTS := int64(0)
	currentTS := int64(1000)
	percent := estimateOverallPostProcessingPercent(tblStatus, prevTS, currentTS)
	s.Require().Equal(0.0, percent)
	tblStatus.PostImportCompletedPercent = uint32(percent)

	tblStatus.TaskPhase = cloudv1.ImportTaskPhasePostProcessing
	prevTS = int64(1000)
	currentTS = int64(1010)
	percent = estimateOverallPostProcessingPercent(tblStatus, prevTS, currentTS)
	s.Require().Equal(1.0, percent)
	tblStatus.PostImportCompletedPercent = uint32(percent)

	tblStatus.TaskPhase = cloudv1.ImportTaskPhasePostProcessing
	prevTS = int64(1010)
	currentTS = int64(1020)
	percent = estimateOverallPostProcessingPercent(tblStatus, prevTS, currentTS)
	// 1 + 100 * 10 / (500KiB * 100 / 500KiB) == 16.625
	s.Require().Equal(11.0, percent)
	tblStatus.PostImportCompletedPercent = uint32(percent)

	tblStatus.TaskPhase = cloudv1.ImportTaskPhasePostProcessing
	prevTS = int64(1020)
	currentTS = int64(1030)
	percent = estimateOverallPostProcessingPercent(tblStatus, prevTS, currentTS)
	// 11 + round((100 * 10 / (500KiB * 100 / 500KiB)) == 21
	s.Require().Equal(21.0, percent)
	tblStatus.PostImportCompletedPercent = uint32(percent)

	tblStatus.Phase = cloudv1.ImportCompleted
	percent = estimateOverallPostProcessingPercent(tblStatus, prevTS, currentTS)
	s.Require().Equal(100.0, percent)
	tblStatus.PostImportCompletedPercent = uint32(percent)

	tblStatus.Phase = cloudv1.ImportFailed
	percent = estimateOverallPostProcessingPercent(tblStatus, prevTS, currentTS)
	s.Require().Equal(0.0, percent)
}

func (s *importProgressSuite) TestGetTableStatusForTiDBBackend() {
	type testCase struct {
		CaseName            string
		TestTableName       string
		TaskProgressStr     string
		ExpectedTableStatus *cloudv1.CurrentTableStatus
	}
	for _, tc := range []testCase{
		{
			CaseName:      "importing",
			TestTableName: "`mytest`.`customer`",
			TaskProgressStr: fmt.Sprintf(`{
  "t": {
    "%s": {
      "w": 183890,
      "z": 52428233,
      "s": 1
    },
    "%s": {
      "w": 183890,
      "z": 52428233,
      "s": 1
    },
    "%s": {
      "w": 58,
      "z": 58,
      "s": 2
    }
  },
  "s": 1
}`, "`mytest`.`customer2`", "`mytest`.`customer`", "`mytest`.`testtbl`"),
			ExpectedTableStatus: &cloudv1.CurrentTableStatus{
				Name:                       "`mytest`.`customer`",
				Phase:                      cloudv1.TableImportPhaseImporting,
				Size:                       52428233,
				ProcessedSize:              183890,
				CompletedPercent:           uint32(183890 * 100 / 52428233),
				PostImportCompletedPercent: 0,
			},
		},
		{
			CaseName:      "importing2",
			TestTableName: "`mytest`.`customer`",
			TaskProgressStr: fmt.Sprintf(`{
  "t": {
    "%s": {
      "w": 1837832,
      "z": 52428233,
      "s": 1
    },
    "%s": {
      "w": 1837832,
      "z": 52428233,
      "s": 1
    },
    "%s": {
      "w": 58,
      "z": 58,
      "s": 2
    }
  },
  "s": 1
}`, "`mytest`.`customer2`", "`mytest`.`customer`", "`mytest`.`testtbl`"),
			ExpectedTableStatus: &cloudv1.CurrentTableStatus{
				Name:                       "`mytest`.`customer`",
				Phase:                      cloudv1.TableImportPhaseImporting,
				Size:                       52428233,
				ProcessedSize:              1837832,
				CompletedPercent:           uint32(1837832 * 100 / 52428233),
				PostImportCompletedPercent: 0,
			},
		},
		{
			CaseName:      "completed",
			TestTableName: "`mytest`.`customer`",
			TaskProgressStr: fmt.Sprintf(`{
  "t": {
    "%s": {
      "w": 1837832,
      "z": 52428233,
      "s": 1
    },
    "%s": {
      "w": 52428233,
      "z": 52428233,
      "s": 2
    },
    "%s": {
      "w": 58,
      "z": 58,
      "s": 2
    }
  },
  "s": 1
}`, "`mytest`.`customer2`", "`mytest`.`customer`", "`mytest`.`testtbl`"),
			ExpectedTableStatus: &cloudv1.CurrentTableStatus{
				Name:                       "`mytest`.`customer`",
				Phase:                      cloudv1.TableImportPhaseComplete,
				Size:                       52428233,
				ProcessedSize:              52428233,
				CompletedPercent:           100,
				PostImportCompletedPercent: 100,
			},
		},
		{
			CaseName:      "failed_halfway",
			TestTableName: "`mytest`.`customer`",
			TaskProgressStr: fmt.Sprintf(`{
  "t": {
    "%s": {
      "w": 1837832,
      "z": 52428233,
      "s": 1
    },
    "%s": {
      "w": 1837832,
      "z": 52428233,
      "s": 2,
      "m": "some error happen"
    },
    "%s": {
      "w": 58,
      "z": 58,
      "s": 2
    }
  },
  "s": 1
}`, "`mytest`.`customer2`", "`mytest`.`customer`", "`mytest`.`testtbl`"),
			ExpectedTableStatus: &cloudv1.CurrentTableStatus{
				Name:                       "`mytest`.`customer`",
				Phase:                      cloudv1.TableImportPhaseFailed,
				Size:                       52428233,
				ProcessedSize:              1837832,
				CompletedPercent:           uint32(1837832 * 100 / 52428233),
				PostImportCompletedPercent: 0,
				Message:                    "some error happen",
			},
		},
	} {
		var task TaskProgress
		s.Require().NoErrorf(json.Unmarshal([]byte(tc.TaskProgressStr), &task), "case name: %s", tc.CaseName)
		tblProgress, ok := task.Tables[tc.TestTableName]
		s.Require().True(ok, "case name: %s", tc.CaseName)
		s.Require().NotNil(tblProgress, "case name: %s", tc.CaseName)
		tblStatus, err := getTableStatusForTiDBBackend(tc.TestTableName, tblProgress)
		s.Require().NoError(err, "case name: %s", tc.CaseName)
		s.Require().NotNil(tblStatus, "case name: %s", tc.CaseName)
		s.Require().Equal(tc.ExpectedTableStatus, tblStatus, "case name: %s", tc.CaseName)
	}
}

func (s *importProgressSuite) TestRandomRound() {
	floorCount := 0
	ceilCount := 0
	for i := 0; i < 100; i++ {
		randInt := randomRound(3.2)
		s.T().Logf("random round result: %d", randInt)
		switch randInt {
		case 3:
			floorCount++
		case 4:
			ceilCount++
		default:
			s.FailNow("unexpected random round result")
		}
		time.Sleep(1 * time.Millisecond)
	}
	s.Require().GreaterOrEqual(floorCount, ceilCount)

	floorCount = 0
	ceilCount = 0
	for i := 0; i < 100; i++ {
		randInt := randomRound(0.8)
		s.T().Logf("random round result: %d", randInt)
		switch randInt {
		case 0:
			floorCount++
		case 1:
			ceilCount++
		default:
			s.FailNow("unexpected random round result")
		}
		time.Sleep(1 * time.Millisecond)
	}
	s.Require().GreaterOrEqual(ceilCount, floorCount)
}