package server

import (
	"strings"

	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	apicommon "github.com/tidbcloud/infra-provider/api/tidbcloud/common"
	"github.com/tidbcloud/infra-provider/pkg/gkeheterogeneous"
)

type gkeHeterogeneousNodePoolBindingState struct {
	computeClassName    string
	bindings            []cloudv1.GKEComputeClassNodePoolBinding
	hasExplicitBindings bool
}

func loadGKEHeterogeneousNodePoolBindingState(cluster *cloudv1.Cluster, intentConfig gkeHeterogeneousIntentConfig, tidbGroupID string) (*gkeHeterogeneousNodePoolBindingState, error) {
	baseInstanceType := baseGKEHeterogeneousInstanceType(cluster, intentConfig.componentType, tidbGroupID)
	if baseInstanceType == "" {
		return nil, status.Error(codes.FailedPrecondition, "base instance type is empty")
	}
	requireExplicitAnnotation := false
	switch intentConfig.componentType {
	case apicommon.ComponentType_TIKV:
		requireExplicitAnnotation = cluster.TiKVGKEComputeClassIntentConfigured()
	case apicommon.ComponentType_TIDB:
		requireExplicitAnnotation = cluster.TiDBGKEComputeClassIntentConfigured(tidbGroupID)
	}
	annotationKey := gkeHeterogeneousIntentAnnotationKey(cluster, intentConfig, tidbGroupID)
	bindingAnnotation, hasExplicitBindings, err := cloudv1.ClusterGKEComputeClassBindingAnnotation(
		cluster,
		annotationKey,
		baseInstanceType,
		gkeheterogeneous.NodePoolName(cluster, baseInstanceType),
	)
	if err != nil {
		return nil, status.Error(codes.FailedPrecondition, err.Error())
	}
	if requireExplicitAnnotation && !hasExplicitBindings {
		return nil, status.Errorf(codes.FailedPrecondition, "gke compute class binding annotation %s is required when heterogeneous intent is enabled", annotationKey)
	}
	return &gkeHeterogeneousNodePoolBindingState{
		computeClassName:    bindingAnnotation.ComputeClassName,
		bindings:            bindingAnnotation.NodePools,
		hasExplicitBindings: hasExplicitBindings,
	}, nil
}

func (s *gkeHeterogeneousNodePoolBindingState) InstanceTypes() []string {
	if !s.hasExplicitBindings {
		return nil
	}
	out := make([]string, 0, len(s.bindings))
	for _, binding := range s.bindings {
		out = append(out, binding.InstanceType)
	}
	return out
}

func (s *gkeHeterogeneousNodePoolBindingState) NodePoolNames() []string {
	if !s.hasExplicitBindings {
		return nil
	}
	out := make([]string, 0, len(s.bindings))
	for _, binding := range s.bindings {
		if binding.NodePoolName == "" {
			continue
		}
		out = append(out, binding.NodePoolName)
	}
	return out
}

func (s *gkeHeterogeneousNodePoolBindingState) ExplicitBindings() []cloudv1.GKEComputeClassNodePoolBinding {
	if !s.hasExplicitBindings {
		return nil
	}
	out := make([]cloudv1.GKEComputeClassNodePoolBinding, 0, len(s.bindings))
	for _, binding := range s.bindings {
		if binding.NodePoolName == "" {
			continue
		}
		out = append(out, binding)
	}
	return out
}

func (s *gkeHeterogeneousNodePoolBindingState) IsBaseBinding(binding cloudv1.GKEComputeClassNodePoolBinding) bool {
	return strings.TrimSpace(binding.Role) == cloudv1.GKEComputeClassNodePoolRoleBase
}

