package tests

import (
	"fmt"
	"testing"
	"time"

	"github.com/google/uuid"
	"github.com/pingcap/dbaas/billing2/config"
	"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/handler/cachehandler"
	"github.com/pingcap/dbaas/billing2/handler/centralhandler"
	"github.com/pingcap/dbaas/billing2/handler/infrahandler"
	"github.com/pingcap/dbaas/billing2/model/billingconfig"
	"github.com/pingcap/dbaas/billing2/model/costexplorer"
	"github.com/pingcap/dbaas/billing2/utils/panicIf"
	"github.com/pingcap/dbaas/billing2/utils/tests"
	"github.com/pingcap/dbaas/billing2/utils/tests/datainit"
	"github.com/shopspring/decimal"
	"github.com/stretchr/testify/suite"
)

type MockInfraHandlerImpl struct {
	infrahandler.InfraHandler
}

func (m *MockInfraHandlerImpl) ListClusters(tenantID uint) (ret []*forms.Cluster, err error) {
	return []*forms.Cluster{
		{
			ID:       datainit.DefaultUnitCostsDedicatedAWSClusterID,
			Name:     fmt.Sprint(constant.UnknownClusterName, 1),
			TenantID: tenantID,
		},
		{
			ID:       datainit.DefaultUnitCostsDedicatedGCPClusterID,
			Name:     fmt.Sprint(constant.UnknownClusterName, 1),
			TenantID: tenantID,
		},
		{
			ID:       10,
			Name:     fmt.Sprint(constant.UnknownClusterName, 10),
			TenantID: tenantID,
		},
	}, nil
}

type MockCentralHandlerImpl struct {
	centralhandler.CentralHandler
}

func (m *MockCentralHandlerImpl) GetTenant(tenantID uint) (ret *forms.Tenant, err error) {
	return &forms.Tenant{
		TenantID: tenantID,
		Status:   centralhandler.TenantStatusActive,
	}, nil
}

func (*MockCentralHandlerImpl) ListProjects(tenantID uint, containsDeleted bool) (ret []*forms.Project, err error) {
	return []*forms.Project{
		{
			ID:       1,
			TenantID: tenantID,
			Name:     "Project 1",
		},
		{
			ID:       2,
			TenantID: tenantID,
			Name:     "Project 2",
		},
		{
			ID:       3,
			TenantID: tenantID,
			Name:     "Project 1",
		},
	}, nil
}

func (*MockCentralHandlerImpl) GetProjectInfo(projectID uint) (*forms.Project, error) {
	return &forms.Project{
		ID:   projectID,
		Name: "Project",
	}, nil
}

type CostExplorerTrendsTestSuite struct {
	tests.BillingTestSuite
}

func TestCostExplorerTrendsSuite(t *testing.T) {
	suite.Run(t, new(CostExplorerTrendsTestSuite))
}

func (suite *CostExplorerTrendsTestSuite) SetupTest() {
	suite.BillingTestSuite.SetupTest()
	_, err := datainit.CreateBillingConfig(1, entity.BillingConfigType_DirectCustomer)
	panicIf.Error(err)
	datainit.CreateCategoryLabels()
	datainit.CreateDefaultUnits()
	cachehandler.InitWithLocal()
	config.Get().Set(config.EnableCache, true)
	infrahandler.Set(&MockInfraHandlerImpl{})
	centralhandler.SetCentralHandler(&MockCentralHandlerImpl{})
}

func (suite *CostExplorerTrendsTestSuite) TearDownTest() {
	suite.BillingTestSuite.SetupTest()
	config.Get().Set(config.EnableCache, false)
}

