package labels

import (
	"strconv"
	"testing"

	"github.com/pingcap/dbaas/billing2/constant"
	"github.com/pingcap/dbaas/billing2/forms"
	"github.com/pingcap/dbaas/billing2/handler/centralhandler"
	"github.com/pingcap/dbaas/billing2/handler/infrahandler"
	"github.com/stretchr/testify/suite"
)

type MockCentralHandlerImpl struct {
	centralhandler.CentralHandlerImpl
}

type MockInfraHandlerImpl struct {
	infrahandler.InfraHandlerImpl
}

func (*MockCentralHandlerImpl) GetProjectName(projectID uint) (string, error) {
	return strconv.FormatUint(uint64(projectID), 10), nil
}

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

func (*MockInfraHandlerImpl) ListClusters(tenantID uint) (ret []*forms.Cluster, err error) {
	return []*forms.Cluster{}, nil
}

func (*MockInfraHandlerImpl) GetCluster(tenantID, clusterID uint) (ret *forms.Cluster, err error) {
	return &forms.Cluster{}, nil
}

type LabelsTestSuite struct {
	suite.Suite
}

func TestLabelsTestSuite(t *testing.T) {
	suite.Run(t, new(LabelsTestSuite))
}

func (suite *LabelsTestSuite) TestWithProjectName() {
	assert := suite.Assert()
	centralhandler.SetCentralHandler(&MockCentralHandlerImpl{})
	infrahandler.Set(&MockInfraHandlerImpl{})
	labels := map[string]string{constant.ProjectLabelKey: "123"}
	labels = WithProjectName(labels)
	assert.Equal(map[string]string{constant.ProjectLabelKey: "123", constant.ProjectNameLabelKey: "123"}, labels)
	labels = WithProjectName(labels)
	assert.Equal(map[string]string{constant.ProjectLabelKey: "123", constant.ProjectNameLabelKey: "123"}, labels)
}

func (suite *LabelsTestSuite) TestWithProjectNameExists() {
	assert := suite.Assert()
	centralhandler.SetCentralHandler(&MockCentralHandlerImpl{})
	infrahandler.Set(&MockInfraHandlerImpl{})
	labels := map[string]string{constant.ProjectLabelKey: "123", constant.ProjectNameLabelKey: "456"}
	labels = WithProjectName(labels)
	assert.Equal(map[string]string{constant.ProjectLabelKey: "123", constant.ProjectNameLabelKey: "123"}, labels)
	labels = WithProjectName(labels)
	assert.Equal(map[string]string{constant.ProjectLabelKey: "123", constant.ProjectNameLabelKey: "123"}, labels)
}

// Lake-reported data carries warehouse in labels but no clusterID.
// WithClusterName should mirror warehouse into clusterNameLabelKey so
// downstream bill export and dashboard aggregates can group Lake costs
// by warehouse. The Lake branch also tags the cost with
// service_plan: "Lake" so downstream aggregates surface a service_plan
// field for Lake costs. The aggregate bucket name (project,
// MSPCustomerID) is the form "warehouse-{provider}-{region}-{name}"
// so summary-label buckets separate Lake costs across cloud providers
// and regions; clusterName stays as just the warehouse name so the
// bill export / CSV "Cluster Name" column is unchanged.
// The fallback only fires for Lake SKUs (databend-*).
func (suite *LabelsTestSuite) TestWithClusterName_LakeWarehouseFallback() {
	assert := suite.Assert()
	centralhandler.SetCentralHandler(&MockCentralHandlerImpl{})
	infrahandler.Set(&MockInfraHandlerImpl{})
	const lakeUnit = "databend-compute-business:aws:us-east-2:large:per-second"

	// Case 1: clusterName is just the warehouse name (CSV uses this),
	// while project / MSPCustomerID use the form
	// "warehouse-{provider}-{region}-{name}". service_plan is "Lake".
	labels := map[string]string{"warehouse": "default"}
	result := WithClusterName(1, lakeUnit, labels)
	assert.Equal("default", result[constant.ClusterNameLabelKey])
	assert.Equal("warehouse-aws-us-east-2-default", result[constant.ProjectNameLabelKey])
	assert.Equal("warehouse-aws-us-east-2-default", result[constant.MSPCustomerIDLabel])
	assert.Equal("default", result["warehouse"]) // original labels are preserved
	assert.Equal("Lake", result[constant.ServicePlanLabelKey])

	// Case 2: clusterName is preset upstream; the fallback must not overwrite it.
	// project / MSPCustomerID are still mirrored to the aggregated bucket name.
	labels2 := map[string]string{"warehouse": "default", constant.ClusterNameLabelKey: "preset-name"}
	result2 := WithClusterName(1, lakeUnit, labels2)
	assert.Equal("preset-name", result2[constant.ClusterNameLabelKey])
	assert.Equal("warehouse-aws-us-east-2-default", result2[constant.ProjectNameLabelKey])
	assert.Equal("warehouse-aws-us-east-2-default", result2[constant.MSPCustomerIDLabel])
	assert.Equal("Lake", result2[constant.ServicePlanLabelKey])

	// Case 3: empty warehouse string does not write clusterName,
	// project / MSPCustomerID fall back to "Lake" (no provider/region
	// when there is no warehouse). service_plan is still "Lake".
	labels3 := map[string]string{"warehouse": ""}
	result3 := WithClusterName(1, lakeUnit, labels3)
	_, ok := result3[constant.ClusterNameLabelKey]
	assert.False(ok)
	assert.Equal("Lake", result3[constant.ProjectNameLabelKey])
	assert.Equal("Lake", result3[constant.MSPCustomerIDLabel])
	assert.Equal("Lake", result3[constant.ServicePlanLabelKey])

	// Case 4: missing warehouse key does not write clusterName,
	// project / MSPCustomerID fall back to "Lake".
	labels4 := map[string]string{}
	result4 := WithClusterName(1, lakeUnit, labels4)
	_, ok = result4[constant.ClusterNameLabelKey]
	assert.False(ok)
	assert.Equal("Lake", result4[constant.ProjectNameLabelKey])
	assert.Equal("Lake", result4[constant.MSPCustomerIDLabel])
	assert.Equal("Lake", result4[constant.ServicePlanLabelKey])

	// Case 5: project/MSPCustomerID preset upstream; the fallback must
	// not overwrite them. clusterName is still mirrored to the warehouse.
	labels5 := map[string]string{
		"warehouse":                  "default",
		constant.ProjectNameLabelKey: "upstream-project",
		"MSPCustomerID":              "upstream-customer",
	}
	result5 := WithClusterName(1, lakeUnit, labels5)
	assert.Equal("default", result5[constant.ClusterNameLabelKey])
	assert.Equal("upstream-project", result5[constant.ProjectNameLabelKey])
	assert.Equal("upstream-customer", result5[constant.MSPCustomerIDLabel])
}

