package pimport

import (
	"testing"

	"github.com/docker/go-units"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/dataflow-service/idl/pbgen/proto/dataflow"
)

func TestGetOverallImportProgress(t *testing.T) {
	type testCase struct {
		CaseName               string
		IsInPreparation        bool
		ImportPercent          uint32
		PostImportPercent      uint32
		MaxPreparationPercent  float64
		ElapsedSeconds         uint32
		DataSize               uint64
		EstimatedMaxDataSize   uint64
		ExpectedOverallPercent int
	}
	for _, tc := range []testCase{
		{
			CaseName:               "preparing_local_file",
			IsInPreparation:        true,
			ImportPercent:          0,
			MaxPreparationPercent:  40,
			ElapsedSeconds:         5,
			DataSize:               0,
			EstimatedMaxDataSize:   50 * units.MiB,
			ExpectedOverallPercent: int(5 * 40 / 45),
		},
		{
			CaseName:               "preparing_overtime_local_file",
			IsInPreparation:        true,
			ImportPercent:          0,
			MaxPreparationPercent:  40,
			ElapsedSeconds:         45,
			DataSize:               0,
			EstimatedMaxDataSize:   50 * units.MiB,
			ExpectedOverallPercent: 40,
		},
		{
			CaseName:               "importing_just_started_local_file",
			IsInPreparation:        false,
			ImportPercent:          0,
			MaxPreparationPercent:  40,
			ElapsedSeconds:         50,
			DataSize:               25 * units.MiB,
			EstimatedMaxDataSize:   50 * units.MiB,
			ExpectedOverallPercent: int(50 * 100 / (expectedPreparationSeconds + ((25 * units.MiB) / (350 * units.KiB)))),
		},
		{
			CaseName:               "importing_local_file",
			IsInPreparation:        false,
			ImportPercent:          60,
			MaxPreparationPercent:  40,
			ElapsedSeconds:         60,
			DataSize:               25 * units.MiB,
			EstimatedMaxDataSize:   50 * units.MiB,
			ExpectedOverallPercent: int(40 + float64(100-40)/100*float64(60)),
		},
		{
			CaseName:               "post_processing_local_file",
			IsInPreparation:        false,
			ImportPercent:          60,
			PostImportPercent:      10,
			MaxPreparationPercent:  40,
			ElapsedSeconds:         70,
			DataSize:               25 * units.MiB,
			EstimatedMaxDataSize:   50 * units.MiB,
			ExpectedOverallPercent: 40 + 60*0.6,
		},
		{
			CaseName:               "complete_local_file",
			IsInPreparation:        false,
			ImportPercent:          100,
			MaxPreparationPercent:  40,
			ElapsedSeconds:         80,
			DataSize:               25 * units.MiB,
			EstimatedMaxDataSize:   50 * units.MiB,
			ExpectedOverallPercent: 100,
		},
		{
			CaseName:               "preparing_dedicated",
			IsInPreparation:        true,
			ImportPercent:          0,
			MaxPreparationPercent:  10,
			ElapsedSeconds:         5,
			DataSize:               0,
			EstimatedMaxDataSize:   1 * units.TiB,
			ExpectedOverallPercent: 1,
		},
		{
			CaseName:               "preparing_overtime_dedicated",
			IsInPreparation:        true,
			ImportPercent:          0,
			MaxPreparationPercent:  10,
			ElapsedSeconds:         45,
			DataSize:               0,
			EstimatedMaxDataSize:   1 * units.TiB,
			ExpectedOverallPercent: 10,
		},
		{
			CaseName:               "importing_just_started_dedicated",
			IsInPreparation:        false,
			ImportPercent:          0,
			MaxPreparationPercent:  10,
			ElapsedSeconds:         50,
			DataSize:               50 * units.GiB,
			EstimatedMaxDataSize:   1 * units.TiB,
			ExpectedOverallPercent: 10,
		},
		{
			CaseName:               "importing_dedicated",
			IsInPreparation:        false,
			ImportPercent:          60,
			MaxPreparationPercent:  10,
			ElapsedSeconds:         60,
			DataSize:               50 * units.GiB,
			EstimatedMaxDataSize:   1 * units.TiB,
			ExpectedOverallPercent: 10 + 60*0.9,
		},
		{
			CaseName:               "post_processing_dedicated",
			IsInPreparation:        false,
			ImportPercent:          60,
			PostImportPercent:      10,
			MaxPreparationPercent:  10,
			ElapsedSeconds:         70,
			DataSize:               50 * units.GiB,
			EstimatedMaxDataSize:   1 * units.TiB,
			ExpectedOverallPercent: 10 + 60*0.9,
		},
		{
			CaseName:               "complete_dedicated",
			IsInPreparation:        false,
			ImportPercent:          100,
			MaxPreparationPercent:  10,
			ElapsedSeconds:         80,
			DataSize:               50 * units.GiB,
			EstimatedMaxDataSize:   1 * units.TiB,
			ExpectedOverallPercent: 100,
		},
	} {
		overallPercent := getOverallImportProgress(
			tc.IsInPreparation,
			&dataflow.ImportGetResp{
				CompletedPercent:           tc.ImportPercent,
				PostImportCompletedPercent: tc.PostImportPercent,
				ElapsedTimeSeconds:         tc.ElapsedSeconds,
				TotalSize:                  tc.DataSize,
			},
			tc.MaxPreparationPercent,
			tc.EstimatedMaxDataSize,
		)
		require.Equalf(t, tc.ExpectedOverallPercent, int(overallPercent), "test case: %s", tc.CaseName)
	}
}

