package server

import (
	"context"
	"fmt"
	"sort"
	"strings"

	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	dplanemgrv1 "github.com/tidbcloud/cluster-service-ng/api/dataplane-manager/v1"
	"github.com/tidbcloud/cluster-service-ng/config"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	"github.com/tidbcloud/cluster-service-ng/pkg/utils"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"google.golang.org/grpc/codes"
	"google.golang.org/protobuf/proto"
)

const (
	nodeSpecDropUnsupportedNodeType = "unsupported_node_type"
	nodeSpecDropInvalidCapacity     = "invalid_capacity"
	nodeSpecDropInvalidAvailable    = "invalid_available_capacity"
	nodeSpecDropInvalidLocalDisk    = "invalid_local_disk"
	nodeSpecDropDuplicate           = "duplicate"
)

type normalizedNodeSpecKey struct {
	cpu               uint32
	memoryGi          uint32
	availableCPU      uint32
	availableMemoryMi uint32
	withLocalDisk     bool
	localDiskSizeGi   uint32
}

func (s *serverImpl) ListNodeSpecs(ctx context.Context, tenantID string, componentType entity.ComponentType) (*entity.ListNodeSpecsResponse, error) {
	component := componentType.String()
	if tenantID == "" {
		return nil, resourcePoolStatusError(codes.InvalidArgument, reasonNodeSpecsUnavailable, "tenant_id is required")
	}
	if componentType != apiv1.ComponentType_TIKV && componentType != apiv1.ComponentType_TIFLASH_COMPUTE {
		return nil, resourcePoolStatusError(codes.InvalidArgument, reasonInvalidBootstrapCapacity,
			fmt.Sprintf("component_type %s is not supported", component))
	}

	ruConfig, err := s.resolveSharedPoolRuConfig(ctx, tenantID)
	if err != nil {
		return nil, err
	}
	options, err := buildCreationNodeCountOptions(componentType, ruConfig)
	if err != nil {
		return nil, err
	}

	nodeSpecs, dropped, err := s.listDefaultNodeSpecs(ctx)
	if err != nil {
		tilog.Error(ctx, "failed to list deployment node specs", zap.String("component", component), zap.Error(err))
		return nil, resourcePoolStatusError(codes.Unavailable, reasonNodeSpecsUnavailable, "node specs are unavailable")
	}
	logNodeSpecDrops(ctx, dropped)
	nodeSpecs = filterNodeSpecsForComponent(nodeSpecs, componentType, ruConfig)
	if len(nodeSpecs) == 0 {
		tilog.Warn(ctx, "no usable deployment node specs", zap.String("tenantId", tenantID), zap.String("component", component))
		return nil, resourcePoolStatusError(codes.FailedPrecondition, reasonNodeSpecsUnavailable, "no usable node specs are configured")
	}

	storageSizePolicy := s.storageSizePolicy
	if storageSizePolicy == (config.StorageSizePolicy{}) {
		storageSizePolicy = config.DefaultStorageSizePolicy()
	}
	nodeSpecOptions := buildNodeSpecOptions(nodeSpecs, componentType, storageSizePolicy)
	if len(nodeSpecOptions) == 0 {
		tilog.Warn(ctx, "no usable deployment node spec options", zap.String("tenantId", tenantID), zap.String("component", component))
		return nil, resourcePoolStatusError(codes.FailedPrecondition, reasonNodeSpecsUnavailable, "no usable node spec options are configured")
	}

	tilog.Info(ctx, "listed deployment node specs", zap.String("tenantId", tenantID), zap.String("component", component), zap.Int("resultCount", len(nodeSpecOptions)))
	return &entity.ListNodeSpecsResponse{
		NodeSpecOptions:  nodeSpecOptions,
		NodeCountOptions: options,
	}, nil
}

