package spu

import (
	"context"

	"github.com/samber/lo"
	"github.com/tidbcloud/pkgv2/db"

	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/cluster"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/repository/region"
	"github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/repository/spu"
	svccommon "github.com/tidbcloud/tidb-management-service/dedicated-global-service/domain/service/common"
	"github.com/tidbcloud/tidb-management-service/utilities/bizerr"
	featgatesvc "github.com/tidbcloud/tidb-management-service/utilities/featuregate"
)

type ServiceSpu interface {
	List(ctx context.Context, listReq *globaldedicatedv1.ListSpusRequest) ([]*entity.Spu, bizerr.BizError)
	Get(ctx context.Context, getReq *globaldedicatedv1.GetSpuRequest) (*entity.Spu, bizerr.BizError)
	Delete(ctx context.Context, id uint64) bizerr.BizError
	Create(ctx context.Context, createReq *globaldedicatedv1.Spu) (*entity.Spu, bizerr.BizError)
	Update(
		ctx context.Context,
		updateReq *globaldedicatedv1.BatchUpsertSpusRequest_Update,
	) (*entity.Spu, bizerr.BizError)
}

type spuServiceImpl struct {
	regionRepo  region.Repo
	spuRepo     spu.RepoSpu
	skuRepo     spu.SkuRepo
	clusterRepo cluster.Repo

	skuSvc SkuService

	featureGate featgatesvc.Service
	transaction db.Transactioner
}

func NewSpuService(
	spuRepo spu.RepoSpu,
	skuRepo spu.SkuRepo,
	regionRepo region.Repo,
	clusterRepo cluster.Repo,
	skuSvc SkuService,
	featureGate featgatesvc.Service,
	transaction db.Transactioner,
) ServiceSpu {
	return &spuServiceImpl{
		regionRepo:  regionRepo,
		spuRepo:     spuRepo,
		skuRepo:     skuRepo,
		clusterRepo: clusterRepo,
		skuSvc:      skuSvc,
		featureGate: featureGate,
		transaction: transaction,
	}
}

func (s *spuServiceImpl) List(
	ctx context.Context,
	request *globaldedicatedv1.ListSpusRequest,
) ([]*entity.Spu, bizerr.BizError) {
	ctx, span := tracer.Start(ctx, "SpuDomainService.List")
	defer span.End()

	req, ok := request.Request.(*globaldedicatedv1.ListSpusRequest_WithSku)
	if ok {
		return s.listWithSku(ctx, req.WithSku)
	}

	filter := request.Request.(*globaldedicatedv1.ListSpusRequest_WithoutSku).WithoutSku

	out, err := s.spuRepo.Query(ctx, &spu.QuerySpuFilter{
		ComponentTypes: filter.GetComponentTypes(),
		Name:           filter.GetName(),
		Versions:       filter.GetVersions(),
	})
	if err != nil {
		return nil, err
	}

	// For each spu, load the mappings.
	for _, it := range out {
		resp, err := s.skuRepo.QuerySpuMappings(ctx, []uint64{it.Id})
		if err != nil {
			return nil, err
		}
		mappings, exist := resp[it.Id]
		if exist {
			for _, mapping := range mappings {
				it.Mappings = append(it.Mappings, covertMappingToProto(mapping))
			}
		}
	}

	return out, nil
}

