package report

import (
	"context"
	"sort"
	"testing"

	"github.com/pingcap/log"
	"github.com/tidbcloud/serverless-metering/pkg/config"
	"github.com/tidbcloud/serverless-metering/pkg/repository"
	"github.com/tidbcloud/serverless-metering/pkg/store"
	"github.com/tj/assert"
)

func TestDefaultWorker(t *testing.T) {
	config.InitMockConfig()
	ctx := context.TODO()
	storage := store.NewMockStorage()
	repo := repository.NewMockClusterMetaRepo()
	worker, err := newDefaultMeteringWorker(storage, repo)
	assert.Nil(t, err)
	reports, err := worker.genBillingReports(ctx, log.L())
	sort.Slice(reports, func(i, j int) bool { return reports[i].CostID < reports[j].CostID })
	assert.Nil(t, err)
	assert.Equal(t, 11, len(reports))

	assert.Equal(t, "us-east-1/1001/load_data_ru/0", reports[0].CostID)
	assert.Equal(t, int64(19800), reports[0].Cost)

	assert.Equal(t, "us-east-1/1001/privatelink_egress/0", reports[1].CostID)
	assert.Equal(t, int64(7), reports[1].Cost)

	assert.Equal(t, "us-east-1/1001/public_egress/0", reports[2].CostID)
	assert.Equal(t, int64(29), reports[2].Cost)

	assert.Equal(t, "us-east-1/1001/tidb_ru/0", reports[3].CostID)
	assert.Equal(t, int64(79200), reports[3].Cost)

	assert.Equal(t, "us-east-1/1001/tiflash_ru/0", reports[4].CostID)
	assert.Equal(t, int64(29700), reports[4].Cost)

	assert.Equal(t, "us-east-1/1002/load_data_ru/0", reports[5].CostID)
	assert.Equal(t, int64(9900), reports[5].Cost)

	assert.Equal(t, "us-east-1/1002/privatelink_egress/0", reports[6].CostID)
	assert.Equal(t, int64(2), reports[6].Cost)

	assert.Equal(t, "us-east-1/1002/public_egress/0", reports[7].CostID)
	assert.Equal(t, int64(9), reports[7].Cost)

	assert.Equal(t, "us-east-1/1003/load_data_ru/0", reports[8].CostID)
	assert.Equal(t, int64(4900), reports[8].Cost)

	assert.Equal(t, "us-east-1/1003/tiflash_storage/0", reports[9].CostID)
	assert.Equal(t, int64((2000+7000)*5*6), reports[9].Cost)

	assert.Equal(t, "us-east-1/1003/tikv_storage/0", reports[10].CostID)
	assert.Equal(t, int64((2000+7000)*5*6), reports[10].Cost)
}

func TestStorageMediumWorker(t *testing.T) {
	config.InitMockConfig()
	ctx := context.TODO()
	storage := store.NewMockStorage()
	repo := repository.NewMockClusterMetaRepo()
	worker, err := newStorageMediumWorker(storage, repo)
	assert.Nil(t, err)
	reports, err := worker.genBillingReports(ctx, log.L())
	sort.Slice(reports, func(i, j int) bool { return reports[i].CostID < reports[j].CostID })
	assert.Nil(t, err)
	assert.Equal(t, 3, len(reports))

	assert.Equal(t, "us-east-1/1001/tiflash_storage/0", reports[0].CostID)
	assert.Equal(t, int64((60+210)*5*6*2), reports[0].Cost)

	assert.Equal(t, "us-east-1/1001/tikv_storage/0", reports[1].CostID)
	assert.Equal(t, int64((60+210)*5*6*2), reports[1].Cost)

	assert.Equal(t, "us-east-1/1002/tikv_storage/0", reports[2].CostID)
	assert.Equal(t, int64((70+245)*5*6*2), reports[2].Cost)
}