func buildCreationNodeCountOptions(componentType entity.ComponentType, ruConfig *entity.RuConfig) (*entity.NodeCountOptions, error) {
	var options *entity.NodeCountOptions
	maxNodeCount := ruConfig.GetNodeCountQuota(componentType)
	switch componentType {
	case apiv1.ComponentType_TIKV:
		if ruConfig == nil {
			return nil, resourcePoolStatusError(codes.FailedPrecondition, reasonNodeSpecsUnavailable, "RU configuration is unavailable")
		}
		maxNodeCount -= maxNodeCount % 3
		if maxNodeCount < 3 {
			return nil, resourcePoolStatusError(codes.FailedPrecondition, reasonNodeSpecsUnavailable,
				"TiKV node-count quota does not admit the minimum creation count")
		}
		options = &entity.NodeCountOptions{
			DefaultNodeCount: 3,
			MinNodeCount:     3,
			MaxNodeCount:     maxNodeCount,
			Rule: &entity.NodeCountRule{
				Type:  apiv1.NodeCountRuleType_NODE_COUNT_RULE_TYPE_MULTIPLE_OF,
				Value: 3,
			},
		}
	case apiv1.ComponentType_TIFLASH_COMPUTE:
		if maxNodeCount <= 0 {
			return nil, resourcePoolStatusError(codes.FailedPrecondition, reasonNodeSpecsUnavailable,
				"TiFlash node-count quota does not admit the minimum creation count")
		}
		options = &entity.NodeCountOptions{
			DefaultNodeCount: 1,
			MinNodeCount:     1,
			MaxNodeCount:     maxNodeCount,
			Rule: &entity.NodeCountRule{
				Type: apiv1.NodeCountRuleType_NODE_COUNT_RULE_TYPE_NONE,
			},
		}
	default:
		return nil, resourcePoolStatusError(codes.InvalidArgument, reasonInvalidBootstrapCapacity, fmt.Sprintf("component_type %s is not supported", componentType.String()))
	}

	defaultCount := options.GetDefaultNodeCount()
	if err := entity.ValidateDefaultNodeCount(componentType, &defaultCount, options.GetMaxNodeCount()); err != nil {
		return nil, fmt.Errorf("invalid creation node-count options for %s: %w", componentType.String(), err)
	}
	return options, nil
}

func validateCreationNodeCount(componentType entity.ComponentType, nodeCount uint32, ruConfig *entity.RuConfig) error {
	options, err := buildCreationNodeCountOptions(componentType, ruConfig)
	if err != nil {
		return err
	}
	if nodeCount < options.GetMinNodeCount() || nodeCount > options.GetMaxNodeCount() {
		return resourcePoolStatusError(codes.InvalidArgument, reasonInvalidBootstrapCapacity,
			fmt.Sprintf("node_count %d for %s must be between %d and %d", nodeCount, componentType.String(), options.GetMinNodeCount(), options.GetMaxNodeCount()))
	}
	if options.GetRule().GetType() == apiv1.NodeCountRuleType_NODE_COUNT_RULE_TYPE_MULTIPLE_OF &&
		nodeCount%options.GetRule().GetValue() != 0 {
		return resourcePoolStatusError(codes.InvalidArgument, reasonInvalidBootstrapCapacity,
			fmt.Sprintf("node_count %d for %s must be a multiple of %d", nodeCount, componentType.String(), options.GetRule().GetValue()))
	}
	return nil
}

type nodeSpecShape struct {
	cpu      uint32
	memoryGi uint32
}

type storageSizeOptionsKey struct {
	name          entity.StorageName
	medium        entity.StorageMedium
	defaultSizeGi uint32
	minSizeGi     uint32
	maxSizeGi     uint32
}