func (s *spuServiceImpl) listWithSku(
	ctx context.Context,
	filter *globaldedicatedv1.ListSpusRequest_FilterWithSku,
) ([]*entity.Spu, bizerr.BizError) {
	// Check if the request is valid.
	if filter.GetCloudProvider() == globalcommonv1.CloudProvider_NAME_UNSPECIFIED {
		return nil, bizerr.New(bizerr.ErrCodeParams, "cloud provider must be specified")
	}
	if filter.GetRegion() == globalcommonv1.Region_NAME_UNSPECIFIED {
		return nil, bizerr.New(bizerr.ErrCodeParams, "region must be specified")
	}
	if filter.GetTenantId() == 0 {
		return nil, bizerr.New(bizerr.ErrCodeParams, "tenant id must be specified")
	}
	if filter.GetProjectId() == 0 { // project is used by check prefer arm feature
		return nil, bizerr.New(bizerr.ErrCodeParams, "project id must be specified")
	}

	// Get cloud provider.
	cloudProviderEn, err := s.regionRepo.GetCloudProviderById(ctx, filter.GetCloudProvider(), filter.GetTenantId())
	if err != nil {
		return nil, err
	}
	if cloudProviderEn.GetState() != globalcommonv1.CloudProvider_ACTIVE {
		return nil, bizerr.New(bizerr.ErrCodeParams, "inactive cloud provider %s", filter.GetCloudProvider())
	}

	// Load regions.
	regions, err := s.regionRepo.ListRegions(ctx, filter.GetTenantId(), &region.Filter{
		CloudProviderNames: []globalcommonv1.CloudProvider_Name{filter.GetCloudProvider()},
		RegionNames:        []entity.RegionName{filter.GetRegion()},
		OnlyActive:         true,
	})
	if err != nil {
		return nil, err
	}

	// Load existing clusterEn if provided.
	newCluster := filter.GetExistingClusterId() == 0
	clusterEn := &entity.Cluster{Cluster: &globaldedicatedv1.Cluster{
		CloudProviderId: filter.GetCloudProvider(),
		RegionId:        filter.GetRegion(),
		OrganizationId:  filter.GetTenantId(),
		ProjectId:       filter.GetProjectId(),
	}}
	if !newCluster {
		clusterEn, err = s.clusterRepo.GetAndLoadComponents(ctx, filter.GetExistingClusterId())
		if err != nil {
			return nil, err
		}
	}

	// Load spus.
	spus, err := s.spuRepo.Query(ctx, &spu.QuerySpuFilter{ComponentTypes: lo.Keys(entity.ComponentsUserEditable)})
	if err != nil {
		return nil, err
	}

	spuIds := lo.Map(spus, func(item *entity.Spu, index int) uint64 {
		return item.Id
	})

	// Load spu mappings.
	spuMappings, err := s.skuRepo.QuerySpuMappings(ctx, spuIds)
	if err != nil {
		return nil, err
	}

	// Load skus.
	skus, err := s.skuRepo.Query(ctx, &spu.SkuFilter{})
	if err != nil {
		return nil, err
	}

	keyToSku := lo.SliceToMap(skus, func(item *entity.Sku) (string, *entity.Sku) {
		return item.Key, item
	})

	res := make(map[string]*globaldedicatedv1.Spu)

	// For each spu, pick the skus that match the filter.
	for _, p := range spus {
		// For each logical regions.
		for _, r := range regions {
			// Select skus without volume type filter.
			dummyComp := &globaldedicatedv1.Cluster_ClusterComponent{ComponentType: p.GetComponentType()}
			selectSkusRequest := &SelectSkusRequest{
				TenantId:      filter.GetTenantId(),
				ProjectId:     filter.GetProjectId(),
				Provider:      filter.GetCloudProvider(),
				Region:        r.GetId(),
				CpuArch:       svccommon.GetCpuArch(s.featureGate, clusterEn, dummyComp, newCluster),
				CpuGeneration: svccommon.GetCpuGenerationFilter(s.featureGate, clusterEn, dummyComp, newCluster),
			}
			prefetched := &PrefetchedData{
				Spu:         p,
				SpuMappings: spuMappings[p.Id],
				SkuMap:      keyToSku,
			}
			selectedSkus, err := s.skuSvc.SelectSkusWithPrefetchedData(ctx, prefetched, selectSkusRequest)
			if err != nil {
				return nil, err
			}
			if len(selectedSkus) == 0 {
				continue
			}

			// build the global spu.
			if _, ok := res[p.GetKey()]; !ok {
				res[p.GetKey()] = &globaldedicatedv1.Spu{
					Id:            p.Id,
					Key:           p.Key,
					Name:          p.Name,
					DisplayName:   p.DisplayName,
					Version:       p.Version,
					ComponentType: p.ComponentType,
					CpuCore:       p.CpuCore,
					MemorySizeGi:  p.MemorySizeGi,
					DefaultCount:  selectedSkus[0].DefaultNodeCount,
					CloudProvider: filter.GetCloudProvider(),
					Region:        filter.GetRegion(),
				}
			}

			if _, ok := entity.ComponentsHaveStorage[p.GetComponentType()]; !ok {
				continue
			}

			// build storages
			resSpu := res[p.GetKey()]
			for _, sku := range selectedSkus {
				if sku.GetDataStorage().GetVolumeType() != globaldedicatedv1.StorageMu_VOLUME_TYPE_UNSPECIFIED {
					if resSpu.DataStorage == nil {
						resSpu.DataStorage = &globaldedicatedv1.Spu_Storage{}
					}
					fillSpuStorage(resSpu.DataStorage, sku.GetDataStorage(), true)
				}
				if sku.GetRaftlogStorage().GetVolumeType() != globaldedicatedv1.StorageMu_VOLUME_TYPE_UNSPECIFIED {
					if resSpu.RaftlogStorage == nil {
						resSpu.RaftlogStorage = &globaldedicatedv1.Spu_Storage{}
					}
					fillSpuStorage(resSpu.RaftlogStorage, sku.GetRaftlogStorage(), true)
				}

				storageType := resolveSpuStorageType(
					sku.GetDataStorage().GetVolumeType(),
					sku.GetRaftlogStorage().GetVolumeType(),
				)
				// Only AWS TiKV Plus data storage exposes per-storage-type IOPS overrides in storage_type_details.
				if shouldBuildSpuStorageTypeDetails(
					p.GetComponentType(),
					sku.GetCloudProvider(),
					storageType,
				) {
					if resSpu.StorageTypeDetails == nil {
						resSpu.StorageTypeDetails = make(map[string]*globaldedicatedv1.Spu_StorageTypeDetail)
					}
					storageTypeKey := storageType.String()
					if _, ok := resSpu.StorageTypeDetails[storageTypeKey]; !ok {
						resSpu.StorageTypeDetails[storageTypeKey] = &globaldedicatedv1.Spu_StorageTypeDetail{}
					}
					detail := resSpu.StorageTypeDetails[storageTypeKey]
					if sku.GetDataStorage().GetVolumeType() != globaldedicatedv1.StorageMu_VOLUME_TYPE_UNSPECIFIED {
						if detail.DataStorage == nil {
							detail.DataStorage = &globaldedicatedv1.Spu_Storage{}
						}
						fillSpuStorage(detail.DataStorage, sku.GetDataStorage(), false)
					}
				}
				resSpu.SupportedDiskPairs = append(resSpu.SupportedDiskPairs, &globaldedicatedv1.Spu_DiskPair{
					DataVolumeType:    sku.GetDataStorage().GetVolumeType(),
					RaftlogVolumeType: sku.GetRaftlogStorage().GetVolumeType(),
				})
			}
			resSpu.SupportedDiskPairs = lo.UniqBy(
				resSpu.GetSupportedDiskPairs(),
				func(i *globaldedicatedv1.Spu_DiskPair) string {
					return i.GetDataVolumeType().String() + i.GetRaftlogVolumeType().String()
				},
			)
		}
	}

	return lo.Values(res), nil
}

