package resourcepool

import (
	"context"
	"strings"

	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"

	globalcommonv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/common/v1"
	globalngv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/nextgen/v1"
)

const (
	defaultInitialRcu            uint32 = 20000
	defaultTikvRaftlogDiskSizeGi uint32 = 40

	serverEnvDev     = "dev"
	serverEnvStaging = "staging"
	serverEnvProd    = "prod"
)

// getDefaultDemand returns the default demand for creating a resource pool.
// This is a temporary solution until we have a better way to manage resource pool demands.
func getDefaultDemand(ctx context.Context, provider globalcommonv1.CloudProvider_Name) *globalngv1.CreateResourcePoolRequest_Demand {
	switch provider {
	case globalcommonv1.CloudProvider_AWS:
		return &globalngv1.CreateResourcePoolRequest_Demand{
			InitialRcu: defaultInitialRcu,
			TikvData: &globalngv1.CreateResourcePoolRequest_Storage{
				StorageType:        globalngv1.ResourcePoolProfile_GP3,
				DiskSizeGi:         1500,
				DiskIops:           8000,
				DiskThroughputMbps: 400,
			},
			TikvRaftlog: &globalngv1.CreateResourcePoolRequest_Storage{
				StorageType:        globalngv1.ResourcePoolProfile_GP3,
				DiskSizeGi:         defaultTikvRaftlogDiskSizeGi,
				DiskIops:           3000,
				DiskThroughputMbps: 125,
			},
			TiflashCn: &globalngv1.CreateResourcePoolRequest_Storage{
				StorageType:        globalngv1.ResourcePoolProfile_GP3,
				DiskSizeGi:         500,
				DiskIops:           6000,
				DiskThroughputMbps: 625,
			},
			TiflashWn: &globalngv1.CreateResourcePoolRequest_Storage{
				StorageType:        globalngv1.ResourcePoolProfile_GP3,
				DiskSizeGi:         500,
				DiskIops:           6000,
				DiskThroughputMbps: 625,
			},
		}
	case globalcommonv1.CloudProvider_ALICLOUD:
		return &globalngv1.CreateResourcePoolRequest_Demand{
			InitialRcu: defaultInitialRcu,
			TikvData: &globalngv1.CreateResourcePoolRequest_Storage{
				StorageType:        globalngv1.ResourcePoolProfile_CloudAuto,
				DiskSizeGi:         1500,
				DiskIops:           8000,
				DiskThroughputMbps: 0,
			},
			TikvRaftlog: &globalngv1.CreateResourcePoolRequest_Storage{
				StorageType:        globalngv1.ResourcePoolProfile_CloudAuto,
				DiskSizeGi:         defaultTikvRaftlogDiskSizeGi,
				DiskIops:           3000,
				DiskThroughputMbps: 0,
			},
			TiflashCn: &globalngv1.CreateResourcePoolRequest_Storage{
				StorageType:        globalngv1.ResourcePoolProfile_CloudAuto,
				DiskSizeGi:         500,
				DiskIops:           6000,
				DiskThroughputMbps: 0,
			},
			TiflashWn: &globalngv1.CreateResourcePoolRequest_Storage{
				StorageType:        globalngv1.ResourcePoolProfile_CloudAuto,
				DiskSizeGi:         500,
				DiskIops:           6000,
				DiskThroughputMbps: 0,
			},
		}
	case globalcommonv1.CloudProvider_AZURE:
		return &globalngv1.CreateResourcePoolRequest_Demand{
			InitialRcu: defaultInitialRcu,
			TikvData: &globalngv1.CreateResourcePoolRequest_Storage{
				StorageType:        globalngv1.ResourcePoolProfile_PremiumSSDV2,
				DiskSizeGi:         1500,
				DiskIops:           8000,
				DiskThroughputMbps: 400,
			},
			TikvRaftlog: &globalngv1.CreateResourcePoolRequest_Storage{
				StorageType:        globalngv1.ResourcePoolProfile_PremiumSSDV2,
				DiskSizeGi:         defaultTikvRaftlogDiskSizeGi,
				DiskIops:           3000,
				DiskThroughputMbps: 125,
			},
			TiflashCn: &globalngv1.CreateResourcePoolRequest_Storage{
				StorageType:        globalngv1.ResourcePoolProfile_PremiumSSDV2,
				DiskSizeGi:         500,
				DiskIops:           6000,
				DiskThroughputMbps: 625,
			},
			TiflashWn: &globalngv1.CreateResourcePoolRequest_Storage{
				StorageType:        globalngv1.ResourcePoolProfile_PremiumSSDV2,
				DiskSizeGi:         500,
				DiskIops:           6000,
				DiskThroughputMbps: 625,
			},
		}
	default:
		tilog.Warn(ctx, "unsupported provider, return empty demand", zap.String("provider", provider.String()))
	}
	return &globalngv1.CreateResourcePoolRequest_Demand{}
}

