package spu

import (
	"cmp"
	"context"
	"encoding/json"
	"slices"

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

	globalcommonv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/common/v1"
	globaldedicatedv1 "github.com/tidbcloud/tidb-management-service/api/spec/global/dedicated/v1"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/entity"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/repository/spu"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
)

//go:generate go run go.uber.org/mock/mockgen@latest -source=$GOFILE -package=mock_$GOPACKAGE -destination=./mock_$GOPACKAGE/mock_$GOFILE
type SkuService interface {
	// SelectSkusByComponentType selects skus for a component type which is based on the spu, or just the pd sku based on the condition of tikv sku.
	SelectSkusByComponentType(
		ctx context.Context,
		componentType entity.ComponentType,
		request *SelectSkusRequest,
	) ([]*entity.Sku, bizerr.BizError)
	// SelectSkusWithPrefetchedData selects skus for a spu. The pre-fetched data are passed in.
	SelectSkusWithPrefetchedData(
		ctx context.Context,
		prefetched *PrefetchedData,
		request *SelectSkusRequest,
	) ([]*entity.Sku, bizerr.BizError)

	// CRUDs.
	GetSkuByKey(ctx context.Context, key string) (*entity.Sku, bizerr.BizError)
	GetSku(ctx context.Context, req *globaldedicatedv1.GetSkuRequest) (*entity.Sku, bizerr.BizError)
	ListSkus(
		ctx context.Context,
		req *globaldedicatedv1.ListSkusRequest,
	) (*globaldedicatedv1.ListSkusResponse, bizerr.BizError)
	DeleteSku(ctx context.Context, id uint64) bizerr.BizError
	CreateSku(ctx context.Context, req *globaldedicatedv1.Sku) (*entity.Sku, bizerr.BizError)
	UpdateSku(ctx context.Context, req *globaldedicatedv1.BatchUpsertSkusRequest_Update) (*entity.Sku, bizerr.BizError)
}

type SelectSkusRequest struct {
	// Required filters.
	TenantId  uint64
	ProjectId uint64
	Provider  globalcommonv1.CloudProvider_Name
	Region    globalcommonv1.Region_Name

	// Optional filters. If not set, then the filter is ignored.
	VolumeType     *globaldedicatedv1.StorageMu_VolumeType
	RaftVolumeType *globaldedicatedv1.StorageMu_VolumeType
	CpuArch        *globaldedicatedv1.InstanceMu_CpuArch
	CpuGeneration  *globaldedicatedv1.InstanceMu_CpuGeneration

	// Three types of request:
	//   1. TiKVComponent not nil: select **PD** skus based on PD-TiKV mapping.
	//   2. SpuName not nil: select skus based on the spu.
	//   3. TiKVComponent and SpuKey both nil: select all available skus.
	TiKVComponent *entity.DesiredClusterComponent
	SpuName       *string
}

type PrefetchedData struct {
	Spu         *entity.Spu
	SpuMappings []*entity.SpuMapping
	SkuMap      map[string]*entity.Sku
}

func NewSkuService(
	spuRepo spu.RepoSpu,
	skuRepo spu.SkuRepo,
	instanceMuRepo spu.InstanceMuRepo,
	storageMuRepo spu.StorageMuRepo,
) SkuService {
	return &skuServiceImpl{
		spuRepo:        spuRepo,
		skuRepo:        skuRepo,
		instanceMuRepo: instanceMuRepo,
		storageMuRepo:  storageMuRepo,
	}
}

type skuServiceImpl struct {
	spuRepo        spu.RepoSpu
	skuRepo        spu.SkuRepo
	instanceMuRepo spu.InstanceMuRepo
	storageMuRepo  spu.StorageMuRepo
}

func (s *skuServiceImpl) SelectSkusByComponentType(
	ctx context.Context,
	componentType entity.ComponentType,
	request *SelectSkusRequest,
) ([]*entity.Sku, bizerr.BizError) {
	ctx, span := tracer.Start(ctx, "Domain.SelectSkusByComponentType")
	defer span.End()

	if request.TenantId == 0 {
		return nil, bizerr.New(bizerr.ErrCodeParams, "missing tenant id")
	}

	if request.ProjectId == 0 {
		return nil, bizerr.New(bizerr.ErrCodeParams, "missing project id")
	}

	if request.Region == globalcommonv1.Region_NAME_UNSPECIFIED {
		return nil, bizerr.New(bizerr.ErrCodeParams, "missing region")
	}

	if request.Provider == globalcommonv1.CloudProvider_NAME_UNSPECIFIED {
		return nil, bizerr.New(bizerr.ErrCodeParams, "unsupported cloud provider: %s", request.Provider)
	}

	if componentType == globaldedicatedv1.ComponentType_COMPONENT_TYPE_UNSPECIFIED {
		return nil, bizerr.New(bizerr.ErrCodeParams, "unsupported component: %s", componentType)
	}

	if request.SpuName != nil {
		return s.selectSkus(ctx, componentType, request)
	}
	if request.TiKVComponent != nil {
		return s.selectPDSku(ctx, request)
	}

	return nil, bizerr.New(bizerr.ErrCodeUnknown, "unsupported request: %+v", request)
}