func fillSpuStorage(dst *globaldedicatedv1.Spu_Storage, src *globaldedicatedv1.Sku_Storage, includeSize bool) {
	if includeSize {
		dst.DefaultSizeGi = lo.Ternary(dst.DefaultSizeGi == 0, src.GetSizeGiDefault(), dst.DefaultSizeGi)
		dst.MaxSizeGi = lo.Ternary(dst.MaxSizeGi == 0, src.GetSizeGiMax(), dst.MaxSizeGi)
		dst.MinSizeGi = lo.Ternary(dst.MinSizeGi == 0, src.GetSizeGiMin(), dst.MinSizeGi)
	}
	dst.DefaultIops = lo.Ternary(dst.DefaultIops == 0, defaultSkuStorageIops(src), dst.DefaultIops)
	dst.MaxIops = lo.Ternary(dst.MaxIops == 0, src.GetIopsMax(), dst.MaxIops)
	dst.MinIops = lo.Ternary(dst.MinIops == 0, src.GetIopsMin(), dst.MinIops)
}

func defaultSkuStorageIops(storage *globaldedicatedv1.Sku_Storage) int32 {
	switch {
	case storage.GetIopsFixed() != nil:
		return storage.GetIopsFixed().GetIops()
	case storage.GetIopsDynamic() != nil:
		// Dynamic IOPS scale with storage size, so there is no single default IOPS value to expose.
		return 0
	default:
		return 0
	}
}

