package nodecomponent

import (
	"slices"
	"strconv"
	"strings"

	"github.com/tidbcloud/cluster-service-ng/internal/dataplane-exporter/pkg/source/kubernetes"
)

const (
	adminComponentLabelKey   = "tidbcloud.com/component"
	adminComponentLabelValue = "shared"
	dedicatedTaintKey        = "dedicated"
	workerTaintValue         = "worker"

	tenantLabelKey       = "tags.tidbcloud.com/tenant"
	resourcePoolLabelKey = "tags.tidbcloud.com/resource-pool"
	sharedPoolLabelKey   = "tags.tidbcloud.com/shared-pool"
	componentLabelKey    = "tags.tidbcloud.com/component"
	clusterLabelKey      = "tags.tidbcloud.com/cluster"

	NodeClassAdmin  = "admin"
	NodeClassWorker = "worker"

	StateOwned   = "owned"
	StateMixed   = "mixed"
	StateUnknown = "unknown"
	StateIdle    = "idle"
)

var clusterScopedComponents = map[string]struct{}{
	"tidb":    {},
	"tiproxy": {},
}

type NodeAttribution struct {
	NodeName     string
	NodeClass    string
	Tenant       string
	ResourcePool string

	// Component is the concrete component name for owned nodes, and reuses
	// attribution state values like idle/mixed/unknown for non-owned workers so
	// the exported metric can still distinguish unattributed worker capacity.
	Component string
	Cluster   string
	VCPU      float64
	State     string
}

func AttributeNodes(nodes []kubernetes.Node, pods []kubernetes.Pod) []NodeAttribution {
	podsByNode := make(map[string][]kubernetes.Pod, len(pods))
	for _, pod := range pods {
		podsByNode[pod.NodeName] = append(podsByNode[pod.NodeName], pod)
	}

	sortedNodes := slices.Clone(nodes)
	slices.SortFunc(sortedNodes, func(a, b kubernetes.Node) int {
		return strings.Compare(a.Name, b.Name)
	})

	result := make([]NodeAttribution, 0, len(sortedNodes))
	for _, node := range sortedNodes {
		switch {
		case isAdminNode(node):
			result = append(result, NodeAttribution{
				NodeName:     node.Name,
				NodeClass:    NodeClassAdmin,
				Component:    adminComponentLabelValue,
				VCPU:         node.CapacityCPU,
				State:        StateOwned,
				Tenant:       "",
				ResourcePool: "",
				Cluster:      "",
			})
		case isWorkerNode(node):
			result = append(result, attributeWorkerNode(node, podsByNode[node.Name]))
		}
	}

	return result
}

func attributeWorkerNode(node kubernetes.Node, pods []kubernetes.Pod) NodeAttribution {
	attribution := NodeAttribution{
		NodeName:  node.Name,
		NodeClass: NodeClassWorker,
	}

	billablePods := make([]kubernetes.Pod, 0, len(pods))
	for _, pod := range pods {
		if isBillablePod(pod) {
			billablePods = append(billablePods, pod)
		}
	}
	if len(billablePods) == 0 {
		attribution.VCPU = parseNodeSpecCPU(node.NodeSpec)
		attribution.Component = StateIdle
		attribution.State = StateIdle
		return attribution
	}

	tenantValues := uniqueNonEmptyValues(billablePods, func(p kubernetes.Pod) string {
		return strings.TrimSpace(p.Labels[tenantLabelKey])
	})
	poolValues := uniqueNonEmptyValues(billablePods, poolLabelValue)
	if len(tenantValues) == 1 {
		attribution.Tenant = tenantValues[0]
	}
	if len(poolValues) == 1 {
		attribution.ResourcePool = poolValues[0]
	}
	if len(tenantValues) > 1 || len(poolValues) > 1 {
		attribution.Tenant = ""
		attribution.ResourcePool = ""
		attribution.VCPU = workerNodeVCPU(node)
		attribution.Component = StateUnknown
		attribution.State = StateUnknown
		return attribution
	}

	attribution.VCPU = workerNodeVCPU(node)
	if attribution.VCPU == 0 {
		attribution.Component = StateUnknown
		attribution.State = StateUnknown
		return attribution
	}

	componentValues := uniqueNonEmptyValues(billablePods, func(p kubernetes.Pod) string {
		return strings.TrimSpace(p.Labels[componentLabelKey])
	})
	missingComponent := false
	for _, pod := range billablePods {
		if strings.TrimSpace(pod.Labels[componentLabelKey]) == "" {
			missingComponent = true
			break
		}
	}

	switch {
	case len(componentValues) > 1:
		attribution.Component = StateMixed
		attribution.State = StateMixed
	case len(componentValues) == 1 && !missingComponent:
		attribution.Component = componentValues[0]
		attribution.State = StateOwned
		attribution.Cluster = clusterValueForComponent(componentValues[0], billablePods)
	default:
		attribution.Component = StateUnknown
		attribution.State = StateUnknown
	}

	return attribution
}