func buildNodeSpecOptions(
	nodeSpecs []*entity.NodeSpec,
	componentType entity.ComponentType,
	policy config.StorageSizePolicy,
) []*entity.NodeSpecOptions {
	grouped := make(map[nodeSpecShape]map[storageSizeOptionsKey]*entity.StorageSizeOptions)
	for _, nodeSpec := range nodeSpecs {
		if nodeSpec == nil {
			continue
		}
		shape := nodeSpecShape{cpu: nodeSpec.GetCpu(), memoryGi: nodeSpec.GetMemoryGi()}
		storageOptions, ok := grouped[shape]
		if !ok {
			storageOptions = make(map[storageSizeOptionsKey]*entity.StorageSizeOptions)
			grouped[shape] = storageOptions
		}

		switch componentType {
		case apiv1.ComponentType_TIKV:
			addConfiguredStorageSizeOption(storageOptions, apiv1.Component_Storage_RAFTLOG, apiv1.StorageMedium_STORAGE_MEDIUM_CLOUD, policy.TiKV.Raftlog)
			if nodeSpec.GetWithLocalDisk() {
				addLocalStorageSizeOption(storageOptions, apiv1.Component_Storage_DATA, nodeSpec.GetLocalDiskSizeGi())
			} else {
				addConfiguredStorageSizeOption(storageOptions, apiv1.Component_Storage_DATA, apiv1.StorageMedium_STORAGE_MEDIUM_CLOUD, policy.TiKV.Data)
			}
		case apiv1.ComponentType_TIFLASH_COMPUTE:
			if nodeSpec.GetWithLocalDisk() {
				addLocalStorageSizeOption(storageOptions, apiv1.Component_Storage_DATA, nodeSpec.GetLocalDiskSizeGi())
			} else {
				addConfiguredStorageSizeOption(storageOptions, apiv1.Component_Storage_DATA, apiv1.StorageMedium_STORAGE_MEDIUM_CLOUD, policy.TiFlashCompute.Data)
			}
		}
	}

	shapes := make([]nodeSpecShape, 0, len(grouped))
	for shape := range grouped {
		shapes = append(shapes, shape)
	}
	sort.Slice(shapes, func(i, j int) bool {
		if shapes[i].cpu != shapes[j].cpu {
			return shapes[i].cpu < shapes[j].cpu
		}
		return shapes[i].memoryGi < shapes[j].memoryGi
	})

	result := make([]*entity.NodeSpecOptions, 0, len(shapes))
	for _, shape := range shapes {
		storageOptions := make([]*entity.StorageSizeOptions, 0, len(grouped[shape]))
		for _, option := range grouped[shape] {
			storageOptions = append(storageOptions, option)
		}
		sort.Slice(storageOptions, func(i, j int) bool {
			left, right := storageOptions[i], storageOptions[j]
			if left.GetName() != right.GetName() {
				return left.GetName() < right.GetName()
			}
			if left.GetMedium() != right.GetMedium() {
				return left.GetMedium() < right.GetMedium()
			}
			if left.GetMinSizeGi() != right.GetMinSizeGi() {
				return left.GetMinSizeGi() < right.GetMinSizeGi()
			}
			if left.GetDefaultSizeGi() != right.GetDefaultSizeGi() {
				return left.GetDefaultSizeGi() < right.GetDefaultSizeGi()
			}
			return left.GetMaxSizeGi() < right.GetMaxSizeGi()
		})

		result = append(result, &entity.NodeSpecOptions{
			NodeSpec:       &entity.NodeSpec{Cpu: shape.cpu, MemoryGi: shape.memoryGi},
			StorageOptions: storageOptions,
		})
	}
	return result
}

func addConfiguredStorageSizeOption(
	options map[storageSizeOptionsKey]*entity.StorageSizeOptions,
	name entity.StorageName,
	medium entity.StorageMedium,
	sizeRange config.StorageSizeRange,
) {
	addStorageSizeOption(options, storageSizeOptionsKey{
		name:          name,
		medium:        medium,
		defaultSizeGi: sizeRange.DefaultSizeGi,
		minSizeGi:     sizeRange.MinSizeGi,
		maxSizeGi:     sizeRange.MaxSizeGi,
	})
}

func addLocalStorageSizeOption(
	options map[storageSizeOptionsKey]*entity.StorageSizeOptions,
	name entity.StorageName,
	sizeGi uint32,
) {
	addStorageSizeOption(options, storageSizeOptionsKey{
		name:          name,
		medium:        apiv1.StorageMedium_STORAGE_MEDIUM_LOCAL,
		defaultSizeGi: sizeGi,
		minSizeGi:     sizeGi,
		maxSizeGi:     sizeGi,
	})
}

func addStorageSizeOption(
	options map[storageSizeOptionsKey]*entity.StorageSizeOptions,
	key storageSizeOptionsKey,
) {
	if _, exists := options[key]; exists {
		return
	}
	options[key] = &entity.StorageSizeOptions{
		Name:          key.name,
		Medium:        key.medium,
		DefaultSizeGi: key.defaultSizeGi,
		MinSizeGi:     key.minSizeGi,
		MaxSizeGi:     key.maxSizeGi,
	}
}