func shouldBuildSpuStorageTypeDetails(
	componentType globaldedicatedv1.ComponentType,
	cloudProvider globalcommonv1.CloudProvider_Name,
	storageType globaldedicatedv1.Spu_StorageType,
) bool {
	return componentType == globaldedicatedv1.ComponentType_TIKV &&
		cloudProvider == globalcommonv1.CloudProvider_AWS &&
		storageType == globaldedicatedv1.Spu_Plus
}

// storageTypeVolumePair identifies one data/raftlog volume combination for SPU storage classification.
type storageTypeVolumePair struct {
	dataVt    globaldedicatedv1.StorageMu_VolumeType
	raftlogVt globaldedicatedv1.StorageMu_VolumeType
}

// storageTypeVolumePairs classifies data/raftlog volume pairs into the public SPU storage types.
var storageTypeVolumePairs = map[storageTypeVolumePair]globaldedicatedv1.Spu_StorageType{
	{dataVt: globaldedicatedv1.StorageMu_GP3, raftlogVt: globaldedicatedv1.StorageMu_VOLUME_TYPE_UNSPECIFIED}:                globaldedicatedv1.Spu_Basic,
	{dataVt: globaldedicatedv1.StorageMu_GP3, raftlogVt: globaldedicatedv1.StorageMu_GP3}:                                    globaldedicatedv1.Spu_Standard,
	{dataVt: globaldedicatedv1.StorageMu_GP3, raftlogVt: globaldedicatedv1.StorageMu_IO2}:                                    globaldedicatedv1.Spu_Performance,
	{dataVt: globaldedicatedv1.StorageMu_IO2, raftlogVt: globaldedicatedv1.StorageMu_VOLUME_TYPE_UNSPECIFIED}:                globaldedicatedv1.Spu_Plus,
	{dataVt: globaldedicatedv1.StorageMu_PD_SSD, raftlogVt: globaldedicatedv1.StorageMu_VOLUME_TYPE_UNSPECIFIED}:             globaldedicatedv1.Spu_Basic,
	{dataVt: globaldedicatedv1.StorageMu_PD_SSD, raftlogVt: globaldedicatedv1.StorageMu_PD_SSD}:                              globaldedicatedv1.Spu_Standard,
	{dataVt: globaldedicatedv1.StorageMu_HYPERDISK_BALANCED, raftlogVt: globaldedicatedv1.StorageMu_VOLUME_TYPE_UNSPECIFIED}: globaldedicatedv1.Spu_Basic,
	{dataVt: globaldedicatedv1.StorageMu_HYPERDISK_BALANCED, raftlogVt: globaldedicatedv1.StorageMu_HYPERDISK_BALANCED}:      globaldedicatedv1.Spu_Standard,
	{dataVt: globaldedicatedv1.StorageMu_PREMIUMV2, raftlogVt: globaldedicatedv1.StorageMu_VOLUME_TYPE_UNSPECIFIED}:          globaldedicatedv1.Spu_Basic,
	{dataVt: globaldedicatedv1.StorageMu_PREMIUMV2, raftlogVt: globaldedicatedv1.StorageMu_PREMIUMV2}:                        globaldedicatedv1.Spu_Standard,
	{dataVt: globaldedicatedv1.StorageMu_PREMIUMV2, raftlogVt: globaldedicatedv1.StorageMu_ULTRASSD}:                         globaldedicatedv1.Spu_Performance,
	{dataVt: globaldedicatedv1.StorageMu_ULTRASSD, raftlogVt: globaldedicatedv1.StorageMu_VOLUME_TYPE_UNSPECIFIED}:           globaldedicatedv1.Spu_Plus,
}

