package tests

import (
	"github.com/google/uuid"
	"github.com/pingcap/dbaas/billing2/constant"
	"github.com/pingcap/dbaas/billing2/db"
	"github.com/pingcap/dbaas/billing2/db/entity"
	"github.com/pingcap/dbaas/billing2/forms"
	"github.com/pingcap/dbaas/billing2/model/billing"
	"github.com/pingcap/dbaas/billing2/model/billingconfig"
	"github.com/pingcap/dbaas/billing2/model/costexplorer"
	"github.com/pingcap/dbaas/billing2/utils/tests"
	"github.com/pingcap/dbaas/billing2/utils/tests/datainit"
	"github.com/shopspring/decimal"
)

// streamTestCategory is intentionally NOT a registered category item, so GetServiceName falls
// back to TrimSuffix(category,"Cost") == "StreamTestCat" deterministically.
const streamTestCategory = "StreamTestCat"
const streamTestService = "StreamTestCat"

func (suite *CostExplorerSuite) insertDetailCost(tenantID, billingConfigID uint, ts int64, unit string, labels map[string]string, source string, raw, cost decimal.Decimal, internal bool) {
	suite.Assert().NoError((&entity.UserUnitCost{
		CostID:                uuid.NewString(),
		TenantID:              tenantID,
		BindedBillingConfigID: billingConfigID,
		CostTimeStamp:         ts,
		Unit:                  unit,
		Cost:                  1,
		CostPrice:             cost,
		RawPrice:              raw,
		Labels:                labels,
		ConsumedSource:        source,
		Status:                entity.UserUnitCostStatus_DONE,
		Internal:              internal,
	}).Create(db.GetDB()))
}

func (suite *CostExplorerSuite) createStreamTestUnit(key string) {
	_, err := (&billing.UnitModel{}).CreateUnit(&forms.UnitInfo{
		UnitKey:   key,
		Category:  streamTestCategory,
		Provider:  "AWS",
		Region:    "us-east-1",
		Price:     decimal.NewFromInt(1),
		PriceType: entity.PriceType_Normal,
	})
	suite.Assert().NoError(err)
}

// seedDetailDataset inserts, for the given tenant, 4 rows on BillingDate:
//   - cluster 101, invoice, raw 10.009 / cost 9.009
//   - cluster 102, invoice, raw 20.009 / cost 18.009
//   - no cluster, credits, raw  5.009 / cost  4.009
//   - cluster 102, internal,raw   100 / cost  100   (must be excluded everywhere)
func (suite *CostExplorerSuite) seedDetailDataset(tenantID uint, unit string) {
	bcID, err := (&billingconfig.BillingConfigModel{}).GetActiveBillingConfigID(tenantID)
	suite.Assert().NoError(err)
	ts := datainit.BillingDate.Unix()
	suite.insertDetailCost(tenantID, bcID, ts, unit, map[string]string{constant.ClusterLabelKey: "101"}, entity.ConsumedSource_Invoice, decimal.NewFromFloat(10.009), decimal.NewFromFloat(9.009), false)
	suite.insertDetailCost(tenantID, bcID, ts, unit, map[string]string{constant.ClusterLabelKey: "102"}, entity.ConsumedSource_Invoice, decimal.NewFromFloat(20.009), decimal.NewFromFloat(18.009), false)
	suite.insertDetailCost(tenantID, bcID, ts, unit, map[string]string{}, entity.ConsumedSource_Credits, decimal.NewFromFloat(5.009), decimal.NewFromFloat(4.009), false)
	suite.insertDetailCost(tenantID, bcID, ts, unit, map[string]string{constant.ClusterLabelKey: "102"}, entity.ConsumedSource_Invoice, decimal.NewFromInt(100), decimal.NewFromInt(100), true)
}