func (suite *CostExplorerTrendsTestSuite) TestGenerateCostTrends() {
	assert := suite.Assert()
	var err error
	tenantID := uint(30008)
	_, err = datainit.CreateBillingConfig(tenantID, entity.BillingConfigType_DirectCustomer)
	assert.NoError(err)
	billingConfigID, err := (&billingconfig.BillingConfigModel{}).GetActiveBillingConfigID(tenantID)
	panicIf.Error(err)
	datainit.CreateDefaultUserUnitCosts(tenantID, datainit.BillingDate, 60)
	start := datainit.BillingDate
	end := datainit.BillingDate.AddDate(0, 1, 0)
	timeGroup := forms.CostExploreTrendsTimeGroupDay
	internal := false
	query := &forms.CostExploreTrendsInput{
		CostExplorerTrendsFilter: forms.CostExplorerTrendsFilter{
			StartDate: &start,
			EndDate:   &end,
			Internal:  &internal,
		},
		TimeGroup:   &timeGroup,
		BreakdownBy: []forms.CostExplorerTrendsBreakdownKey{forms.CostExploreTrendsBreakdownKeyService, forms.CostExploreTrendsBreakdownKeyRegion, forms.CostExploreTrendsBreakdownKeyCluster, forms.CostExploreTrendsBreakdownKeyProject},
	}
	r, err := costexplorer.CostTrends(tenantID, billingConfigID, query)
	assert.NoError(err)
	assert.Len(r.Elements, 30)
	assert.Equal(decimal.RequireFromString("350"), r.RawPrice)
	assert.Equal(decimal.RequireFromString("120"), r.Discount)
	assert.Equal(decimal.RequireFromString("70"), r.CreditsUsage)
	assert.Equal(decimal.RequireFromString("160"), r.CostPrice)
	for i, elem := range r.Elements {
		if datainit.BillingDate.AddDate(0, 0, i).Day() <= 10 {
			assert.Equal(decimal.Zero, elem.CostPrice)
		}
		for _, breakdown := range elem.Breakdown {
			assert.NotNil(breakdown.Service)
			assert.Equal(costexplorer.FormatServiceName(&breakdown.CostExplorerService), *breakdown.Service)
			assert.NotNil(breakdown.ProjectID)
			assert.NotNil(breakdown.ProjectName)
			if breakdown.ClusterID == nil {
				assert.Equal("Data Transfer", *breakdown.Category)
			} else if *breakdown.ClusterID == 1 {
				assert.Equal("TiDB Dedicated", *breakdown.Product)
			} else if *breakdown.ClusterID == 10 {
				assert.Equal("TiDB Serverless", *breakdown.Product)
			}
			if *breakdown.ProjectID == 1 {
				assert.Equal("Project 1", *breakdown.ProjectName)
			}
		}
	}
	timeGroup = forms.CostExploreTrendsTimeGroupMonth
	query.TimeGroup = &timeGroup
	start = time.Date(2021, 10, 17, 0, 0, 0, 0, time.UTC)
	end = time.Date(2022, 1, 1, 0, 0, 0, 0, time.UTC)
	query.StartDate = &start
	query.EndDate = &end
	r, err = costexplorer.CostTrends(tenantID, billingConfigID, query)
	assert.NoError(err)
	assert.Len(r.Elements, 3)
	for i, billingDate := range []string{"2021-10-01", "2021-11-01", "2021-12-01"} {
		assert.Equal(billingDate, r.Elements[i].BillingDate)
	}
	c, err := costexplorer.CostTrendsCSV(tenantID, billingConfigID, &forms.CostExploreTrendsCSVInput{CostExploreTrendsInput: *query})
	assert.NoError(err)
	target := `service,region,cluster,project,total_cost,2021-10-01,2021-11-01,2021-12-01
TiDB Dedicated - Data Transfer - Same Region,Others,Others,Project 1,80,0,40,40
TiDB Dedicated - Node Compute - TiDB,us-east-1,unknown1,Project 1,40,0,20,20
TiDB Dedicated - Node Compute - TiDB,us-west-2,unknown1,Project 1,40,0,20,20
TiDB Dedicated - Node Compute - TiKV,us-east-1,unknown1,Project 1,40,0,20,20
TiDB Dedicated - Node Compute - TiKV,us-west-2,unknown1,Project 1,40,0,20,20
TiDB Serverless - Request Units,us-east-1,unknown10,Project 1,40,0,20,20
TiDB Serverless - Row-based Storage,us-east-1,unknown10,Project 1,40,0,20,20
Total,0,0,0,320,0,160,160
`
	assert.Equal(target, string(c))
	c2Targets := `charge_type,subtotal,2021-10-01,2021-11-01,2021-12-01
Total Cost,320,0,160,160
Running Total,700,0,350,350
Discount,-240,0,-120,-120
Credits,-140,0,-70,-70
`
	query.BreakdownBy = []forms.CostExplorerTrendsBreakdownKey{}
	c2, err := costexplorer.CostTrendsCSV(tenantID, billingConfigID, &forms.CostExploreTrendsCSVInput{CostExploreTrendsInput: *query})
	assert.NoError(err)
	assert.Equal(c2Targets, string(c2))
	c3Targets := `charge_type,subtotal,2020-10-01,2020-11-01,2020-12-01
Total Cost,0,0,0,0
Running Total,0,0,0,0
Discount,0,0,0,0
Credits,0,0,0,0
`
	start = time.Date(2020, 10, 1, 0, 0, 0, 0, time.UTC)
	end = time.Date(2021, 1, 1, 0, 0, 0, 0, time.UTC)
	query.StartDate = &start
	query.EndDate = &end
	c3, err := costexplorer.CostTrendsCSV(tenantID, billingConfigID, &forms.CostExploreTrendsCSVInput{CostExploreTrendsInput: *query})
	assert.NoError(err)
	assert.Equal(c3Targets, string(c3))
	start = time.Date(2021, 10, 1, 0, 0, 0, 0, time.UTC)
	end = time.Date(2022, 1, 1, 0, 0, 0, 0, time.UTC)
	query.StartDate = &start
	query.EndDate = &end
	query.TimeGroup = nil
	c4, err := costexplorer.CostTrendsCSV(tenantID, billingConfigID, &forms.CostExploreTrendsCSVInput{CostExploreTrendsInput: *query})
	assert.NoError(err)
	c4Targets := `charge_type,subtotal
Total Cost,320
Running Total,700
Discount,-240
Credits,-140
`
	assert.Equal(c4Targets, string(c4))
	ent := entity.UserUnitCost{
		CostID:                uuid.NewString(),
		TenantID:              tenantID,
		BindedBillingConfigID: billingConfigID,
		CostTimeStamp:         time.Date(2021, 10, 2, 0, 0, 0, 0, time.UTC).Unix(),
		Unit:                  datainit.AWSNodeUnit,
		CostPrice:             decimal.NewFromInt(1),
		RawPrice:              decimal.NewFromInt(2),
		Labels:                map[string]string{"component": constant.TiDBComp, constant.ClusterLabelKey: fmt.Sprint(3), constant.ProjectLabelKey: fmt.Sprint(3)},
		ConsumedSource:        entity.ConsumedSource_Invoice,
		Status:                entity.UserUnitCostStatus_DONE,
		Internal:              false,
	}
	assert.Nil(ent.Create(db.GetDB()))
	query.BreakdownBy = []forms.CostExplorerTrendsBreakdownKey{forms.CostExploreTrendsBreakdownKeyCluster, forms.CostExploreTrendsBreakdownKeyProject}
	c5, err := costexplorer.CostTrendsCSV(tenantID, billingConfigID, &forms.CostExploreTrendsCSVInput{CostExploreTrendsInput: *query})
	assert.NoError(err)
	c5Targets := `cluster,project,total_cost
unknown1,Project 1,80
unknown10,Project 1,80
unknown1,Project 1,80
3,Project 1,1
Others,Project 1,80
Total,0,321
`
	assert.Equal(c5Targets, string(c5))
	query.Providers = []forms.CostExplorerProvider{
		{
			Name:    "AWS",
			Regions: []string{"Others"},
		},
	}
	c6, err := costexplorer.CostTrendsCSV(tenantID, billingConfigID, &forms.CostExploreTrendsCSVInput{CostExploreTrendsInput: *query})
	assert.NoError(err)
	c6Targets := `cluster,project,total_cost
Others,Project 1,80
Total,0,80
`
	assert.Equal(c6Targets, string(c6))
	c6D, err := costexplorer.CostTrendsCSV(tenantID, billingConfigID, &forms.CostExploreTrendsCSVInput{CostExploreTrendsInput: *query, DisplayCreditsAndDiscountOnCSV: true})
	assert.NoError(err)
	c6DTargets := `cluster,project,running_total,discounts,credits,total_cost
Others,Project 1,120,0,-40,80
Total,0,120,0,-40,80
`
	assert.Equal(c6DTargets, string(c6D))
}