func (s *skuServiceImpl) SelectSkusWithPrefetchedData(
	ctx context.Context,
	container *PrefetchedData,
	request *SelectSkusRequest,
) ([]*entity.Sku, bizerr.BizError) {
	// Filter skus by tenant, volume type and region. Return the grouped-by-region mappings.
	candidates := make([]*entity.Sku, 0)
	for _, mapping := range container.SpuMappings {
		theSku := container.SkuMap[mapping.SkuKey]

		// check if the sku is open to the tenant.
		// if the tenant id is 0, it means the tenant is not specified, so the sku is open to all tenants.
		if !mapping.TenantIDs[0] && !mapping.TenantIDs[request.TenantId] {
			continue
		}

		// check if the sku is in the provider
		if theSku.GetCloudProvider() != request.Provider {
			continue
		}

		// check if the sku is in the region
		if !mapping.Regions[globalcommonv1.Region_NAME_UNSPECIFIED] && !mapping.Regions[request.Region] {
			continue
		}

		if request.CpuArch != nil && theSku.GetCpuArch() != *request.CpuArch {
			continue
		}

		if request.CpuGeneration != nil && theSku.GetCpuGeneration() != *request.CpuGeneration {
			continue
		}

		// check if the sku has required volume type
		// FIXME(@yifan): temporary workaround for the hyperdisk-balanced volume type.
		//   Storage type Standard, Basic and Performance cannot be mapped to Hyperdisk Balanced now.
		//   1. For data storage volume type filtering:
		//     - If a volume type is specified in the request and doesn't match the SKU's data storage volume type:
		//       - If the requested volume type is PD_SSD but the SKU uses HYPERDISK_BALANCED, consider it a match (don't continue)
		//       - Otherwise, exclude this SKU (continue to next iteration)
		//   2. For raft log storage volume type filtering:
		//     - Similarly checks if the requested raft volume type is specified and doesn't match
		//     - Same special case: If PD_SSD is requested, also accept HYPERDISK_BALANCED
		if request.VolumeType != nil && theSku.GetDataStorage().GetVolumeType() != *request.VolumeType {
			if *request.VolumeType != globaldedicatedv1.StorageMu_PD_SSD ||
				theSku.GetDataStorage().GetVolumeType() != globaldedicatedv1.StorageMu_HYPERDISK_BALANCED {
				continue
			}
		}
		if request.RaftVolumeType != nil && theSku.GetRaftlogStorage().GetVolumeType() != *request.RaftVolumeType {
			if *request.RaftVolumeType != globaldedicatedv1.StorageMu_PD_SSD ||
				theSku.GetRaftlogStorage().GetVolumeType() != globaldedicatedv1.StorageMu_HYPERDISK_BALANCED {
				continue
			}
		}

		candidates = append(candidates, theSku)
	}

	// sort skus by priority, ascending.
	// Usually there is only one sku in the candidates, no need to sort.
	slices.SortFunc(candidates, func(a, b *entity.Sku) int {
		return int(a.GetPriority() - b.GetPriority())
	})

	return candidates, nil
}

func (s *skuServiceImpl) GetSkuByKey(ctx context.Context, key string) (*entity.Sku, bizerr.BizError) {
	ctx, span := tracer.Start(ctx, "Domain.GetSku")
	defer span.End()

	sku, err := s.skuRepo.LoadByKeys(ctx, []string{key})
	if err != nil {
		return nil, err
	}

	return sku[key], nil
}

func (s *skuServiceImpl) selectSkus(
	ctx context.Context,
	componentType globaldedicatedv1.ComponentType,
	request *SelectSkusRequest,
) ([]*entity.Sku, bizerr.BizError) {
	ctx, span := tracer.Start(ctx, "Domain.SelectTierUnitsExcludePD2")
	defer span.End()

	if componentType == globaldedicatedv1.ComponentType_COMPONENT_TYPE_UNSPECIFIED {
		return nil, bizerr.New(bizerr.ErrCodeParams, "unsupported component: %s", componentType)
	}

	spus, err := s.spuRepo.Query(ctx, &spu.QuerySpuFilter{
		ComponentTypes: []globaldedicatedv1.ComponentType{componentType},
		Name:           *request.SpuName,
	})
	if err != nil {
		return nil, err
	}
	if len(spus) == 0 {
		return nil, bizerr.New(
			bizerr.ErrCodeInvalidNodeSize,
			"node_size(%s) of the node_type(%s) is not available in specified region and cloud provider",
			*request.SpuName,
			componentType.String(),
		)
	}

	out := make([]*entity.Sku, 0)

	for i := range spus {
		spuId := spus[i].Id
		spuMappings, err := s.skuRepo.QuerySpuMappings(ctx, []uint64{spuId})
		if err != nil {
			return nil, err
		}

		skus, err := s.skuRepo.Query(ctx, &spu.SkuFilter{
			Component: componentType,
			SpuId:     lo.ToPtr(spuId),
		})
		if err != nil {
			return nil, err
		}

		container := &PrefetchedData{
			Spu:         spus[i],
			SpuMappings: spuMappings[spuId],
			SkuMap: lo.SliceToMap(skus, func(item *entity.Sku) (string, *entity.Sku) {
				return item.Key, item
			}),
		}

		sku, err := s.SelectSkusWithPrefetchedData(ctx, container, request)
		if err != nil {
			return nil, err
		}
		if sku != nil {
			out = append(out, sku...)
		}
	}

	if len(out) == 0 {
		r, _ := json.Marshal(request)
		return nil, bizerr.New(bizerr.ErrCodeParams, "unable to select sku. component=%s, req=%s", componentType, r)
	}

	return out, nil
}