// getRcuDrivenDefaultDemand applies RCU-driven storage defaults only for known server envs.
// Unknown envs intentionally fall back to provider-only defaults for compatibility.
func getRcuDrivenDefaultDemand(ctx context.Context,
	provider globalcommonv1.CloudProvider_Name,
	env string,
	initialRcu uint32,
) *globalngv1.CreateResourcePoolRequest_Demand {
	demand := getDefaultDemand(ctx, provider)
	demand.InitialRcu = initialRcu

	switch strings.ToLower(strings.TrimSpace(env)) {
	case serverEnvDev, serverEnvStaging:
		setStorageDefault(demand.TikvData, 100, 3000, 125)
		setStorageDefault(demand.TikvRaftlog, defaultTikvRaftlogDiskSizeGi, 3000, 125)
		setTiflashDefault(demand)
	case serverEnvProd:
		// RCU boundary values belong to the next tier, so 50000 uses the higher storage default.
		if initialRcu < 50000 {
			setStorageDefault(demand.TikvData, 500, 3000, 125)
			setStorageDefault(demand.TikvRaftlog, defaultTikvRaftlogDiskSizeGi, 3000, 125)
		} else {
			setStorageDefault(demand.TikvData, 1000, 4000, 300)
			setStorageDefault(demand.TikvRaftlog, defaultTikvRaftlogDiskSizeGi, 4000, 125)
		}
		setTiflashDefault(demand)
	default:
		tilog.Warn(ctx, "unsupported server env, use provider default demand",
			zap.String("env", env),
			zap.String("provider", provider.String()),
		)
	}

	return demand
}

func setTiflashDefault(demand *globalngv1.CreateResourcePoolRequest_Demand) {
	setStorageDefault(demand.TiflashCn, 500, 3000, 125)
	setStorageDefault(demand.TiflashWn, 500, 3000, 125)
}

func setStorageDefault(storage *globalngv1.CreateResourcePoolRequest_Storage,
	sizeGi uint32,
	iops uint32,
	throughputMbps uint32,
) {
	if storage == nil {
		return
	}
	storage.DiskSizeGi = sizeGi
	storage.DiskIops = iops
	storage.DiskThroughputMbps = throughputMbps
}

func getDefaultCidr(ctx context.Context) string {
	return "10.0.0.0/16"
}
func getDefaultK8sServiceCidr(ctx context.Context) string {
	return "172.20.0.0/16"
}

// getDefaultInstanceLimit returns the default instance limit
func getDefaultInstanceLimit(ctx context.Context) uint64 {
	return 50
}