func (suite *CostExplorerTrendsTestSuite) TestClusterBreakdownFilterForDisabledTenants() {
	assert := suite.Assert()
	var err error
	tenantID := uint(30009)
	_, err = datainit.CreateBillingConfig(tenantID, entity.BillingConfigType_DirectCustomer)
	assert.NoError(err)
	billingConfigID, err := (&billingconfig.BillingConfigModel{}).GetActiveBillingConfigID(tenantID)
	panicIf.Error(err)
	datainit.CreateDefaultUserUnitCosts(tenantID, datainit.BillingDate, 60)
	start := datainit.BillingDate
	end := datainit.BillingDate.AddDate(0, 1, 0)
	timeGroup := forms.CostExploreTrendsTimeGroupDay
	internal := false

	config.Get().Set(config.DisableClusterDisplayTenants, []string{fmt.Sprint(tenantID)})
	defer config.Get().Set(config.DisableClusterDisplayTenants, []string{})

	query := &forms.CostExploreTrendsInput{
		CostExplorerTrendsFilter: forms.CostExplorerTrendsFilter{
			StartDate: &start,
			EndDate:   &end,
			Internal:  &internal,
		},
		TimeGroup:   &timeGroup,
		BreakdownBy: []forms.CostExplorerTrendsBreakdownKey{forms.CostExploreTrendsBreakdownKeyService, forms.CostExploreTrendsBreakdownKeyRegion, forms.CostExploreTrendsBreakdownKeyCluster, forms.CostExploreTrendsBreakdownKeyProject},
	}
	r, err := costexplorer.CostTrends(tenantID, billingConfigID, query)
	assert.NoError(err)
	assert.Len(r.Elements, 30)

	for _, elem := range r.Elements {
		for _, breakdown := range elem.Breakdown {
			assert.NotNil(breakdown.Service)
			assert.NotNil(breakdown.ProjectID)
			assert.NotNil(breakdown.ProjectName)
			assert.Nil(breakdown.ClusterID)
			assert.Nil(breakdown.ClusterName)
		}
	}
}