// TestStreamingExportBillsDetailsMultiTenant verifies the streamed GetExportBillsDetailsByTenant
// aggregates correctly across multiple tenants/clusters and excludes internal costs.
func (suite *CostExplorerSuite) TestStreamingExportBillsDetailsMultiTenant() {
	assert := suite.Assert()
	_, err := datainit.CreateBillingConfig(2, entity.BillingConfigType_DirectCustomer)
	assert.NoError(err)
	unit := "stream-export-" + uuid.NewString()
	suite.createStreamTestUnit(unit)
	suite.seedDetailDataset(1, unit)
	suite.seedDetailDataset(2, unit)

	model := &costexplorer.CostExplorerModel{}
	details, err := model.GetExportBillsDetailsByTenant([]uint{1, 2}, nil, datainit.BillingDate.Year(), int(datainit.BillingDate.Month()), nil)
	assert.NoError(err)

	// 3 external rows per tenant (cluster 101 invoice, cluster 102 invoice, no-cluster credits).
	streamRows := filterExportByService(details)
	assert.Len(streamRows, 6)

	for _, tenantID := range []uint{1, 2} {
		byCluster := map[string]forms.ExportUsageDetails{}
		for _, r := range streamRows {
			if r.TenantID == tenantID {
				byCluster[r.ClusterID] = r
				assert.Equal(streamTestService, r.Service)
				assert.Equal("AWS", r.CloudProvider)
				assert.Equal("us-east-1", r.Region)
			}
		}
		assert.Len(byCluster, 3)
		// Export keeps raw (untruncated) sums.
		assert.True(decimal.NewFromFloat(10.009).Equal(byCluster["101"].Amount))
		assert.True(decimal.NewFromFloat(9.009).Equal(byCluster["101"].BillableCost))
		assert.True(decimal.NewFromFloat(20.009).Equal(byCluster["102"].Amount))
		assert.True(decimal.NewFromFloat(18.009).Equal(byCluster["102"].BillableCost))
		assert.True(decimal.NewFromFloat(5.009).Equal(byCluster[""].Amount))
		assert.True(byCluster[""].BillableCost.IsZero())
		assert.True(decimal.NewFromFloat(4.009).Equal(byCluster[""].CreditUsage))
	}
}

// TestStreamingPartnerBillsDetails verifies the streamed GetPartnerBillsDetails truncation and
// per-customer attribution, with internal costs excluded.
func (suite *CostExplorerSuite) TestStreamingPartnerBillsDetails() {
	assert := suite.Assert()
	tests.MockCentralHandler()
	_, err := datainit.CreateBillingConfig(2, entity.BillingConfigType_DirectCustomer)
	assert.NoError(err)
	unit := "stream-partner-" + uuid.NewString()
	suite.createStreamTestUnit(unit)
	suite.seedDetailDataset(1, unit)
	suite.seedDetailDataset(2, unit)

	model := &costexplorer.CostExplorerModel{}
	res, err := model.GetPartnerBillsDetails([]uint{1, 2}, nil, datainit.BillingDate.Year(), int(datainit.BillingDate.Month()), true)
	assert.NoError(err)

	rows := filterPartnerByService(res.BillsDetails, streamTestService)
	// Partner details dedup on cluster/project NAME (empty without enrichment), so the two
	// invoice clusters merge; the credits row stays separate => 2 rows per customer.
	assert.Len(rows, 4)
	for _, customerID := range []string{"1", "2"} {
		var invoiceRow, creditsRow forms.PartnerBillsDetail
		count := 0
		for _, r := range rows {
			if r.CustomerID != customerID {
				continue
			}
			count++
			if r.Credits.IsZero() {
				invoiceRow = r
			} else {
				creditsRow = r
			}
		}
		assert.Equal(2, count)
		// Partner details truncate to BillingCostsPrecisionDigits (2) and compute discounts.
		assert.True(decimal.NewFromFloat(30.01).Equal(invoiceRow.RunningTotal))
		assert.True(decimal.NewFromFloat(27.01).Equal(invoiceRow.TotalCost))
		assert.True(decimal.NewFromFloat(3).Equal(invoiceRow.Discounts))
		assert.True(decimal.NewFromInt(5).Equal(creditsRow.RunningTotal))
		assert.True(decimal.NewFromInt(4).Equal(creditsRow.Credits))
		assert.True(decimal.NewFromInt(1).Equal(creditsRow.Discounts))
	}
}

// TestStreamingBillsDetailsClusterFilter verifies single-tenant GetBillsDetails, including the
// explicit cluster filter path (which forces the Aurora cluster index).
func (suite *CostExplorerSuite) TestStreamingBillsDetailsClusterFilter() {
	assert := suite.Assert()
	unit := "stream-bills-" + uuid.NewString()
	suite.createStreamTestUnit(unit)
	suite.seedDetailDataset(1, unit)

	model := &costexplorer.CostExplorerModel{}
	year, month := datainit.BillingDate.Year(), int(datainit.BillingDate.Month())

	// Without name enrichment the two invoice rows share an empty cluster name and collapse;
	// the credits row stays separate => 2 rows for the whole tenant.
	all, err := model.GetBillsDetails(1, 0, year, month, 0)
	assert.NoError(err)
	allRows := filterUsageByService(all.Details, streamTestService)
	assert.Len(allRows, 2)
	var invoiceRow, creditsRow forms.UsageDetails
	for _, r := range allRows {
		if r.CreditUsage.IsZero() {
			invoiceRow = r
		} else {
			creditsRow = r
		}
	}
	assert.True(decimal.NewFromFloat(30.01).Equal(invoiceRow.Amount))
	assert.True(decimal.NewFromFloat(27.01).Equal(invoiceRow.BillableCost))
	assert.True(decimal.NewFromInt(5).Equal(creditsRow.Amount))
	assert.True(decimal.NewFromInt(4).Equal(creditsRow.CreditUsage))

	// Explicit cluster filter returns only cluster 101's invoice cost.
	cluster, err := model.GetBillsDetails(1, 0, year, month, 101)
	assert.NoError(err)
	clusterRows := filterUsageByService(cluster.Details, streamTestService)
	assert.Len(clusterRows, 1)
	assert.True(decimal.NewFromInt(10).Equal(clusterRows[0].Amount))
	assert.True(decimal.NewFromInt(9).Equal(clusterRows[0].BillableCost))
	assert.True(decimal.NewFromInt(1).Equal(clusterRows[0].DiscountCost))
}