// fillingCreateRequestWithDefault
func fillingCreateRequestWithDefault(ctx context.Context,
	request *globalngv1.CreateResourcePoolRequest,
	env string,
	useRcuDrivenDefault bool,
) {
	provider := request.GetResourcePool().GetProviderId()
	defaultDemand := getDefaultDemand(ctx, provider)

	if request.ResourcePool.Cidr == "" {
		request.ResourcePool.Cidr = getDefaultCidr(ctx)
	}

	if request.ResourcePool.K8SServiceCidr == "" {
		request.ResourcePool.K8SServiceCidr = getDefaultK8sServiceCidr(ctx)
	}

	if request.ResourcePool.InstanceLimit <= 0 {
		request.ResourcePool.InstanceLimit = getDefaultInstanceLimit(ctx)
	}

	if request.Demand == nil {
		request.Demand = &globalngv1.CreateResourcePoolRequest_Demand{}
	}

	if request.Demand.InitialRcu == 0 && defaultDemand.InitialRcu != 0 {
		request.Demand.InitialRcu = defaultDemand.InitialRcu
	}

	if useRcuDrivenDefault {
		defaultDemand = getRcuDrivenDefaultDemand(ctx, provider, env, request.Demand.InitialRcu)
	}

	ensureDemandStorages(request.Demand)

	// The IO2 5000 IOPS default is a TiKV data-only special case.
	if request.Demand.TikvData.StorageType == globalngv1.ResourcePoolProfile_IO2 && defaultDemand.TikvData != nil {
		defaultDemand.TikvData.DiskIops = 5000
	}

	fillingAndCorrect(provider, request.Demand.TikvData, defaultDemand.TikvData)
	fillingAndCorrect(provider, request.Demand.TikvRaftlog, defaultDemand.TikvRaftlog)
	fillingAndCorrect(provider, request.Demand.TiflashCn, defaultDemand.TiflashCn)
	fillingAndCorrect(provider, request.Demand.TiflashWn, defaultDemand.TiflashWn)
}

func ensureDemandStorages(demand *globalngv1.CreateResourcePoolRequest_Demand) {
	if demand.TikvData == nil {
		demand.TikvData = &globalngv1.CreateResourcePoolRequest_Storage{}
	}
	if demand.TikvRaftlog == nil {
		demand.TikvRaftlog = &globalngv1.CreateResourcePoolRequest_Storage{}
	}
	if demand.TiflashCn == nil {
		demand.TiflashCn = &globalngv1.CreateResourcePoolRequest_Storage{}
	}
	if demand.TiflashWn == nil {
		demand.TiflashWn = &globalngv1.CreateResourcePoolRequest_Storage{}
	}
}

func fillingAndCorrect(provider globalcommonv1.CloudProvider_Name,
	storage *globalngv1.CreateResourcePoolRequest_Storage,
	defaultStorage *globalngv1.CreateResourcePoolRequest_Storage) {
	fillingEmptyValues(defaultStorage, storage)
	correctThroughput(provider, storage)
	correctSize(storage)
	correctIops(storage)
}

func fillingEmptyValues(source *globalngv1.CreateResourcePoolRequest_Storage,
	target *globalngv1.CreateResourcePoolRequest_Storage) {
	if source == nil || target == nil {
		return
	}
	if target.StorageType == globalngv1.ResourcePoolProfile_STORAGE_TYPE_UNSPECIFIED {
		target.StorageType = source.StorageType
	}
	if target.DiskSizeGi == 0 {
		target.DiskSizeGi = source.DiskSizeGi
	}
	if target.DiskIops == 0 {
		target.DiskIops = source.DiskIops
	}
	if target.DiskThroughputMbps == 0 {
		target.DiskThroughputMbps = source.DiskThroughputMbps
	}
}

func correctThroughput(provider globalcommonv1.CloudProvider_Name, storage *globalngv1.CreateResourcePoolRequest_Storage) {
	if storage == nil {
		return
	}
	if provider == globalcommonv1.CloudProvider_ALICLOUD {
		storage.DiskThroughputMbps = 0
	}
	if provider == globalcommonv1.CloudProvider_AWS &&
		storage.StorageType == globalngv1.ResourcePoolProfile_IO2 {
		storage.DiskThroughputMbps = 0
	}
	if storage.StorageType == globalngv1.ResourcePoolProfile_LocalSSD {
		storage.DiskThroughputMbps = 0
	}
}

func correctSize(storage *globalngv1.CreateResourcePoolRequest_Storage) {
	if storage == nil {
		return
	}
	if storage.StorageType == globalngv1.ResourcePoolProfile_LocalSSD {
		storage.DiskSizeGi = 0
	}
}

func correctIops(storage *globalngv1.CreateResourcePoolRequest_Storage) {
	if storage == nil {
		return
	}
	if storage.StorageType == globalngv1.ResourcePoolProfile_LocalSSD {
		storage.DiskIops = 0
	}
}