// Multi-cluster warehouses share the same warehouse value across
// warehouseCluster entries, so all small clusters must collapse to the
// same clusterName on bills. All entries carry service_plan: "Lake"
// and the project / MSPCustomerID summary-label entries use the
// aggregated bucket name "warehouse-{provider}-{region}-{name}".
func (suite *LabelsTestSuite) TestWithClusterName_LakeWarehouseMultiCluster() {
	assert := suite.Assert()
	centralhandler.SetCentralHandler(&MockCentralHandlerImpl{})
	infrahandler.Set(&MockInfraHandlerImpl{})
	const lakeUnit = "databend-compute-business:aws:us-east-2:large:per-second"

	labels1 := map[string]string{"warehouse": "aaaa", "warehouseCluster": "cl0001"}
	labels2 := map[string]string{"warehouse": "aaaa", "warehouseCluster": "cl0002"}
	r1 := WithClusterName(1, lakeUnit, labels1)
	r2 := WithClusterName(1, lakeUnit, labels2)
	assert.Equal("aaaa", r1[constant.ClusterNameLabelKey])
	assert.Equal("aaaa", r2[constant.ClusterNameLabelKey])
	assert.Equal("warehouse-aws-us-east-2-aaaa", r1[constant.ProjectNameLabelKey])
	assert.Equal("warehouse-aws-us-east-2-aaaa", r2[constant.ProjectNameLabelKey])
	assert.Equal("warehouse-aws-us-east-2-aaaa", r1[constant.MSPCustomerIDLabel])
	assert.Equal("warehouse-aws-us-east-2-aaaa", r2[constant.MSPCustomerIDLabel])
	assert.Equal("Lake", r1[constant.ServicePlanLabelKey])
	assert.Equal("Lake", r2[constant.ServicePlanLabelKey])
}

// Different cloud provider / region must produce different aggregate
// bucket names so summary-label buckets separate Lake costs across
// providers and regions.
func (suite *LabelsTestSuite) TestWithClusterName_LakeProviderRegionIsolation() {
	assert := suite.Assert()
	centralhandler.SetCentralHandler(&MockCentralHandlerImpl{})
	infrahandler.Set(&MockInfraHandlerImpl{})
	cases := []struct {
		unit, wantSuffix string
	}{
		{"databend-compute-business:aws:us-east-2:large:per-second", "warehouse-aws-us-east-2-default"},
		{"databend-compute-business:aws:us-west-2:large:per-second", "warehouse-aws-us-west-2-default"},
		{"databend-compute-business:gcp:us-central1:large:per-second", "warehouse-gcp-us-central1-default"},
	}
	for _, tc := range cases {
		labels := map[string]string{"warehouse": "default"}
		result := WithClusterName(1, tc.unit, labels)
		assert.Equal("default", result[constant.ClusterNameLabelKey],
			"clusterName must stay as just the warehouse name for unit %q", tc.unit)
		assert.Equal(tc.wantSuffix, result[constant.ProjectNameLabelKey],
			"project bucket name for unit %q", tc.unit)
		assert.Equal(tc.wantSuffix, result[constant.MSPCustomerIDLabel],
			"MSPCustomerID bucket name for unit %q", tc.unit)
	}
}