func TestWorkerNextTs(t *testing.T) {
	worker := meteringWorker{finishedTs: 1732259983, reportIntervalHour: 1}
	assert.Equal(t, int64(1732262400), worker.NextTs())
	worker = meteringWorker{finishedTs: 1732255200, reportIntervalHour: 1}
	assert.Equal(t, int64(1732258800), worker.NextTs())
	worker = meteringWorker{finishedTs: 1732255200, reportIntervalHour: 24}
	assert.Equal(t, int64(1732341600), worker.NextTs())
}

func TestSplitCost(t *testing.T) {
	worker := &meteringWorker{}

	t.Run("nil percent map", func(t *testing.T) {
		result := worker.splitCost(1000, nil)
		assert.Equal(t, 1, len(result))
		assert.Equal(t, int64(1000), result[""])
	})

	t.Run("empty percent map", func(t *testing.T) {
		result := worker.splitCost(1000, map[string]int64{})
		assert.Equal(t, 1, len(result))
		assert.Equal(t, int64(1000), result[""])
	})

	t.Run("single entry", func(t *testing.T) {
		result := worker.splitCost(1000, map[string]int64{"db1": 100})
		assert.Equal(t, 1, len(result))
		assert.Equal(t, int64(1000), result["db1"])
	})

	t.Run("equal split two entries", func(t *testing.T) {
		result := worker.splitCost(1000, map[string]int64{"db1": 50, "db2": 50})
		assert.Equal(t, 2, len(result))
		// Total should equal original cost
		total := result["db1"] + result["db2"]
		assert.Equal(t, int64(1000), total)
		// Each should be approximately half
		assert.Equal(t, int64(500), result["db1"])
		assert.Equal(t, int64(500), result["db2"])
	})

	t.Run("proportional split three entries", func(t *testing.T) {
		result := worker.splitCost(1000, map[string]int64{"db1": 50, "db2": 30, "db3": 20})
		assert.Equal(t, 3, len(result))
		// Total should equal original cost
		total := result["db1"] + result["db2"] + result["db3"]
		assert.Equal(t, int64(1000), total)
		// Check proportions (50%, 30%, 20%)
		assert.True(t, result["db1"] >= 490 && result["db1"] <= 510)
		assert.True(t, result["db2"] >= 290 && result["db2"] <= 310)
		assert.True(t, result["db3"] >= 190 && result["db3"] <= 210)
	})

	t.Run("unequal percentages", func(t *testing.T) {
		result := worker.splitCost(1000, map[string]int64{"db1": 7, "db2": 3})
		assert.Equal(t, 2, len(result))
		// Total should equal original cost
		total := result["db1"] + result["db2"]
		assert.Equal(t, int64(1000), total)
		// db1 should get 70%, db2 should get 30%
		assert.Equal(t, int64(700), result["db1"])
		assert.Equal(t, int64(300), result["db2"])
	})

	t.Run("remainder goes to last entry", func(t *testing.T) {
		// Use a cost that doesn't divide evenly
		result := worker.splitCost(100, map[string]int64{"db1": 1, "db2": 1, "db3": 1})
		assert.Equal(t, 3, len(result))
		// Total should equal original cost exactly
		total := result["db1"] + result["db2"] + result["db3"]
		assert.Equal(t, int64(100), total)
	})

	t.Run("large numbers", func(t *testing.T) {
		result := worker.splitCost(1000000, map[string]int64{"db1": 250, "db2": 750})
		assert.Equal(t, 2, len(result))
		total := result["db1"] + result["db2"]
		assert.Equal(t, int64(1000000), total)
		assert.Equal(t, int64(250000), result["db1"])
		assert.Equal(t, int64(750000), result["db2"])
	})

	t.Run("zero cost", func(t *testing.T) {
		result := worker.splitCost(0, map[string]int64{"db1": 50, "db2": 50})
		assert.Equal(t, 2, len(result))
		total := result["db1"] + result["db2"]
		assert.Equal(t, int64(0), total)
	})

	t.Run("cost with rounding", func(t *testing.T) {
		// 1000 split by percentages that don't divide evenly
		result := worker.splitCost(1000, map[string]int64{"db1": 3, "db2": 3, "db3": 3})
		assert.Equal(t, 3, len(result))
		// Total must still equal original cost (rounding handled by remainder)
		total := result["db1"] + result["db2"] + result["db3"]
		assert.Equal(t, int64(1000), total)
	})
}

