package rcuquota

import (
	"context"
	"errors"
	"strings"

	"github.com/tidbcloud/pkgv2/db"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"gorm.io/gorm"

	globalcommonv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/common/v1"
	globalngv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/nextgen/v1"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/entity"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/domain/repository/rcuquota"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/infra/repoimpl/schema"
	"github.com/tidbcloud/tidb-management-service/nextgen-global-service/pkg/config"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
)

type rcuQuotaRepoImpl struct {
	dbProvider               db.Provider
	defaultRcuQuota          globalngv1.RcuQuota
	defaultRcuQuotaOverrides map[string]*globalngv1.RcuQuota
}

func NewRcuQuotaRepo(p db.Provider, cfg *config.Config) rcuquota.Repo {
	return &rcuQuotaRepoImpl{
		dbProvider: p,
		defaultRcuQuota: globalngv1.RcuQuota{
			InstanceMaxRcu:             uint64(cfg.BizConfig.MaxRCUQuotaPerCluster),
			TotalMaxRcu:                uint64(cfg.BizConfig.TotalMaxRCU),
			BurstyRatio:                cfg.BizConfig.BurstyRatio,
			IdleResourcePoolTtlInHours: uint64(cfg.BizConfig.IdleResourcePoolTtlInHours),
			VcpusLimitPerResourcePool:  uint64(cfg.BizConfig.VcpusLimitPerResourcePool),
			MinimalRcuBaseline:         uint64(cfg.BizConfig.MinimalRcuBaseline),
			MinimalRcuCoefficient:      cfg.BizConfig.MinimalRcuCoefficient,
			ScaleDownCooldownMinutes:   uint64(cfg.BizConfig.ScaleDownCooldownMinutes),
		},
		defaultRcuQuotaOverrides: defaultRcuQuotaOverrides(&cfg.BizConfig),
	}
}

func defaultRcuQuotaOverrides(cfg *config.BizConfig) map[string]*globalngv1.RcuQuota {
	overrides := map[string]*globalngv1.RcuQuota{}
	for servicePlanKey, quotaConfig := range cfg.RcuQuotaConfigs {
		servicePlan := canonicalServicePlanKey(servicePlanKey)
		quota := &globalngv1.RcuQuota{
			InstanceMaxRcu:             uint64(cfg.MaxRCUQuotaPerCluster),
			TotalMaxRcu:                uint64(cfg.TotalMaxRCU),
			BurstyRatio:                cfg.BurstyRatio,
			IdleResourcePoolTtlInHours: uint64(cfg.IdleResourcePoolTtlInHours),
			VcpusLimitPerResourcePool:  uint64(cfg.VcpusLimitPerResourcePool),
			MinimalRcuBaseline:         uint64(cfg.MinimalRcuBaseline),
			MinimalRcuCoefficient:      cfg.MinimalRcuCoefficient,
			ScaleDownCooldownMinutes:   uint64(cfg.ScaleDownCooldownMinutes),
		}
		applyRcuQuotaConfig(quota, &quotaConfig)
		overrides[servicePlan] = quota
	}
	return overrides
}

func canonicalServicePlanKey(servicePlanKey string) string {
	if _, ok := globalcommonv1.ServicePlan_value[servicePlanKey]; ok {
		return servicePlanKey
	}
	for canonicalKey := range globalcommonv1.ServicePlan_value {
		if strings.EqualFold(canonicalKey, servicePlanKey) {
			return canonicalKey
		}
	}
	return servicePlanKey
}

func applyRcuQuotaConfig(quota *globalngv1.RcuQuota, quotaConfig *config.RcuQuotaConfig) {
	if quotaConfig.MaxRCUQuotaPerCluster != nil {
		quota.InstanceMaxRcu = uint64(*quotaConfig.MaxRCUQuotaPerCluster)
	}
	if quotaConfig.TotalMaxRCU != nil {
		quota.TotalMaxRcu = uint64(*quotaConfig.TotalMaxRCU)
	}
	if quotaConfig.MinimalRcuBaseline != nil {
		quota.MinimalRcuBaseline = uint64(*quotaConfig.MinimalRcuBaseline)
	}
	if quotaConfig.MinimalRcuCoefficient != nil {
		quota.MinimalRcuCoefficient = *quotaConfig.MinimalRcuCoefficient
	}
	if quotaConfig.BurstyRatio != nil {
		quota.BurstyRatio = *quotaConfig.BurstyRatio
	}
	if quotaConfig.VcpusLimitPerResourcePool != nil {
		quota.VcpusLimitPerResourcePool = uint64(*quotaConfig.VcpusLimitPerResourcePool)
	}
	if quotaConfig.IdleResourcePoolTtlInHours != nil {
		quota.IdleResourcePoolTtlInHours = uint64(*quotaConfig.IdleResourcePoolTtlInHours)
	}
	if quotaConfig.ScaleDownCooldownMinutes != nil {
		quota.ScaleDownCooldownMinutes = uint64(*quotaConfig.ScaleDownCooldownMinutes)
	}
}