func TestEstimateImportProgressByElapsedTime(t *testing.T) {
	type testCase struct {
		CaseName               string
		ElapsedSeconds         uint32
		DataSize               uint64
		ExpectedOverallPercent int
	}
	for _, tc := range []testCase{
		{
			CaseName:               "just started",
			ElapsedSeconds:         5,
			DataSize:               50 * units.MiB,
			ExpectedOverallPercent: int(5 * 100 / (expectedPreparationSeconds + ((50 * units.MiB) / (350 * units.KiB)))),
		},
		{
			CaseName:               "prepare overtime",
			ElapsedSeconds:         48,
			DataSize:               50 * units.MiB,
			ExpectedOverallPercent: int(48 * 100 / (expectedPreparationSeconds + ((50 * units.MiB) / (350 * units.KiB)))),
		},
		{
			CaseName:               "importing small data",
			ElapsedSeconds:         47,
			DataSize:               100 * units.KiB,
			ExpectedOverallPercent: 99,
		},
		{
			CaseName:               "importing big data",
			ElapsedSeconds:         55,
			DataSize:               25 * units.MiB,
			ExpectedOverallPercent: int(55 * 100 / (expectedPreparationSeconds + ((25 * units.MiB) / (350 * units.KiB)))),
		},
	} {
		estimatedProgress := estimateImportProgressByElapsedTime(tc.ElapsedSeconds, tc.DataSize)
		require.Equalf(t, tc.ExpectedOverallPercent, int(estimatedProgress), "test case: %s", tc.CaseName)
	}
}

func TestGetRealValidationItems(t *testing.T) {
	type testCase struct {
		OriginalItems []dataflow.ImportValidationItem
		ExpectedItems []dataflow.ImportValidationItem
	}
	for _, tc := range []testCase{
		{
			OriginalItems: []dataflow.ImportValidationItem{
				dataflow.ImportValidationItem_SOURCE_ACCESS_CHECK,
			},
			ExpectedItems: []dataflow.ImportValidationItem{
				dataflow.ImportValidationItem_SOURCE_ACCESS_CHECK,
			},
		},
		{
			OriginalItems: []dataflow.ImportValidationItem{
				dataflow.ImportValidationItem_IMPORT_PRECHECK,
			},
			ExpectedItems: []dataflow.ImportValidationItem{
				dataflow.ImportValidationItem_IMPORT_PRECHECK,
			},
		},
		{
			OriginalItems: []dataflow.ImportValidationItem{
				dataflow.ImportValidationItem_IMPORT_PRECHECK,
				dataflow.ImportValidationItem_UNKNOWN,
				dataflow.ImportValidationItem(999),
			},
			ExpectedItems: []dataflow.ImportValidationItem{
				dataflow.ImportValidationItem_IMPORT_PRECHECK,
			},
		},
	} {
		actualItems := getRealValidationItems(tc.OriginalItems)
		require.Equal(t, tc.ExpectedItems, actualItems)
	}
}