func parseNodeSpecCPU(nodeSpec string) float64 {
	nodeSpec = strings.TrimSpace(nodeSpec)
	if nodeSpec == "" {
		return 0
	}

	// Dataplane worker nodes use a node-spec label like "8C16G"; we only need
	// the vCPU prefix before "C" and accept lowercase variants defensively.
	upperNodeSpec := strings.ToUpper(nodeSpec)
	idx := strings.Index(upperNodeSpec, "C")
	if idx <= 0 {
		return 0
	}

	vcpu, err := strconv.Atoi(nodeSpec[:idx])
	if err != nil {
		return 0
	}

	return float64(vcpu)
}

func isAdminNode(node kubernetes.Node) bool {
	return node.Labels[adminComponentLabelKey] == adminComponentLabelValue
}

func isWorkerNode(node kubernetes.Node) bool {
	for _, taint := range node.Taints {
		if taint.Key == dedicatedTaintKey && taint.Value == workerTaintValue {
			return true
		}
	}

	return false
}

func workerNodeVCPU(node kubernetes.Node) float64 {
	return parseNodeSpecCPU(node.NodeSpec)
}

func isBillablePod(pod kubernetes.Pod) bool {
	if pod.OwnerKind == "DaemonSet" {
		return false
	}
	if pod.Phase == "Succeeded" || pod.Phase == "Failed" {
		return false
	}

	return strings.TrimSpace(pod.Labels[tenantLabelKey]) != "" &&
		poolLabelValue(pod) != ""
}

func poolLabelValue(pod kubernetes.Pod) string {
	if value := strings.TrimSpace(pod.Labels[resourcePoolLabelKey]); value != "" {
		return value
	}

	return strings.TrimSpace(pod.Labels[sharedPoolLabelKey])
}

func clusterValueForComponent(component string, pods []kubernetes.Pod) string {
	if _, ok := clusterScopedComponents[component]; !ok {
		return ""
	}

	clusterMissing := false
	for _, pod := range pods {
		if strings.TrimSpace(pod.Labels[componentLabelKey]) != component {
			continue
		}
		if strings.TrimSpace(pod.Labels[clusterLabelKey]) == "" {
			clusterMissing = true
			break
		}
	}
	if clusterMissing {
		return ""
	}

	clusterValues := uniqueNonEmptyValues(pods, func(p kubernetes.Pod) string {
		if strings.TrimSpace(p.Labels[componentLabelKey]) != component {
			return ""
		}
		return strings.TrimSpace(p.Labels[clusterLabelKey])
	})
	if len(clusterValues) != 1 {
		return ""
	}

	return clusterValues[0]
}

func uniqueNonEmptyValues[T any](items []T, getValue func(T) string) []string {
	seen := make(map[string]struct{}, len(items))
	values := make([]string, 0, len(items))
	for _, item := range items {
		value := getValue(item)
		if value == "" {
			continue
		}
		if _, ok := seen[value]; ok {
			continue
		}
		seen[value] = struct{}{}
		values = append(values, value)
	}

	slices.Sort(values)
	return values
}