// TestStreamingDailyCostsExcludingUnits exercises the AR-table streaming source (#2): excluded
// units are filtered out, the slice wrapper stays in parity, and the empty range is a no-op.
func (suite *CostExplorerSuite) TestStreamingDailyCostsExcludingUnits() {
	assert := suite.Assert()
	unitA := "stream-daily-a-" + uuid.NewString()
	unitB := "stream-daily-b-" + uuid.NewString()
	suite.createStreamTestUnit(unitA)
	suite.createStreamTestUnit(unitB)
	bcID, err := (&billingconfig.BillingConfigModel{}).GetActiveBillingConfigID(1)
	assert.NoError(err)

	day0 := datainit.BillingDate
	day1 := datainit.BillingDate.AddDate(0, 0, 1)
	suite.insertDetailCost(1, bcID, day0.Unix(), unitA, map[string]string{}, entity.ConsumedSource_Invoice, decimal.NewFromInt(3), decimal.NewFromInt(2), false)
	suite.insertDetailCost(1, bcID, day1.Unix(), unitA, map[string]string{}, entity.ConsumedSource_Invoice, decimal.NewFromInt(4), decimal.NewFromInt(3), false)
	suite.insertDetailCost(1, bcID, day0.Unix(), unitB, map[string]string{}, entity.ConsumedSource_Invoice, decimal.NewFromInt(9), decimal.NewFromInt(9), false)
	// internal must be excluded by the WHERE internal = false clause.
	suite.insertDetailCost(1, bcID, day0.Unix(), unitA, map[string]string{}, entity.ConsumedSource_Invoice, decimal.NewFromInt(50), decimal.NewFromInt(50), true)

	model := &costexplorer.CostExplorerModel{}
	rangeEnd := datainit.BillingDate.AddDate(0, 0, 2)

	var streamed []entity.CostsDetailsQuery
	err = model.ForEachDailyCostByDateRangeExcludingUnits(1, bcID, datainit.BillingDate, rangeEnd, []string{unitB}, func(c entity.CostsDetailsQuery) error {
		streamed = append(streamed, c)
		return nil
	})
	assert.NoError(err)
	// Only the two external unitA rows survive (unitB excluded, internal filtered).
	assert.Len(streamed, 2)
	total := decimal.Zero
	for _, c := range streamed {
		assert.Equal(unitA, c.Unit)
		total = total.Add(c.CostPrice)
	}
	assert.True(decimal.NewFromInt(5).Equal(total))

	// Slice wrapper must stay in parity with the streaming source.
	sliceRows, err := model.GetDailyCostsByDateRangeExcludingUnits(1, bcID, datainit.BillingDate, rangeEnd, []string{unitB})
	assert.NoError(err)
	assert.Len(sliceRows, 2)

	// Empty range: the handler is never invoked and no error is returned.
	emptyStart := datainit.BillingDate.AddDate(0, -2, 0)
	called := false
	err = model.ForEachDailyCostByDateRangeExcludingUnits(1, bcID, emptyStart, emptyStart.AddDate(0, 0, 1), nil, func(c entity.CostsDetailsQuery) error {
		called = true
		return nil
	})
	assert.NoError(err)
	assert.False(called)
}

func filterExportByService(details []forms.ExportUsageDetails) []forms.ExportUsageDetails {
	out := make([]forms.ExportUsageDetails, 0, len(details))
	for _, d := range details {
		if d.Service == streamTestService {
			out = append(out, d)
		}
	}
	return out
}

func filterPartnerByService(details []forms.PartnerBillsDetail, service string) []forms.PartnerBillsDetail {
	out := make([]forms.PartnerBillsDetail, 0, len(details))
	for _, d := range details {
		if d.ServicePathName == service {
			out = append(out, d)
		}
	}
	return out
}

func filterUsageByService(details []forms.UsageDetails, service string) []forms.UsageDetails {
	out := make([]forms.UsageDetails, 0, len(details))
	for _, d := range details {
		if d.Service == service {
			out = append(out, d)
		}
	}
	return out
}