// Lake Cloud Service Fee and Lake Storage are intentionally NOT
// handled by the Lake branch — only Lake Compute is. They fall
// through to the clusterID path, which returns early because they
// have no clusterID, so the warehouse / service_plan / summary-label
// mirroring is skipped. This test pins that boundary.
func (suite *LabelsTestSuite) TestWithClusterName_LakeOtherSkusBypassFallback() {
	assert := suite.Assert()
	centralhandler.SetCentralHandler(&MockCentralHandlerImpl{})
	infrahandler.Set(&MockInfraHandlerImpl{})
	cases := []struct {
		name, unit string
	}{
		{"cloud-service-fee", "databend-cloudservice-fee:based-on-warehouse:per-day"},
		{"storage", "databend-storage-write:aws:us-east-2:gb-hour"},
		{"storage-read", "databend-storage-read:aws:us-west-2:gb-hour"},
	}
	for _, tc := range cases {
		labels := map[string]string{"warehouse": "default"}
		result := WithClusterName(1, tc.unit, labels)
		// Labels must be untouched — no clusterName, no service_plan,
		// no project / MSPCustomerID mirroring.
		_, ok := result[constant.ClusterNameLabelKey]
		assert.False(ok, "%s SKUs must not write clusterName", tc.name)
		_, ok = result[constant.ProjectNameLabelKey]
		assert.False(ok, "%s SKUs must not write project", tc.name)
		_, ok = result[constant.MSPCustomerIDLabel]
		assert.False(ok, "%s SKUs must not write MSPCustomerID", tc.name)
		_, ok = result[constant.ServicePlanLabelKey]
		assert.False(ok, "%s SKUs must not tag service_plan", tc.name)
		// Original labels are preserved.
		assert.Equal("default", result["warehouse"], "warehouse label must be preserved for %s", tc.name)
	}
}

// Critical scope test: when unit does not have the databend-compute
// prefix, the warehouse label must never be mirrored into clusterName
// and service_plan must not be tagged "Lake". Includes other Lake
// SKUs (storage, cloudservice-fee) plus TiDB cluster units and an
// empty unit as a safety check.
func (suite *LabelsTestSuite) TestWithClusterName_NonLakeUnitBypassesFallback() {
	assert := suite.Assert()
	centralhandler.SetCentralHandler(&MockCentralHandlerImpl{})
	infrahandler.Set(&MockInfraHandlerImpl{})
	cases := []string{
		"tidb:dedicated:compute:hour",
		"serverless:row-storage:gb-hour",
		"essential:compute:hour",
		"databend-storage-write:aws:us-east-2:gb-hour",
		"databend-cloudservice-fee:based-on-warehouse:per-day",
		"", // safety: empty unit must not panic
	}
	for _, unit := range cases {
		labels := map[string]string{"warehouse": "should-not-leak"}
		result := WithClusterName(1, unit, labels)
		_, ok := result[constant.ClusterNameLabelKey]
		assert.False(ok, "unit %q should not trigger Lake fallback", unit)
		assert.NotEqual("Lake", result[constant.ServicePlanLabelKey],
			"unit %q must not tag service_plan as Lake", unit)
	}
}

// Regression: when a non-Lake unit has both clusterID and warehouse
// in labels, the original clusterID path must take precedence and the
// warehouse value must not leak into clusterName. service_plan must
// also not be tagged "Lake" because the Lake branch is skipped.
func (suite *LabelsTestSuite) TestWithClusterName_LakeFallbackDoesNotAffectClusterIDPath() {
	assert := suite.Assert()
	centralhandler.SetCentralHandler(&MockCentralHandlerImpl{})
	infrahandler.Set(&MockInfraHandlerImpl{})
	labels := map[string]string{constant.ClusterLabelKey: "999", "warehouse": "default"}
	result := WithClusterName(1, "tidb:dedicated:compute:hour", labels)
	// The mock GetCluster returns an empty Cluster{Name:""}, so the
	// clusterID path writes "" — but the warehouse value "default"
	// must never reach clusterName via the Lake fallback.
	assert.NotEqual("default", result[constant.ClusterNameLabelKey])
	assert.NotEqual("Lake", result[constant.ServicePlanLabelKey])
}