package labels

import (
	"errors"
	"fmt"
	"strconv"
	"strings"

	"github.com/pingcap/dbaas/billing2/handler/tidbmgmtsvchandler"

	"github.com/tidbcloud/account/idl/pbgen/proto/account"

	"github.com/pingcap/dbaas/billing2/config"

	"github.com/avast/retry-go"
	"github.com/pingcap/dbaas/billing2/billing_errors"
	"github.com/pingcap/dbaas/billing2/constant"
	"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/utils/log"
	"go.uber.org/zap"
)

func CopyMap(in map[string]string) (out map[string]string) {
	out = make(map[string]string, len(in))
	for key, val := range in {
		out[key] = val
	}
	return
}

func GetProjectName(projectID uint) (string, error) {
	projectInfo, err := GetProjectInfo(projectID)
	if err != nil {
		return "", err
	}
	return projectInfo.Name, nil
}

func GetProjectInfo(projectID uint) (*forms.Project, error) {
	cache := cachehandler.Get()
	var project forms.Project
	found := cache.Get("GetProjectInfo", &project, projectID)
	if found {
		log.Info("hits project info cache", zap.Uint("projectID", projectID), zap.Any("project", project))
		return &project, nil
	}
	central := centralhandler.GetCentralHandler()
	projectInfo, err := central.GetProjectInfo(projectID)
	if err != nil {
		return nil, err
	}
	cache.Set("GetProjectInfo", *projectInfo, projectID)
	return projectInfo, err
}

func GetClusterWithCache(tenantID, clusterID uint) (*forms.Cluster, error) {
	objNotExisted := false
	cache := cachehandler.Get()
	var cacheResult struct {
		Cluster       *forms.Cluster
		ObjNotExisted bool
	}
	var err error
	var cluster *forms.Cluster
	funcName := "GetCluster"
	found := cache.Get(funcName, &cacheResult, tenantID, clusterID)
	if found {
		objNotExisted = cacheResult.ObjNotExisted
		cluster = cacheResult.Cluster
	} else {
		if config.Get().GetBool(config.EnableUseMgmtGetClusterInfo) {
			mgmtClient, err := tidbmgmtsvchandler.Get()
			if err != nil {
				log.Error("fail to GetClusterWithCache get mgmt client error", zap.Uint("clusterID", clusterID), zap.Uint("tenantID", tenantID), zap.Error(err))
				return nil, err
			}
			cluster, err = mgmtClient.GetClusterByID(uint64(tenantID), uint64(clusterID))
			if err != nil {
				if errors.Is(err, billing_errors.ObjectNotExistsError) {
					objNotExisted = true
				} else {
					return nil, err
				}
			}
		} else {
			cluster, err = infrahandler.Get().GetCluster(tenantID, uint(clusterID))
			if err != nil {
				if errors.Is(err, billing_errors.ObjectNotExistsError) {
					objNotExisted = true
				} else {
					return nil, err
				}
			}
		}

		cacheResult.Cluster = cluster
		cacheResult.ObjNotExisted = objNotExisted
		cache.Set(funcName, cacheResult, tenantID, clusterID)
	}
	if objNotExisted {
		return nil, billing_errors.ObjectNotExistsError
	}
	return cluster, nil
}

func WithProjectName(in map[string]string) (out map[string]string) {
	log.Info("start to add projectName")
	out = CopyMap(in)
	projectIDStr, ok := out[constant.ProjectLabelKey]
	if !ok {
		log.Info("projectID not existed, ignore it")
		return
	}
	projectID, err := strconv.ParseUint(projectIDStr, 10, 64)
	if err != nil || projectID < 1 {
		log.Error("failed to parse projectID", zap.Error(err), zap.String("projectID", projectIDStr), zap.Any("labels", in))
		return
	}
	var projectName string
	err = retry.Do(
		func() (err error) {
			projectName, err = GetProjectName(uint(projectID))
			return err
		},
		retry.RetryIf(func(err error) bool { return !errors.Is(err, billing_errors.ObjectNotExistsError) }),
	)
	if err != nil {
		log.Error("failed to get projectName", zap.Error(err), zap.String("projectID", projectIDStr), zap.Any("labels", in))
		return
	}
	if rawProjectName, ok := out[constant.ProjectNameLabelKey]; projectName == "" && ok && rawProjectName != "" {
		log.Warn("get empty project name", zap.String("projectID", projectIDStr))
		return
	}
	log.Info("success to get project name", zap.String("projectID", projectIDStr), zap.String("projectName", projectName))
	out[constant.ProjectNameLabelKey] = projectName
	return
}