func filterNodeSpecsForComponent(nodeSpecs []*entity.NodeSpec, componentType entity.ComponentType, ruConfig *entity.RuConfig) []*entity.NodeSpec {
	if ruConfig == nil {
		return nil
	}
	componentConfig := ruConfig.GetComponents()[strings.ToLower(componentType.String())]
	if componentConfig == nil {
		return nil
	}

	configuredShapes := make(map[nodeSpecShape]struct{}, len(componentConfig.GetNodeSpecRuRange()))
	for configuredSpec := range componentConfig.GetNodeSpecRuRange() {
		cpu, memoryGi, err := utils.ParseNodeSpec(configuredSpec)
		if err != nil || cpu <= 0 || memoryGi <= 0 {
			continue
		}
		configuredShapes[nodeSpecShape{cpu: uint32(cpu), memoryGi: uint32(memoryGi)}] = struct{}{}
	}

	filtered := make([]*entity.NodeSpec, 0, len(nodeSpecs))
	for _, nodeSpec := range nodeSpecs {
		if _, ok := configuredShapes[nodeSpecShape{cpu: nodeSpec.GetCpu(), memoryGi: nodeSpec.GetMemoryGi()}]; ok {
			filtered = append(filtered, nodeSpec)
		}
	}
	return filtered
}

func (s *serverImpl) listDefaultNodeSpecs(ctx context.Context) ([]*entity.NodeSpec, map[string]int, error) {
	resp, err := s.dataplaneClient.ListNodePoolConfigs(ctx, &dplanemgrv1.ListNodePoolConfigsRequest{})
	if err != nil {
		return nil, nil, err
	}
	return normalizeNodePoolConfigs(resp.GetNodePoolConfigs())
}

func normalizeNodePoolConfigs(configs []*dplanemgrv1.NodePoolConfig) ([]*entity.NodeSpec, map[string]int, error) {
	dropped := make(map[string]int)
	seen := make(map[normalizedNodeSpecKey]struct{})
	nodeSpecs := make([]*entity.NodeSpec, 0, len(configs))

	for _, cfg := range configs {
		if cfg == nil || cfg.GetNodeType() != dplanemgrv1.NodeType_Dedicated {
			dropped[nodeSpecDropUnsupportedNodeType]++
			continue
		}
		if cfg.GetCpuCore() == 0 || cfg.GetMemoryGi() == 0 {
			dropped[nodeSpecDropInvalidCapacity]++
			continue
		}
		availableCPU := cfg.GetAvailableCpuMilliCore() / 1000
		if availableCPU == 0 || cfg.GetAvailableMemoryMi() == 0 {
			dropped[nodeSpecDropInvalidAvailable]++
			continue
		}
		if cfg.GetWithLocalDisk() && cfg.GetLocalDiskSizeGi() == 0 {
			dropped[nodeSpecDropInvalidLocalDisk]++
			continue
		}

		maxCPUMilli := cfg.GetMaxCpuMilli()
		if maxCPUMilli == 0 {
			maxCPUMilli = cfg.GetAvailableCpuMilliCore()
		}
		maxMemoryMi := cfg.GetMaxMemoryMi()
		if maxMemoryMi == 0 {
			maxMemoryMi = cfg.GetAvailableMemoryMi()
		}

		key := normalizedNodeSpecKey{
			cpu:               cfg.GetCpuCore(),
			memoryGi:          cfg.GetMemoryGi(),
			availableCPU:      availableCPU,
			availableMemoryMi: cfg.GetAvailableMemoryMi(),
			withLocalDisk:     cfg.GetWithLocalDisk(),
			localDiskSizeGi:   cfg.GetLocalDiskSizeGi(),
		}
		if _, ok := seen[key]; ok {
			dropped[nodeSpecDropDuplicate]++
			continue
		}
		seen[key] = struct{}{}
		nodeSpecs = append(nodeSpecs, &entity.NodeSpec{
			Cpu:               key.cpu,
			MemoryGi:          key.memoryGi,
			AvailableCpu:      key.availableCPU,
			AvailableCpuMilli: cfg.GetAvailableCpuMilliCore(),
			AvailableMemoryMi: key.availableMemoryMi,
			MaxCpuMilli:       maxCPUMilli,
			MaxMemoryMi:       maxMemoryMi,
			WithLocalDisk:     key.withLocalDisk,
			LocalDiskSizeGi:   key.localDiskSizeGi,
		})
	}

	sort.Slice(nodeSpecs, func(i, j int) bool {
		left, right := nodeSpecs[i], nodeSpecs[j]
		switch {
		case left.GetCpu() != right.GetCpu():
			return left.GetCpu() < right.GetCpu()
		case left.GetMemoryGi() != right.GetMemoryGi():
			return left.GetMemoryGi() < right.GetMemoryGi()
		case left.GetWithLocalDisk() != right.GetWithLocalDisk():
			return !left.GetWithLocalDisk()
		case left.GetLocalDiskSizeGi() != right.GetLocalDiskSizeGi():
			return left.GetLocalDiskSizeGi() < right.GetLocalDiskSizeGi()
		case left.GetAvailableCpu() != right.GetAvailableCpu():
			return left.GetAvailableCpu() < right.GetAvailableCpu()
		case left.GetAvailableMemoryMi() != right.GetAvailableMemoryMi():
			return left.GetAvailableMemoryMi() < right.GetAvailableMemoryMi()
		default:
			return false
		}
	})

	return nodeSpecs, dropped, nil
}