func buildGKEHeterogeneousIntentBindingAnnotation(cluster *cloudv1.Cluster, intentConfig gkeHeterogeneousIntentConfig, tidbGroupID, computeClassName string, nodePoolNames, instanceTypes []string) (string, cloudv1.GKEComputeClassBindingAnnotation, error) {
	baseInstanceType := strings.TrimSpace(baseGKEHeterogeneousInstanceType(cluster, intentConfig.componentType, tidbGroupID))
	if baseInstanceType == "" {
		return "", cloudv1.GKEComputeClassBindingAnnotation{}, status.Errorf(codes.FailedPrecondition, "base instance type is empty for %s heterogeneous deploy", intentConfig.componentType)
	}

	extraInstanceTypes := make([]string, 0, len(instanceTypes))
	for _, instanceType := range instanceTypes {
		instanceType = strings.TrimSpace(instanceType)
		if instanceType != "" {
			extraInstanceTypes = append(extraInstanceTypes, instanceType)
		}
	}
	if len(extraInstanceTypes) != 1 {
		return "", cloudv1.GKEComputeClassBindingAnnotation{}, status.Errorf(codes.InvalidArgument, "instance_types must contain exactly one extra instance type")
	}
	if extraInstanceTypes[0] == baseInstanceType {
		return "", cloudv1.GKEComputeClassBindingAnnotation{}, status.Errorf(codes.InvalidArgument, "instance_types must contain exactly one extra instance type")
	}

	// Intent is persisted before extra NodePool resources are reconciled, so
	// target scope must be complete up front.
	if _, err := requireGKEHeterogeneousNodePoolScopeLabels(cluster); err != nil {
		return "", cloudv1.GKEComputeClassBindingAnnotation{}, err
	}

	orderedInstanceTypes := []string{baseInstanceType, extraInstanceTypes[0]}
	generatedNodePoolNames := []string{
		gkeheterogeneous.NodePoolName(cluster, baseInstanceType),
		gkeheterogeneous.NodePoolName(cluster, extraInstanceTypes[0]),
	}
	normalizedNodePoolNames, err := resolveGKEHeterogeneousIntentNodePoolNames(nodePoolNames, generatedNodePoolNames)
	if err != nil {
		return "", cloudv1.GKEComputeClassBindingAnnotation{}, err
	}

	bindings, err := cloudv1.NormalizeGKEComputeClassOrderedBindings([]cloudv1.GKEComputeClassNodePoolBinding{
		{
			Role:         cloudv1.GKEComputeClassNodePoolRoleBase,
			InstanceType: orderedInstanceTypes[0],
			NodePoolName: normalizedNodePoolNames[0],
		},
		{
			Role:         cloudv1.GKEComputeClassNodePoolRoleExtra,
			InstanceType: orderedInstanceTypes[1],
			NodePoolName: normalizedNodePoolNames[1],
		},
	})
	if err != nil {
		return "", cloudv1.GKEComputeClassBindingAnnotation{}, status.Error(codes.InvalidArgument, err.Error())
	}
	expectedComputeClassName, err := cloudv1.GKEComputeClassNameForBindings(bindings)
	if err != nil {
		return "", cloudv1.GKEComputeClassBindingAnnotation{}, status.Error(codes.InvalidArgument, err.Error())
	}
	if strings.TrimSpace(computeClassName) != "" {
		normalizedComputeClassName, normalizeErr := cloudv1.NormalizeGKEComputeClassName(computeClassName)
		if normalizeErr != nil {
			return "", cloudv1.GKEComputeClassBindingAnnotation{}, status.Error(codes.InvalidArgument, normalizeErr.Error())
		}
		if normalizedComputeClassName != expectedComputeClassName {
			return "", cloudv1.GKEComputeClassBindingAnnotation{}, status.Errorf(codes.InvalidArgument, "compute_class_name %s does not match generated compute class name %s", normalizedComputeClassName, expectedComputeClassName)
		}
	}

	bindingAnnotation := cloudv1.GKEComputeClassBindingAnnotation{
		ComputeClassName: expectedComputeClassName,
		NodePools:        bindings,
	}
	return baseInstanceType, bindingAnnotation, nil
}

func resolveGKEHeterogeneousIntentNodePoolNames(requestedNodePoolNames, expectedNodePoolNames []string) ([]string, error) {
	normalizedNodePoolNames, err := cloudv1.NormalizeGKENodePoolNames(requestedNodePoolNames)
	if err != nil {
		return nil, status.Error(codes.InvalidArgument, err.Error())
	}
	if len(normalizedNodePoolNames) == 0 {
		return expectedNodePoolNames, nil
	}
	if len(normalizedNodePoolNames) != len(expectedNodePoolNames) {
		return nil, status.Errorf(codes.InvalidArgument, "node_pool_names has %d items, want %d to match base instance type plus instance_types", len(normalizedNodePoolNames), len(expectedNodePoolNames))
	}
	for i, expectedNodePoolName := range expectedNodePoolNames {
		if normalizedNodePoolNames[i] != expectedNodePoolName {
			return nil, status.Errorf(codes.InvalidArgument, "node_pool_names[%d]=%s does not match generated nodepool name %s", i, normalizedNodePoolNames[i], expectedNodePoolName)
		}
	}
	return normalizedNodePoolNames, nil
}