// isLakeComputeUnit reports whether the unit belongs to the Lake
// Compute SKU family. Defined by the "databend-compute" prefix —
// see LakeComputeUnitPrefix. Other Lake SKUs (storage, cloud service
// fee) intentionally fall through to the default clusterID path.
func isLakeComputeUnit(unit string) bool {
	return strings.HasPrefix(unit, constant.LakeComputeUnitPrefix)
}

// parseLakeProviderRegion extracts the provider and region from a
// Lake Compute unit string. The unit format is:
//
//	databend-compute-<sku>:<provider>:<region>:<size>:<unit>
//
// e.g. "databend-compute-business:aws:us-east-2:large:per-second".
// Returns ok=false for any malformed unit. Callers fall back to the
// bucket name without provider/region in that case.
func parseLakeProviderRegion(unit string) (provider, region string, ok bool) {
	parts := strings.Split(unit, ":")
	if len(parts) < 4 {
		return "", "", false
	}
	return parts[1], parts[2], true
}

func WithClusterName(tenantID uint, unit string, in map[string]string) (out map[string]string) {
	log.Info("start to add clusterName")
	out = CopyMap(in)

	// Lake units (databend-*) have no clusterID; mirror the warehouse
	// name into clusterName so bill export and dashboard aggregates
	// roll up Lake costs by warehouse. Multi-cluster warehouses share
	// the same warehouse value across warehouseCluster entries, so all
	// small clusters collapse to the same clusterName. Existing
	// clusterName is preserved. We branch here — and skip the clusterID
	// lookup entirely — so unrelated units that happen to carry a
	// warehouse label never enter this path.
	if isLakeComputeUnit(unit) {
		// Bucket name used for the summary-label aggregation. Falls back
		// to "Lake" when the cost has no warehouse so the bucket is
		// still distinguishable from the "Other Charges" catch-all.
		bucketName := constant.ServicePlanLake
		if warehouse, ok := out[constant.WarehouseLabelKey]; ok && warehouse != "" {
			bucketName = warehouse
			if existing, ok := out[constant.ClusterNameLabelKey]; !ok || existing == "" {
				log.Info("lake warehouse fallback: copy warehouse to clusterName",
					zap.String("unit", unit),
					zap.String("warehouse", warehouse))
				out[constant.ClusterNameLabelKey] = warehouse
			}
		}
		// Tag Lake costs with the service plan so downstream aggregates
		// surface service_plan: "Lake" alongside the per-warehouse
		// clusterName. This labels the cost as belonging to the Lake
		// service tier for UI grouping and analytics.
		out[constant.ServicePlanLabelKey] = constant.ServicePlanLake
		// Aggregate bucket name is the form "warehouse-{provider}-
		// {region}-{name}" where the leading "warehouse" is a literal
		// prefix marking the bucket as a Lake warehouse bucket,
		// provider and region come from the unit, and {name} is the
		// warehouse name. This separates Lake costs across cloud
		// providers and regions in summary-label buckets (project,
		// MSPCustomerID). The clusterName label stays as just the
		// warehouse name so the bill export / CSV "Cluster Name"
		// column is unchanged. When the cost has no warehouse we fall
		// back to the plain bucketName (e.g. "Lake") — provider/region
		// only matter when they belong to a specific warehouse.
		aggBucketName := bucketName
		if warehouse, ok := out[constant.WarehouseLabelKey]; ok && warehouse != "" {
			if provider, region, ok := parseLakeProviderRegion(unit); ok {
				aggBucketName = fmt.Sprintf("warehouse-%s-%s-%s", provider, region, warehouse)
			} else {
				aggBucketName = fmt.Sprintf("warehouse-%s", warehouse)
			}
		}
		// Mirror aggBucketName into the summary-label entries (project,
		// MSPCustomerID) so bill export and dashboard aggregates do not
		// roll Lake costs into the "Other Charges" catch-all bucket
		// under the configured summary labels. Existing values are
		// preserved so upstream labels are never overwritten.
		if _, ok := out[constant.ProjectNameLabelKey]; !ok {
			out[constant.ProjectNameLabelKey] = aggBucketName
		}
		if _, ok := out[constant.MSPCustomerIDLabel]; !ok {
			out[constant.MSPCustomerIDLabel] = aggBucketName
		}
		return
	}

	clusterIDStr, ok := out[constant.ClusterLabelKey]
	if !ok {
		log.Info("clusterID not existed, ignore it")
		return
	}
	clusterID, err := strconv.ParseUint(clusterIDStr, 10, 64)
	if err != nil {
		log.Error("failed to parse clusterID", zap.Error(err), zap.String("clusterID", clusterIDStr))
		return
	}
	var customerID uint64
	customerIDStr, customerTransferredCost := out[constant.MSPCustomerIDLabel]
	if customerTransferredCost {
		customerID, err = strconv.ParseUint(customerIDStr, 10, 64)
		if err != nil {
			log.Error("failed to parse customerID", zap.Error(err), zap.String("customerID", customerIDStr))
			return
		}
	}
	var cluster *forms.Cluster
	err = retry.Do(
		func() (err error) {
			if customerTransferredCost {
				cluster, err = GetClusterWithCache(uint(customerID), uint(clusterID))
			} else {
				cluster, err = GetClusterWithCache(tenantID, uint(clusterID))
			}
			return err
		},
		retry.RetryIf(func(err error) bool { return !errors.Is(err, billing_errors.ObjectNotExistsError) }),
	)
	if err != nil {
		log.Warn("failed to get clusterName", zap.Error(err), zap.String("clusterID", clusterIDStr))
		return
	}
	if rawClusterName, ok := out[constant.ClusterNameLabelKey]; cluster.Name == "" && ok && rawClusterName != "" {
		log.Warn("get empty clusterName", zap.String("clusterID", clusterIDStr))
		return
	}
	log.Info("success to get clusterName", zap.String("clusterID", clusterIDStr), zap.String("clusterName", cluster.Name))
	out[constant.ClusterNameLabelKey] = cluster.Name
	_, ok = out[constant.ProjectLabelKey]
	if !ok && cluster.DeployType == 3 {
		out[constant.ProjectNameLabelKey] = cluster.Name
		if cluster.ServicePlan != "" {
			out[constant.ServicePlanLabelKey] = cluster.ServicePlan
		} else {
			out[constant.ServicePlanLabelKey] = constant.ServicePlanPremium
		}
	}

	projectIDStr, hasProjectID := out[constant.ProjectLabelKey]
	if hasProjectID && projectIDStr != "" {
		projectID, err := strconv.ParseUint(projectIDStr, 10, 64)
		if err != nil {
			log.Error("failed to parse projectID", zap.Error(err), zap.String("projectID", projectIDStr))
		} else {
			projectInfo, err := GetProjectInfo(uint(projectID))
			if err != nil {
				log.Error("failed to get project info", zap.Error(err), zap.String("projectID", projectIDStr))
			} else if projectInfo.Type == int32(account.ProjectType_PROJECT_TYPE_TIDBXVIRTUAL) {
				log.Info("project type is TIDBX, replace project name with cluster name",
					zap.String("projectID", projectIDStr),
					zap.String("clusterName", cluster.Name))
				out[constant.ProjectNameLabelKey] = cluster.Name
				if cluster.ServicePlan != "" {
					out[constant.ServicePlanLabelKey] = cluster.ServicePlan
				} else {
					out[constant.ServicePlanLabelKey] = constant.ServicePlanPremium
				}
			}
		}
	}

	return
}