func (s *skuServiceImpl) selectPDSku(ctx context.Context, request *SelectSkusRequest) ([]*entity.Sku, bizerr.BizError) {
	ctx, span := tracer.Start(ctx, "Domain.SelectPDTierUnit")
	defer span.End()

	if request.TiKVComponent.GetComponentType() != globaldedicatedv1.ComponentType_TIKV {
		return nil, bizerr.New(bizerr.ErrCodeParams, "pd depend on tikv component")
	}

	label, err := buildLabelBuilder(
		request.TenantId,
		request.Provider,
		request.Region,
		globaldedicatedv1.ComponentType_PD,
		request.CpuArch,
	).WithTiKVSku(request.TiKVComponent.Sku.Key).
		Build()
	if err != nil {
		return nil, err
	}

	out, err := s.matchedSkus(ctx, label, request.Provider, globaldedicatedv1.ComponentType_PD)
	if err != nil {
		return nil, err
	}

	slices.SortFunc(out, comparePrioritySku)

	if len(out) == 0 {
		r, _ := json.Marshal(request)
		return nil, bizerr.New(bizerr.ErrCodeParams, "unable to select sku for pd. req=%s", r)
	}

	return out, nil
}

// prioritySku is a slice type ([]*globaldedicatedv1.Sku) for SKU candidates.
//
// Sorting logic (comparePrioritySku):
//
// 1. If two SKUs have the same CPU arch, the one with higher Priority value comes first (descending order).
// 2. If CPU arch differs and one is ARM64, ARM64 is always ordered before the other arch.
// 3. Otherwise (non-ARM64 comparison), higher Priority still comes first.
//
// Effect:
//
// - After sorting, the best candidate is at index 0.
// - ARM64 SKUs are preferred when available.
// - Within the same arch group, larger Priority ranks higher.
type prioritySku []*globaldedicatedv1.Sku

func comparePrioritySku(a, b *globaldedicatedv1.Sku) int {
	if a.CpuArch == b.CpuArch {
		return cmp.Compare(b.Priority, a.Priority)
	}

	if a.CpuArch == globaldedicatedv1.InstanceMu_ARM64 {
		return -1
	}

	if b.CpuArch == globaldedicatedv1.InstanceMu_ARM64 {
		return 1
	}

	return cmp.Compare(b.Priority, a.Priority)
}

func (s *skuServiceImpl) matchedSkus(
	ctx context.Context,
	componentLabels Labels,
	provider globalcommonv1.CloudProvider_Name,
	componentType globaldedicatedv1.ComponentType,
) (prioritySku, bizerr.BizError) {
	filter := &spu.SkuFilter{
		Component: componentType,
		Provider:  provider,
	}

	skus, err := s.skuRepo.Query(ctx, filter)
	if err != nil {
		return nil, err
	}

	out := prioritySku{}
	for _, sku := range skus {
		ss, err := newSkuSelector(sku)
		if err != nil {
			tilog.Warn(ctx, "build tierunit selector failed", zap.Uint64("tierunit", sku.Id), zap.Error(err))
			continue
		}

		if ss.matches(componentLabels) {
			out = append(out, sku)
		}
	}

	slices.SortFunc(out, comparePrioritySku)
	return out, nil
}

func buildLabelBuilder(
	tenantId uint64,
	provider globalcommonv1.CloudProvider_Name,
	region globalcommonv1.Region_Name,
	componentType globaldedicatedv1.ComponentType,
	cpuArch *globaldedicatedv1.InstanceMu_CpuArch,
) *componentLabelsBuilder {
	builder := newLabelBuilder().
		WithTenant(tenantId).
		WithComponent(componentType).
		WithRegion(region).
		WithProvider(provider)
	if cpuArch != nil {
		builder.WithCpuArch(*cpuArch)
	}

	return builder
}