package imports

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

	"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}"
	tblCkpt01JsonStr := "{\"Status\":30,\"AllocBase\":0,\"Engines\":{\"-1\":{\"Status\":30,\"Chunks\":[]},\"0\":{\"Status\":90,\"Chunks\":[{\"Key\":{\"Path\":\"edstats.footnote.csv\",\"Offset\":0},\"FileMeta\":{\"Path\":\"edstats.footnote.csv\",\"Type\":4,\"Compression\":0,\"SortKey\":\"\",\"FileSize\":33060496},\"ColumnPermutation\":[],\"Chunk\":{\"Offset\":0,\"EndOffset\":33060496,\"PrevRowIDMax\":0,\"RowIDMax\":8265124,\"Columns\":null},\"Checksum\":{\"checksum\":0,\"size\":0,\"kvs\":0},\"Timestamp\":1660457473}]}},\"TableID\":0,\"Checksum\":{\"checksum\":0,\"size\":0,\"kvs\":0}}"
	var task TaskProgress
	s.Require().NoError(json.Unmarshal([]byte(taskProgressRespStr), &task))
	var tblCkpt01 TableCheckpoint
	s.Require().NoError(json.Unmarshal([]byte(tblCkpt01JsonStr), &tblCkpt01))
	s.T().Logf("parsed table checkpoint: %+v\n", tblCkpt01)
	tblCkptMap := map[string]*TableCheckpoint{}
	tblCkptMap["`edstats`.`footnote`"] = &tblCkpt01
	var importStatus cloudv1.ImportStatus
	// to simulate `footnote` is in pre-processing
	task.Tables["`edstats`.`footnote`"].TotalWritten = 0
	saveImportStatus(&importStatus, &task, tblCkptMap)
	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)
	}

	time.Sleep(2 * time.Second)
	saveImportStatus(&importStatus, &task, tblCkptMap)
	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)
		s.Require().Equal(int64(2), tblStatus.TimeRecord.PreProcessDuration)
	}

	importStatusJson, err := json.MarshalIndent(&importStatus, "", "  ")
	s.Require().NoError(err)
	s.T().Logf("import status: %+v\n", string(importStatusJson))
	time.Sleep(2 * time.Second)
	task.Tables["`edstats`.`footnote`"].TotalWritten = 33060496
	saveImportStatus(&importStatus, &task, tblCkptMap)
	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(uint32(maxPreProcessedPortion+maxImportDataPortion)*100/(maxPreProcessedPortion+maxImportDataPortion+maxBuildIndexPortion), tblStatus.CompletedPercent)
		s.Require().Equal(int64(2), tblStatus.TimeRecord.PreProcessDuration)
		s.Require().Equal(int64(2), tblStatus.TimeRecord.ImportDuration)
	}

	importStatusJson, err = json.MarshalIndent(&importStatus, "", "  ")
	s.Require().NoError(err)
	s.T().Logf("import status 2: %+v\n", string(importStatusJson))
}

func (s *importProgressSuite) TestCalculateTableCompletePercent() {
	tblStatus := &cloudv1.CurrentTableStatus{
		Name:          "`edstats`.`footnote`",
		Phase:         cloudv1.TableImportPhaseImporting,
		Size:          33060496,
		ProcessedSize: 33060496,
		TimeRecord: cloudv1.TableImportTimeRecord{
			ImportStartTime:      1000,
			PostProcessStartTime: 0,
			PreProcessDuration:   60,
			ImportDuration:       10,
			BuildIndexDuration:   0,
			PostProcessDuration:  0,
			AnalyzeDuration:      0,
		},
	}
	percent := calculateTableCompletePercent(tblStatus)
	s.Require().Equal(int(maxPreProcessedPortion+10*maxPreProcessedPortion/60)*100/(maxPreProcessedPortion+maxImportDataPortion+maxBuildIndexPortion), int(percent))
}

func (s *importProgressSuite) TestUpdateTimeRecord() {
	curTime := int64(1000)
	timeRecord := &cloudv1.TableImportTimeRecord{}
	updateTimeRecord(timeRecord, curTime, cloudv1.TableImportPhasePreProcessing)
	s.Require().Equal(int64(1000), timeRecord.ImportStartTime)

	curTime += 10
	updateTimeRecord(timeRecord, curTime, cloudv1.TableImportPhasePreProcessing)
	s.Require().Equal(int64(10), timeRecord.PreProcessDuration)
	curTime += 10
	updateTimeRecord(timeRecord, curTime, cloudv1.TableImportPhasePreProcessing)
	s.Require().Equal(int64(20), timeRecord.PreProcessDuration)

	curTime += 10
	updateTimeRecord(timeRecord, curTime, cloudv1.TableImportPhaseImporting)
	s.Require().Equal(int64(10), timeRecord.ImportDuration)
	curTime += 10
	updateTimeRecord(timeRecord, curTime, cloudv1.TableImportPhaseImporting)
	s.Require().Equal(int64(20), timeRecord.ImportDuration)

	curTime += 10
	updateTimeRecord(timeRecord, curTime, cloudv1.TableImportPhaseImported)
	s.Require().Equal(int64(0), timeRecord.BuildIndexDuration)
	curTime += 10
	updateTimeRecord(timeRecord, curTime, cloudv1.TableImportPhaseImported)
	s.Require().Equal(int64(0), timeRecord.BuildIndexDuration)

	curTime += 10
	updateTimeRecord(timeRecord, curTime, cloudv1.TableImportPhasePostProcessing)
	s.Require().Equal(curTime, timeRecord.PostProcessStartTime)
	s.Require().Equal(int64(0), timeRecord.PostProcessDuration)
	curTime += 10
	updateTimeRecord(timeRecord, curTime, cloudv1.TableImportPhasePostProcessing)
	s.Require().Equal(int64(10), timeRecord.PostProcessDuration)

	curTime += 10
	updateTimeRecord(timeRecord, curTime, cloudv1.TableImportPhaseAnalyzing)
	s.Require().Equal(int64(10), timeRecord.AnalyzeDuration)
	curTime += 10
	updateTimeRecord(timeRecord, curTime, cloudv1.TableImportPhaseAnalyzing)
	s.Require().Equal(int64(20), timeRecord.AnalyzeDuration)

	curTime += 10
	updateTimeRecord(timeRecord, curTime, cloudv1.TableImportPhaseComplete)
	s.Require().Equal(int64(20), timeRecord.AnalyzeDuration)
	s.T().Logf("after updating time record: %+v\n", timeRecord)
}