func (r *rcuQuotaRepoImpl) Save(ctx context.Context, en *entity.RcuQuota) bizerr.BizError {
	if en == nil {
		return bizerr.New(bizerr.ErrCodeParams, "rcu quota entity is nil")
	}
	if en.OrganizationId == 0 {
		return bizerr.New(bizerr.ErrCodeParams, "tenant ID is required")
	}

	wdb := r.dbProvider.WriteDB(ctx).Model(&schema.RcuQuota{})
	schemaQuota := entityToSchema(en)

	result := wdb.Where("tenant_id = ? and service_plan = ?", en.OrganizationId, en.ServicePlan.String()).Updates(schemaQuota)
	if result.Error != nil {
		return bizerr.WrapGormError(result.Error, "save tenant RCU quota")
	}
	if result.RowsAffected == 0 {
		// If no rows were updated, it means the quota does not exist, so create it
		if err := wdb.Create(schemaQuota).Error; err != nil {
			return bizerr.WrapGormError(err, "create tenant RCU quota")
		}
	}

	return nil
}

func (r *rcuQuotaRepoImpl) Load(ctx context.Context, tenantID uint64, servicePlan string) (*entity.RcuQuota, bizerr.BizError) {
	quota, err := r.getFromDB(ctx, tenantID, servicePlan)
	if err != nil {
		if err.ErrCode() == bizerr.ErrCodeRecordNotFound {
			tilog.Info(ctx, "rcu quota not found, will return default quota",
				zap.Uint64("tenant id", tenantID),
				zap.String("service plan", servicePlan))
			return r.buildRcuQuotaForTenant(ctx, tenantID, servicePlan)
		}
		return nil, err
	}
	return quota, nil
}

func (r *rcuQuotaRepoImpl) getFromDB(ctx context.Context, tenantID uint64, servicePlan string) (*entity.RcuQuota, bizerr.BizError) {
	rdb := r.dbProvider.ReadDB(ctx)
	var quota *schema.RcuQuota

	if err := rdb.Where("tenant_id = ? and service_plan = ?", tenantID, servicePlan).First(&quota).Error; err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return nil, bizerr.New(bizerr.ErrCodeRecordNotFound, "rcu quota not found")
		}

		return nil, bizerr.WrapGormError(err, "get tenant quota")
	}

	return schemaToEntity(quota), nil
}

func (r *rcuQuotaRepoImpl) buildRcuQuotaForTenant(ctx context.Context, tenantID uint64, servicePlan string) (*entity.RcuQuota, bizerr.BizError) {
	tilog.Debug(ctx, "init rcu quota with default value", zap.Uint64("tenant id", tenantID), zap.String("service plan", servicePlan))
	defaultRcuQuota := r.defaultRcuQuotaForServicePlan(servicePlan)
	return &entity.RcuQuota{
		OrganizationId:             tenantID,
		ServicePlan:                globalcommonv1.ServicePlan(globalcommonv1.ServicePlan_value[servicePlan]),
		InstanceMaxRcu:             defaultRcuQuota.InstanceMaxRcu,
		TotalMaxRcu:                defaultRcuQuota.TotalMaxRcu,
		BurstyRatio:                defaultRcuQuota.BurstyRatio,
		IdleResourcePoolTtlInHours: defaultRcuQuota.IdleResourcePoolTtlInHours,
		VcpusLimitPerResourcePool:  defaultRcuQuota.VcpusLimitPerResourcePool,
		MinimalRcuBaseline:         defaultRcuQuota.MinimalRcuBaseline,
		MinimalRcuCoefficient:      defaultRcuQuota.MinimalRcuCoefficient,
		ScaleDownCooldownMinutes:   defaultRcuQuota.ScaleDownCooldownMinutes,
	}, nil
}

func (r *rcuQuotaRepoImpl) defaultRcuQuotaForServicePlan(servicePlan string) *globalngv1.RcuQuota {
	if quota, ok := r.defaultRcuQuotaOverrides[servicePlan]; ok {
		return quota
	}
	return &r.defaultRcuQuota
}