package nodecomponent

import (
	"testing"

	"github.com/stretchr/testify/require"

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

func workerNode(name string, capacity float64, nodeSpec string) kubernetes.Node {
	return kubernetes.Node{
		Name:        name,
		CapacityCPU: capacity,
		NodeSpec:    nodeSpec,
		Taints: []kubernetes.Taint{{
			Key:   dedicatedTaintKey,
			Value: workerTaintValue,
		}},
	}
}

func TestAttributeNodes(t *testing.T) {
	tests := []struct {
		name     string
		nodes    []kubernetes.Node
		pods     []kubernetes.Pod
		expected []NodeAttribution
	}{
		{
			name: "admin node uses capacity cpu and shared component",
			nodes: []kubernetes.Node{{
				Name: "admin-1",
				Labels: map[string]string{
					adminComponentLabelKey: adminComponentLabelValue,
				},
				CapacityCPU: 4,
			}},
			expected: []NodeAttribution{{
				NodeName:     "admin-1",
				NodeClass:    NodeClassAdmin,
				Component:    adminComponentLabelValue,
				VCPU:         4,
				State:        StateOwned,
				Cluster:      "",
				Tenant:       "",
				ResourcePool: "",
			}},
		},
		{
			name:  "worker node unique tidb component and cluster",
			nodes: []kubernetes.Node{workerNode("worker-1", 6, "8C16G")},
			pods: []kubernetes.Pod{{
				Name:     "tidb-0",
				NodeName: "worker-1",
				Labels: map[string]string{
					tenantLabelKey:       "t1",
					resourcePoolLabelKey: "rp1",
					componentLabelKey:    "tidb",
					clusterLabelKey:      "10903881287785916858",
				},
				OwnerKind: "StatefulSet",
				Phase:     "Running",
			}},
			expected: []NodeAttribution{{
				NodeName:     "worker-1",
				NodeClass:    NodeClassWorker,
				Tenant:       "t1",
				ResourcePool: "rp1",
				Component:    "tidb",
				Cluster:      "10903881287785916858",
				VCPU:         8,
				State:        StateOwned,
			}},
		},
		{
			name:  "worker node with multiple components becomes mixed",
			nodes: []kubernetes.Node{workerNode("worker-1", 6, "8C16G")},
			pods: []kubernetes.Pod{
				{
					Name:     "tidb-0",
					NodeName: "worker-1",
					Labels: map[string]string{
						tenantLabelKey:       "t1",
						resourcePoolLabelKey: "rp1",
						componentLabelKey:    "tidb",
					},
					OwnerKind: "StatefulSet",
					Phase:     "Running",
				},
				{
					Name:     "tikv-0",
					NodeName: "worker-1",
					Labels: map[string]string{
						tenantLabelKey:       "t1",
						resourcePoolLabelKey: "rp1",
						componentLabelKey:    "tikv",
					},
					OwnerKind: "StatefulSet",
					Phase:     "Running",
				},
			},
			expected: []NodeAttribution{{
				NodeName:     "worker-1",
				NodeClass:    NodeClassWorker,
				Tenant:       "t1",
				ResourcePool: "rp1",
				Component:    StateMixed,
				VCPU:         8,
				State:        StateMixed,
			}},
		},
		{
			name:  "worker node with billable pod but missing component becomes unknown",
			nodes: []kubernetes.Node{workerNode("worker-1", 6, "8C16G")},
			pods: []kubernetes.Pod{{
				Name:     "mystery-0",
				NodeName: "worker-1",
				Labels: map[string]string{
					tenantLabelKey:       "t1",
					resourcePoolLabelKey: "rp1",
				},
				OwnerKind: "StatefulSet",
				Phase:     "Running",
			}},
			expected: []NodeAttribution{{
				NodeName:     "worker-1",
				NodeClass:    NodeClassWorker,
				Tenant:       "t1",
				ResourcePool: "rp1",
				Component:    StateUnknown,
				VCPU:         8,
				State:        StateUnknown,
			}},
		},
		{
			name:     "worker node with no billable pod becomes idle",
			nodes:    []kubernetes.Node{workerNode("worker-1", 6, "8C16G")},
			pods:     []kubernetes.Pod{{Name: "ds-0", NodeName: "worker-1", OwnerKind: "DaemonSet", Phase: "Running"}},
			expected: []NodeAttribution{{NodeName: "worker-1", NodeClass: NodeClassWorker, Component: StateIdle, VCPU: 8, State: StateIdle}},
		},
		{
			name:  "worker node with unique tiproxy component but conflicting cluster labels keeps cluster empty",
			nodes: []kubernetes.Node{workerNode("worker-1", 6, "8C16G")},
			pods: []kubernetes.Pod{
				{
					Name:     "tiproxy-0",
					NodeName: "worker-1",
					Labels: map[string]string{
						tenantLabelKey:       "t1",
						resourcePoolLabelKey: "rp1",
						componentLabelKey:    "tiproxy",
						clusterLabelKey:      "c1",
					},
					OwnerKind: "StatefulSet",
					Phase:     "Running",
				},
				{
					Name:     "tiproxy-1",
					NodeName: "worker-1",
					Labels: map[string]string{
						tenantLabelKey:       "t1",
						resourcePoolLabelKey: "rp1",
						componentLabelKey:    "tiproxy",
						clusterLabelKey:      "c2",
					},
					OwnerKind: "StatefulSet",
					Phase:     "Running",
				},
			},
			expected: []NodeAttribution{{
				NodeName:     "worker-1",
				NodeClass:    NodeClassWorker,
				Tenant:       "t1",
				ResourcePool: "rp1",
				Component:    "tiproxy",
				Cluster:      "",
				VCPU:         8,
				State:        StateOwned,
			}},
		},
		{
			name:  "worker node with unique tidb component but missing cluster label keeps cluster empty",
			nodes: []kubernetes.Node{workerNode("worker-1", 6, "8C16G")},
			pods: []kubernetes.Pod{
				{
					Name:     "tidb-0",
					NodeName: "worker-1",
					Labels: map[string]string{
						tenantLabelKey:       "t1",
						resourcePoolLabelKey: "rp1",
						componentLabelKey:    "tidb",
						clusterLabelKey:      "c1",
					},
					OwnerKind: "StatefulSet",
					Phase:     "Running",
				},
				{
					Name:     "tidb-1",
					NodeName: "worker-1",
					Labels: map[string]string{
						tenantLabelKey:       "t1",
						resourcePoolLabelKey: "rp1",
						componentLabelKey:    "tidb",
					},
					OwnerKind: "StatefulSet",
					Phase:     "Running",
				},
			},
			expected: []NodeAttribution{{
				NodeName:     "worker-1",
				NodeClass:    NodeClassWorker,
				Tenant:       "t1",
				ResourcePool: "rp1",
				Component:    "tidb",
				Cluster:      "",
				VCPU:         8,
				State:        StateOwned,
			}},
		},
		{
			name:  "worker node with conflicting tenant or pool becomes unknown",
			nodes: []kubernetes.Node{workerNode("worker-1", 6, "8C16G")},
			pods: []kubernetes.Pod{
				{
					Name:     "tidb-0",
					NodeName: "worker-1",
					Labels: map[string]string{
						tenantLabelKey:       "t1",
						resourcePoolLabelKey: "rp1",
						componentLabelKey:    "tidb",
					},
					OwnerKind: "StatefulSet",
					Phase:     "Running",
				},
				{
					Name:     "tidb-1",
					NodeName: "worker-1",
					Labels: map[string]string{
						tenantLabelKey:       "t2",
						resourcePoolLabelKey: "rp2",
						componentLabelKey:    "tidb",
					},
					OwnerKind: "StatefulSet",
					Phase:     "Running",
				},
			},
			expected: []NodeAttribution{{
				NodeName:     "worker-1",
				NodeClass:    NodeClassWorker,
				Component:    StateUnknown,
				VCPU:         8,
				State:        StateUnknown,
				Cluster:      "",
				Tenant:       "",
				ResourcePool: "",
			}},
		},
		{
			name:  "worker node with conflicting tenants and unique pool clears partial labels",
			nodes: []kubernetes.Node{workerNode("worker-1", 6, "8C16G")},
			pods: []kubernetes.Pod{
				{
					Name:     "tidb-0",
					NodeName: "worker-1",
					Labels: map[string]string{
						tenantLabelKey:       "t1",
						resourcePoolLabelKey: "rp1",
						componentLabelKey:    "tidb",
					},
					OwnerKind: "StatefulSet",
					Phase:     "Running",
				},
				{
					Name:     "tidb-1",
					NodeName: "worker-1",
					Labels: map[string]string{
						tenantLabelKey:       "t2",
						resourcePoolLabelKey: "rp1",
						componentLabelKey:    "tidb",
					},
					OwnerKind: "StatefulSet",
					Phase:     "Running",
				},
			},
			expected: []NodeAttribution{{
				NodeName:     "worker-1",
				NodeClass:    NodeClassWorker,
				Component:    StateUnknown,
				VCPU:         8,
				State:        StateUnknown,
				Cluster:      "",
				Tenant:       "",
				ResourcePool: "",
			}},
		},
		{
			name:  "worker node with unique tenant and conflicting pools clears partial labels",
			nodes: []kubernetes.Node{workerNode("worker-1", 6, "8C16G")},
			pods: []kubernetes.Pod{
				{
					Name:     "tidb-0",
					NodeName: "worker-1",
					Labels: map[string]string{
						tenantLabelKey:       "t1",
						resourcePoolLabelKey: "rp1",
						componentLabelKey:    "tidb",
					},
					OwnerKind: "StatefulSet",
					Phase:     "Running",
				},
				{
					Name:     "tidb-1",
					NodeName: "worker-1",
					Labels: map[string]string{
						tenantLabelKey:       "t1",
						resourcePoolLabelKey: "rp2",
						componentLabelKey:    "tidb",
					},
					OwnerKind: "StatefulSet",
					Phase:     "Running",
				},
			},
			expected: []NodeAttribution{{
				NodeName:     "worker-1",
				NodeClass:    NodeClassWorker,
				Component:    StateUnknown,
				VCPU:         8,
				State:        StateUnknown,
				Cluster:      "",
				Tenant:       "",
				ResourcePool: "",
			}},
		},
		{
			name:  "worker node with missing node spec becomes unknown and zero vcpu",
			nodes: []kubernetes.Node{workerNode("worker-1", 6, "")},
			pods: []kubernetes.Pod{{
				Name:     "tidb-0",
				NodeName: "worker-1",
				Labels: map[string]string{
					tenantLabelKey:       "t1",
					resourcePoolLabelKey: "rp1",
					componentLabelKey:    "tidb",
				},
				OwnerKind: "StatefulSet",
				Phase:     "Running",
			}},
			expected: []NodeAttribution{{
				NodeName:     "worker-1",
				NodeClass:    NodeClassWorker,
				Tenant:       "t1",
				ResourcePool: "rp1",
				Component:    StateUnknown,
				VCPU:         0,
				State:        StateUnknown,
			}},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			require.Equal(t, tt.expected, AttributeNodes(tt.nodes, tt.pods))
		})
	}
}