// TestTrendsDayGroupSplitByDay verifies that TimeGroup=day splits by day — the most common
// trends query ("last month, daily"). Each per-day sub-query scans one day of user_unit_costs
// so even tenants with tens of thousands of clusters complete well under 30s.
func (suite *CostExplorerTrendsTestSuite) TestTrendsDayGroupSplitByDay() {
	assert := suite.Assert()
	tenantID := uint(30009)
	_, err := datainit.CreateBillingConfig(tenantID, entity.BillingConfigType_DirectCustomer)
	assert.NoError(err)
	billingConfigID, err := (&billingconfig.BillingConfigModel{}).GetActiveBillingConfigID(tenantID)
	panicIf.Error(err)
	// 3 days of data — should produce 3 per-day sub-queries
	datainit.CreateDefaultUserUnitCosts(tenantID, datainit.BillingDate, 3)
	start := datainit.BillingDate
	end := datainit.BillingDate.AddDate(0, 0, 3)
	timeGroup := forms.CostExploreTrendsTimeGroupDay
	internal := false
	query := &forms.CostExploreTrendsInput{
		CostExplorerTrendsFilter: forms.CostExplorerTrendsFilter{
			StartDate: &start,
			EndDate:   &end,
			Internal:  &internal,
		},
		TimeGroup:   &timeGroup,
		BreakdownBy: []forms.CostExplorerTrendsBreakdownKey{forms.CostExploreTrendsBreakdownKeyService},
	}
	r, err := costexplorer.CostTrends(tenantID, billingConfigID, query)
	assert.NoError(err)
	assert.NotNil(r)
	assert.Len(r.Elements, 3)
	// Days ≤ 10 are credits (CostPrice = 0), so assert on RawPrice instead
	assert.True(r.RawPrice.GreaterThan(decimal.Zero),
		"3-day day-grouped query should produce non-zero raw price")
}