func TestGenBillingReports(t *testing.T) {
	config.InitMockConfig()
	ctx := context.TODO()

	t.Run("normal cluster without special handling", func(t *testing.T) {
		storage := store.NewMockStorage()
		repo := repository.NewMockClusterMetaRepo()
		worker, err := newDefaultMeteringWorker(storage, repo)
		assert.Nil(t, err)

		reports, err := worker.genBillingReports(ctx, log.L())
		sort.Slice(reports, func(i, j int) bool { return reports[i].CostID < reports[j].CostID })
		assert.Nil(t, err)
		assert.Equal(t, 11, len(reports))

		// Verify detailed cost and costID for each report
		assert.Equal(t, "us-east-1/1001/load_data_ru/0", reports[0].CostID)
		assert.Equal(t, int64(19800), reports[0].Cost)
		assert.Equal(t, uint64(11), reports[0].TenantID)

		assert.Equal(t, "us-east-1/1001/privatelink_egress/0", reports[1].CostID)
		assert.Equal(t, int64(7), reports[1].Cost)
		assert.Equal(t, uint64(11), reports[1].TenantID)

		assert.Equal(t, "us-east-1/1001/public_egress/0", reports[2].CostID)
		assert.Equal(t, int64(29), reports[2].Cost)
		assert.Equal(t, uint64(11), reports[2].TenantID)

		assert.Equal(t, "us-east-1/1001/tidb_ru/0", reports[3].CostID)
		assert.Equal(t, int64(79200), reports[3].Cost)
		assert.Equal(t, uint64(11), reports[3].TenantID)

		assert.Equal(t, "us-east-1/1001/tiflash_ru/0", reports[4].CostID)
		assert.Equal(t, int64(29700), reports[4].Cost)
		assert.Equal(t, uint64(11), reports[4].TenantID)

		assert.Equal(t, "us-east-1/1002/load_data_ru/0", reports[5].CostID)
		assert.Equal(t, int64(9900), reports[5].Cost)
		assert.Equal(t, uint64(12), reports[5].TenantID)

		assert.Equal(t, "us-east-1/1002/privatelink_egress/0", reports[6].CostID)
		assert.Equal(t, int64(2), reports[6].Cost)
		assert.Equal(t, uint64(12), reports[6].TenantID)

		assert.Equal(t, "us-east-1/1002/public_egress/0", reports[7].CostID)
		assert.Equal(t, int64(9), reports[7].Cost)
		assert.Equal(t, uint64(12), reports[7].TenantID)

		assert.Equal(t, "us-east-1/1003/load_data_ru/0", reports[8].CostID)
		assert.Equal(t, int64(4900), reports[8].Cost)
		assert.Equal(t, uint64(13), reports[8].TenantID)

		assert.Equal(t, "us-east-1/1003/tiflash_storage/0", reports[9].CostID)
		assert.Equal(t, int64((2000+7000)*5*6), reports[9].Cost)
		assert.Equal(t, uint64(13), reports[9].TenantID)

		assert.Equal(t, "us-east-1/1003/tikv_storage/0", reports[10].CostID)
		assert.Equal(t, int64((2000+7000)*5*6), reports[10].Cost)
		assert.Equal(t, uint64(13), reports[10].TenantID)

		// Verify each report has required fields
		for _, report := range reports {
			assert.NotEmpty(t, report.CostID)
			assert.True(t, report.TenantID > 0)
			assert.NotEmpty(t, report.Sku)
			assert.True(t, report.Cost >= 0)
			assert.NotEmpty(t, report.Labels)
			assert.True(t, report.ShouldLog)
		}
	})

	t.Run("ru_version=2 uses ruv2 only", func(t *testing.T) {
		storage := store.NewMockStorage()
		repo := repository.NewMockClusterMetaRepoWithRUVersions(map[string]string{"1": "2"})
		worker, err := newDefaultMeteringWorker(storage, repo)
		assert.Nil(t, err)

		reports, err := worker.genBillingReports(ctx, log.L())
		sort.Slice(reports, func(i, j int) bool { return reports[i].CostID < reports[j].CostID })
		assert.Nil(t, err)
		assert.Equal(t, 11, len(reports))

		foundTidbRUV2 := false
		foundTiflashRUV2 := false
		for _, report := range reports {
			switch report.CostID {
			case "us-east-1/1001/tidb_ru_v2/0":
				foundTidbRUV2 = true
				assert.Equal(t, int64(9900), report.Cost)
				assert.Equal(t, "2", report.Labels[BillingLabelRUVersion])
			case "us-east-1/1001/tiflash_ru_v2/0":
				foundTiflashRUV2 = true
				assert.Equal(t, int64(9900), report.Cost)
				assert.Equal(t, "2", report.Labels[BillingLabelRUVersion])
			}
		}

		assert.True(t, foundTidbRUV2, "Expected ruv2 metered tidb_ru_v2 report")
		assert.True(t, foundTiflashRUV2, "Expected ruv2 metered tiflash_ru_v2 report")
	})

	t.Run("db-level RU splitting for enabled tenants", func(t *testing.T) {
		storage := store.NewMockStorage()
		repo := repository.NewMockClusterMetaRepo()
		worker, err := newDefaultMeteringWorker(storage, repo)
		assert.Nil(t, err)

		// Enable DB-level RU for tenant 11 (cluster 1001)
		worker.dbRuTenantIDs = map[uint64]struct{}{11: {}}

		reports, err := worker.genBillingReports(ctx, log.L())
		sort.Slice(reports, func(i, j int) bool { return reports[i].CostID < reports[j].CostID })
		assert.Nil(t, err)

		// Based on mock data: cluster 1001 has exec counts for test_db and another_db (9900 each)
		// The IsRu() method returns true for RU metrics AND egress metrics
		// So all RU and egress metrics will be split by database
		// Total load_data_ru: 19800, splits to: test_db=9900, another_db=9900
		// Total tidb_ru: 79200, splits to: test_db=39600, another_db=39600
		// Total tiflash_ru: 29700, splits to: test_db=14850, another_db=14850
		// Total privatelink_egress: 7, splits to: test_db=4, another_db=3
		// Total public_egress: 29, splits to: test_db=15, another_db=14
		// Expected: 5 metrics * 2 databases = 10 DB-split reports for tenant 11

		// Verify specific DB-split reports with exact costs
		foundReports := make(map[string]bool)
		expectedReports := map[string]int64{
			"us-east-1/1001/load_data_ru/0/db_test_db":          9900,
			"us-east-1/1001/load_data_ru/0/db_another_db":       9900,
			"us-east-1/1001/tidb_ru/0/db_test_db":               39600,
			"us-east-1/1001/tidb_ru/0/db_another_db":            39600,
			"us-east-1/1001/tiflash_ru/0/db_test_db":            14850,
			"us-east-1/1001/tiflash_ru/0/db_another_db":         14850,
			"us-east-1/1001/privatelink_egress/0/db_test_db":    4,
			"us-east-1/1001/privatelink_egress/0/db_another_db": 3,
			"us-east-1/1001/public_egress/0/db_test_db":         15,
			"us-east-1/1001/public_egress/0/db_another_db":      14,
		}

		tenant11DBSplitCount := 0
		for _, report := range reports {
			if report.TenantID == 11 {
				if expectedCost, exists := expectedReports[report.CostID]; exists {
					foundReports[report.CostID] = true
					assert.Equal(t, expectedCost, report.Cost, "Cost mismatch for %s", report.CostID)
					tenant11DBSplitCount++
					// Verify required fields
					assert.NotEmpty(t, report.Sku)
					assert.NotEmpty(t, report.Labels)
				}
			}
		}

		// Verify all expected reports were found
		assert.Equal(t, 10, tenant11DBSplitCount, "Expected 10 DB-split reports for tenant 11")
		for costID := range expectedReports {
			assert.True(t, foundReports[costID], "Expected to find report: %s", costID)
		}

		// Verify total costs are preserved (sum of DB splits should equal original)
		// load_data_ru: 9900 + 9900 = 19800
		// tidb_ru: 39600 + 39600 = 79200
		// tiflash_ru: 14850 + 14850 = 29700
		// privatelink_egress: 4 + 3 = 7
		// public_egress: 15 + 14 = 29
	})

	t.Run("POC tenant cost splitting", func(t *testing.T) {
		storage := store.NewMockStorage()
		repo := repository.NewMockClusterMetaRepo()
		worker, err := newDefaultMeteringWorker(storage, repo)
		assert.Nil(t, err)

		// Enable POC for tenant 11 (cluster 1001) with 30% free
		worker.pocTenantIDs = map[uint64]struct{}{11: {}}
		worker.pocFreePercent = 30

		reports, err := worker.genBillingReports(ctx, log.L())
		sort.Slice(reports, func(i, j int) bool { return reports[i].CostID < reports[j].CostID })
		assert.Nil(t, err)

		// Find reports for tenant 11 and verify splits
		tenant11Reports := make([]*BillingReport, 0)
		for _, report := range reports {
			if report.TenantID == 11 {
				tenant11Reports = append(tenant11Reports, report)
			}
		}

		// For tenant 11, each usage should be split into two reports: poc (30%) and normal (70%)
		// Expected: load_data_ru, privatelink_egress, public_egress, tidb_ru, tiflash_ru = 5 metrics * 2 = 10 reports
		assert.Equal(t, 10, len(tenant11Reports))

		// Verify specific cost splits for load_data_ru (19800 total)
		// 30% poc = 5940, 70% normal = 13860
		foundLoadDataRUPoc := false
		foundLoadDataRUNormal := false
		for _, report := range tenant11Reports {
			if report.CostID == "us-east-1/1001/load_data_ru/0/poc" {
				foundLoadDataRUPoc = true
				assert.Equal(t, int64(5940), report.Cost)
				assert.Equal(t, "true", report.Labels[BillingLabelBetaKey])
				assert.Equal(t, BillingBetaCategoryPreview, report.Labels[BillingLabelBetaCategoryKey])
			} else if report.CostID == "us-east-1/1001/load_data_ru/0/" {
				foundLoadDataRUNormal = true
				assert.Equal(t, int64(13860), report.Cost)
			}
		}
		assert.True(t, foundLoadDataRUPoc, "Expected POC report for load_data_ru")
		assert.True(t, foundLoadDataRUNormal, "Expected normal report for load_data_ru")

		// Verify specific cost splits for tidb_ru (79200 total)
		// 30% poc = 23760, 70% normal = 55440
		foundTidbRUPoc := false
		foundTidbRUNormal := false
		for _, report := range tenant11Reports {
			if report.CostID == "us-east-1/1001/tidb_ru/0/poc" {
				foundTidbRUPoc = true
				assert.Equal(t, int64(23760), report.Cost)
				assert.Equal(t, "true", report.Labels[BillingLabelBetaKey])
				assert.Equal(t, BillingBetaCategoryPreview, report.Labels[BillingLabelBetaCategoryKey])
			} else if report.CostID == "us-east-1/1001/tidb_ru/0/" {
				foundTidbRUNormal = true
				assert.Equal(t, int64(55440), report.Cost)
			}
		}
		assert.True(t, foundTidbRUPoc, "Expected POC report for tidb_ru")
		assert.True(t, foundTidbRUNormal, "Expected normal report for tidb_ru")

		// Verify the total cost is preserved (sum of all poc + normal reports should equal original)
		totalPocCost := int64(0)
		totalNormalCost := int64(0)
		for _, report := range tenant11Reports {
			if len(report.CostID) > 4 && report.CostID[len(report.CostID)-4:] == "/poc" {
				totalPocCost += report.Cost
			} else {
				totalNormalCost += report.Cost
			}
		}
		// Total should be: 19800 + 7 + 29 + 79200 + 29700 = 128736
		expectedTotal := int64(19800 + 7 + 29 + 79200 + 29700)
		assert.Equal(t, expectedTotal, totalPocCost+totalNormalCost)
	})

	t.Run("verify cost totals remain consistent", func(t *testing.T) {
		storage := store.NewMockStorage()
		repo := repository.NewMockClusterMetaRepo()

		// Create worker without special handling
		worker1, err := newDefaultMeteringWorker(storage, repo)
		assert.Nil(t, err)
		reports1, err := worker1.genBillingReports(ctx, log.L())
		assert.Nil(t, err)

		// Create worker with POC handling for tenant 11
		worker2, err := newDefaultMeteringWorker(storage, repo)
		assert.Nil(t, err)
		worker2.pocTenantIDs = map[uint64]struct{}{11: {}}
		worker2.pocFreePercent = 50
		reports2, err := worker2.genBillingReports(ctx, log.L())
		assert.Nil(t, err)

		// Calculate total cost for tenant 11 from both workers
		total1 := int64(0)
		for _, r := range reports1 {
			if r.TenantID == 11 {
				total1 += r.Cost
			}
		}

		total2 := int64(0)
		for _, r := range reports2 {
			if r.TenantID == 11 {
				total2 += r.Cost
			}
		}

		// Total cost should remain the same despite splitting
		assert.Equal(t, total1, total2, "Total cost should remain consistent with POC splitting")
	})

	t.Run("multiple tenants with mixed configurations", func(t *testing.T) {
		storage := store.NewMockStorage()
		repo := repository.NewMockClusterMetaRepo()
		worker, err := newDefaultMeteringWorker(storage, repo)
		assert.Nil(t, err)

		// Setup: tenant 11 gets POC, tenant 12 gets DB-level RU
		worker.pocTenantIDs = map[uint64]struct{}{11: {}}
		worker.pocFreePercent = 20
		worker.dbRuTenantIDs = map[uint64]struct{}{12: {}}

		reports, err := worker.genBillingReports(ctx, log.L())
		assert.Nil(t, err)

		// Verify we have reports for both tenants
		tenant11Reports := 0
		tenant12Reports := 0
		for _, report := range reports {
			if report.TenantID == 11 {
				tenant11Reports++
			}
			if report.TenantID == 12 {
				tenant12Reports++
			}
		}

		assert.True(t, tenant11Reports > 0, "Expected reports for tenant 11")
		assert.True(t, tenant12Reports > 0, "Expected reports for tenant 12")
	})

	t.Run("verify report structure and labels", func(t *testing.T) {
		storage := store.NewMockStorage()
		repo := repository.NewMockClusterMetaRepo()
		worker, err := newDefaultMeteringWorker(storage, repo)
		assert.Nil(t, err)

		reports, err := worker.genBillingReports(ctx, log.L())
		assert.Nil(t, err)

		for _, report := range reports {
			// Verify CostID format: region/primaryID/metricType/timestamp
			assert.Contains(t, report.CostID, "us-east-1")
			assert.Contains(t, report.CostID, "/")

			// Verify required labels
			assert.NotEmpty(t, report.Labels[BillingLabelTenantID])
			assert.NotEmpty(t, report.Labels[BillingLabelProjectID])
			assert.NotEmpty(t, report.Labels[BillingLabelClusterID])
			assert.Equal(t, "us-east-1", report.Labels[BillingLabelRegion])
			assert.Equal(t, "true", report.Labels[BillingLabelServerless])
			assert.Equal(t, WorkerNameDefault, report.Labels[BillingLabelWorker])

			// Verify SKU format
			assert.Contains(t, report.Sku, "serverless")
			assert.Contains(t, report.Sku, "aws")
			assert.Contains(t, report.Sku, "us-east-1")

			// Verify timestamp
			assert.True(t, report.CostTimeStamp > 0)
		}
	})
}