func resolveSpuStorageType(
	dataVt, raftlogVt globaldedicatedv1.StorageMu_VolumeType,
) globaldedicatedv1.Spu_StorageType {
	if storageType, ok := storageTypeVolumePairs[storageTypeVolumePair{
		dataVt:    dataVt,
		raftlogVt: raftlogVt,
	}]; ok {
		return storageType
	}
	return globaldedicatedv1.Spu_STORAGE_TYPE_UNSPECIFIED
}

func (s *spuServiceImpl) Get(ctx context.Context, req *globaldedicatedv1.GetSpuRequest) (*entity.Spu, bizerr.BizError) {
	ctx, span := tracer.Start(ctx, "SpuDomainService.Get")
	defer span.End()

	getById := func() (*entity.Spu, bizerr.BizError) {
		it, err := s.spuRepo.Get(ctx, req.GetId())
		if err != nil {
			return nil, err
		}
		return it, nil
	}

	getByAttr := func() (*entity.Spu, bizerr.BizError) {
		spus, err := s.listWithSku(ctx, &globaldedicatedv1.ListSpusRequest_FilterWithSku{
			CloudProvider:     req.GetByAttributes().GetCloudProvider(),
			Region:            req.GetByAttributes().GetRegion(),
			TenantId:          req.GetByAttributes().GetTenantId(),
			ProjectId:         req.GetByAttributes().GetProjectId(),
			ExistingClusterId: req.GetByAttributes().GetExistingClusterId(),
		})
		if err != nil {
			return nil, err
		}
		filtered := lo.Filter(spus, func(item *entity.Spu, _ int) bool {
			return item.GetComponentType() == req.GetByAttributes().GetComponentType() &&
				item.GetName() == req.GetByAttributes().GetSpuName()
		})
		if len(filtered) == 0 {
			return nil, bizerr.New(bizerr.ErrCodeRecordNotFound, "spu not found with specified attributes")
		}
		if len(filtered) > 1 {
			return nil, bizerr.New(bizerr.ErrCodeUnknown, "multiple spus found with specified attributes")
		}
		return filtered[0], nil
	}

	var (
		spuRes *entity.Spu
		err    bizerr.BizError
	)
	if req.GetId() != 0 {
		spuRes, err = getById()
	} else {
		spuRes, err = getByAttr()
	}
	if err != nil {
		return nil, err
	}

	resp, err := s.skuRepo.QuerySpuMappings(ctx, []uint64{spuRes.GetId()})
	if err != nil {
		return nil, err
	}
	mappings, exist := resp[req.GetId()]
	if exist {
		for _, mapping := range mappings {
			spuRes.Mappings = append(spuRes.Mappings, covertMappingToProto(mapping))
		}
	}

	return spuRes, nil
}

func (s *spuServiceImpl) Delete(ctx context.Context, id uint64) bizerr.BizError {
	ctx, span := tracer.Start(ctx, "Domain.DeleteTier")
	defer span.End()

	return s.spuRepo.Delete(ctx, id)
}

func (s *spuServiceImpl) Create(ctx context.Context, req *globaldedicatedv1.Spu) (*entity.Spu, bizerr.BizError) {
	ctx, span := tracer.Start(ctx, "SpuDomainService.Create")
	defer span.End()

	if req.GetDisplayName() == "" {
		return nil, bizerr.New(bizerr.ErrCodeParams, "display name is required")
	}
	if req.GetName() == "" {
		return nil, bizerr.New(bizerr.ErrCodeParams, "name is required")
	}
	if req.GetComponentType() == globaldedicatedv1.ComponentType_COMPONENT_TYPE_UNSPECIFIED {
		return nil, bizerr.New(bizerr.ErrCodeParams, "component type is required")
	}
	if req.GetCpuCore() == 0 {
		return nil, bizerr.New(bizerr.ErrCodeParams, "cpu core is required")
	}
	if req.GetMemorySizeGi() == 0 {
		return nil, bizerr.New(bizerr.ErrCodeParams, "memory size is required")
	}
	if err := validSpuMappings(req.GetMappings()); err != nil {
		return nil, err
	}

	// generate key
	req.Key = entity.BuildSpuKey(req.ComponentType, req.Name, req.Version)

	err := s.transaction.Transaction(ctx, func(ctx context.Context) error {
		id, err := s.spuRepo.Save(ctx, req)
		if err != nil {
			return err
		}
		mappings := make([]*entity.SpuMapping, 0, len(req.Mappings))
		for _, mapping := range req.GetMappings() {
			mappings = append(mappings, covertMappingToEntity(mapping, uint64(id)))
		}
		if _, err := s.spuRepo.SaveSpuMappings(ctx, mappings); err != nil {
			return err
		}
		req.Id = uint64(id)
		return nil
	})
	if err != nil {
		return nil, bizerr.WrapError(err, "create spu failed")
	}

	return req, nil
}