func nodeSpecMapFromList(nodeSpecs []*entity.NodeSpec) map[string][]*entity.NodeSpec {
	nodeSpecMap := make(map[string][]*entity.NodeSpec)
	for _, nodeSpec := range nodeSpecs {
		key := entity.GetComponentNodeSpec(nodeSpec)
		nodeSpecMap[key] = append(nodeSpecMap[key], nodeSpec)
	}
	return nodeSpecMap
}

func resolveBootstrapNodeSpec(candidate *entity.NodeSpec, nodeSpecs []*entity.NodeSpec, withLocalDisk bool) (*entity.NodeSpec, bool) {
	if candidate == nil || candidate.GetCpu() == 0 || candidate.GetMemoryGi() == 0 {
		return nil, false
	}
	isLogicalSelection := candidate.GetAvailableCpu() == 0 &&
		candidate.GetAvailableMemoryMi() == 0 &&
		candidate.GetAvailableCpuMilli() == 0 &&
		candidate.GetMaxCpuMilli() == 0 &&
		candidate.GetMaxMemoryMi() == 0 &&
		!candidate.GetWithLocalDisk() &&
		candidate.GetLocalDiskSizeGi() == 0 &&
		candidate.GetNodeVersion() == "" &&
		candidate.GetArch() == ""

	var matches []*entity.NodeSpec
	for _, nodeSpec := range nodeSpecs {
		if nodeSpec == nil ||
			nodeSpec.GetCpu() != candidate.GetCpu() ||
			nodeSpec.GetMemoryGi() != candidate.GetMemoryGi() ||
			nodeSpec.GetWithLocalDisk() != withLocalDisk {
			continue
		}
		if !isLogicalSelection {
			if proto.Equal(candidate, nodeSpec) {
				return proto.Clone(nodeSpec).(*entity.NodeSpec), true
			}
			continue
		}
		matches = append(matches, nodeSpec)
	}
	if len(matches) == 0 {
		return nil, false
	}
	if withLocalDisk {
		localDiskSizes := make(map[uint32]struct{}, len(matches))
		for _, match := range matches {
			localDiskSizes[match.GetLocalDiskSizeGi()] = struct{}{}
		}
		if len(localDiskSizes) != 1 {
			return nil, false
		}
	}
	return proto.Clone(matches[0]).(*entity.NodeSpec), true
}

func logNodeSpecDrops(ctx context.Context, dropped map[string]int) {
	tilog.Info(ctx, "normalized deployment node specs",
		zap.Int("unsupportedNodeType", dropped[nodeSpecDropUnsupportedNodeType]),
		zap.Int("invalidCapacity", dropped[nodeSpecDropInvalidCapacity]),
		zap.Int("invalidAvailableCapacity", dropped[nodeSpecDropInvalidAvailable]),
		zap.Int("invalidLocalDisk", dropped[nodeSpecDropInvalidLocalDisk]),
		zap.Int("duplicates", dropped[nodeSpecDropDuplicate]),
	)
}