// TestTrendsMonthGroupSplitByMonth verifies that TimeGroup=month splits by calendar month
// when the range spans more than one month.
func (suite *CostExplorerTrendsTestSuite) TestTrendsMonthGroupSplitByMonth() {
	assert := suite.Assert()
	tenantID := uint(30010)
	_, err := datainit.CreateBillingConfig(tenantID, entity.BillingConfigType_DirectCustomer)
	assert.NoError(err)
	billingConfigID, err := (&billingconfig.BillingConfigModel{}).GetActiveBillingConfigID(tenantID)
	panicIf.Error(err)
	// 60 days spanning 2-3 months
	datainit.CreateDefaultUserUnitCosts(tenantID, datainit.BillingDate, 60)
	start := datainit.BillingDate
	end := datainit.BillingDate.AddDate(0, 0, 60)
	timeGroup := forms.CostExploreTrendsTimeGroupMonth
	internal := false
	query := &forms.CostExploreTrendsInput{
		CostExplorerTrendsFilter: forms.CostExplorerTrendsFilter{
			StartDate: &start,
			EndDate:   &end,
			Internal:  &internal,
		},
		TimeGroup:   &timeGroup,
		BreakdownBy: []forms.CostExplorerTrendsBreakdownKey{forms.CostExploreTrendsBreakdownKeyService},
	}
	r, err := costexplorer.CostTrends(tenantID, billingConfigID, query)
	assert.NoError(err)
	assert.NotNil(r)
	// Monthly grouping of 60 days should produce 2-3 monthly elements
	assert.GreaterOrEqual(len(r.Elements), 2)
	assert.True(r.RawPrice.GreaterThan(decimal.Zero))
}

// TestTrendsTotalTimeGroupNoSplit verifies that TimeGroup=nil ("total") is NOT split —
// the query is already aggregated to few rows per breakdown key and is fast enough.
func (suite *CostExplorerTrendsTestSuite) TestTrendsTotalTimeGroupNoSplit() {
	assert := suite.Assert()
	tenantID := uint(30011)
	_, err := datainit.CreateBillingConfig(tenantID, entity.BillingConfigType_DirectCustomer)
	assert.NoError(err)
	billingConfigID, err := (&billingconfig.BillingConfigModel{}).GetActiveBillingConfigID(tenantID)
	panicIf.Error(err)
	datainit.CreateDefaultUserUnitCosts(tenantID, datainit.BillingDate, 45)
	start := datainit.BillingDate
	end := datainit.BillingDate.AddDate(0, 0, 45)
	internal := false
	query := &forms.CostExploreTrendsInput{
		CostExplorerTrendsFilter: forms.CostExplorerTrendsFilter{
			StartDate: &start,
			EndDate:   &end,
			Internal:  &internal,
		},
		TimeGroup:   nil,
		BreakdownBy: []forms.CostExplorerTrendsBreakdownKey{forms.CostExploreTrendsBreakdownKeyService},
	}
	r, err := costexplorer.CostTrends(tenantID, billingConfigID, query)
	assert.NoError(err)
	assert.NotNil(r)
	assert.Len(r.Elements, 1)
	assert.True(r.RawPrice.GreaterThan(decimal.Zero))
}