func (s *spuServiceImpl) Update(
	ctx context.Context,
	req *globaldedicatedv1.BatchUpsertSpusRequest_Update,
) (*entity.Spu, bizerr.BizError) {
	ctx, span := tracer.Start(ctx, "SpuDomainService.Update")
	defer span.End()

	var existing *entity.Spu
	err := s.transaction.Transaction(ctx, func(ctx context.Context) error {
		var err bizerr.BizError
		existing, err = s.spuRepo.Get(ctx, req.GetId())
		if err != nil {
			return err
		}

		if req.DisplayName != "" {
			existing.DisplayName = req.DisplayName
		}
		if req.GetDisplayName() != "" {
			existing.DisplayName = req.GetDisplayName()
		}
		if len(req.GetMappings()) > 0 {
			if err := validSpuMappings(req.GetMappings()); err != nil {
				return err
			}
			existing.Mappings = req.GetMappings()
		}

		// Save spu.
		if _, err := s.spuRepo.Save(ctx, existing); err != nil {
			return err
		}

		// Save mappings.
		if _, err := s.spuRepo.SaveSpuMappings(
			ctx,
			lo.Map(existing.GetMappings(), func(i *globaldedicatedv1.Spu_Mapping, _ int) *entity.SpuMapping {
				return covertMappingToEntity(i, req.GetId())
			}),
		); err != nil {
			return err
		}

		return nil
	})
	if err != nil {
		return nil, bizerr.WrapError(err, "update spu failed")
	}
	existing.Mappings = req.Mappings

	return existing, nil
}

func validSpuMappings(mappings []*globaldedicatedv1.Spu_Mapping) bizerr.BizError {
	for _, mapping := range mappings {
		if mapping.GetCloudProvider() == globalcommonv1.CloudProvider_NAME_UNSPECIFIED {
			return bizerr.New(bizerr.ErrCodeParams, "cloud provider is required")
		}
		if mapping.GetSkuKey() == "" {
			return bizerr.New(bizerr.ErrCodeParams, "sku key is required")
		}
	}
	return nil
}

func covertMappingToEntity(mapping *globaldedicatedv1.Spu_Mapping, spuId uint64) *entity.SpuMapping {
	regions := make(map[globalcommonv1.Region_Name]bool, len(mapping.PhysicalRegions))
	for _, r := range mapping.PhysicalRegions {
		regions[r] = true
	}
	tenants := make(map[uint64]bool, len(mapping.OrganizationIds))
	for _, tenant := range mapping.OrganizationIds {
		tenants[tenant] = true
	}

	return &entity.SpuMapping{
		Provider:  mapping.CloudProvider,
		SpuId:     spuId,
		Regions:   regions,
		TenantIDs: tenants,
		SkuKey:    mapping.GetSkuKey(),
	}
}

func covertMappingToProto(mapping *entity.SpuMapping) *globaldedicatedv1.Spu_Mapping {
	var regions []globalcommonv1.Region_Name
	for r, ok := range mapping.Regions {
		if ok {
			regions = append(regions, r)
		}
	}
	var orgs []uint64
	for org, ok := range mapping.TenantIDs {
		if ok {
			orgs = append(orgs, org)
		}
	}
	return &globaldedicatedv1.Spu_Mapping{
		CloudProvider:   mapping.Provider,
		PhysicalRegions: regions,
		OrganizationIds: orgs,
		SkuKey:          mapping.SkuKey,
	}
}