package server

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"

	agentv2 "github.com/tidbcloud/observability/agent-control-server/api/v2"
	"github.com/tidbcloud/observability/agent-control-server/pkg/consts"
)

type filterTestCase struct {
	name     string
	filter   agentFilter
	expected bool
}

func runFilterTests[T any](t *testing.T, agent T, matchFunc func(T, agentFilter) *AgentInfoResponse, tests []filterTestCase) {
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := matchFunc(agent, tt.filter)
			if tt.expected {
				assert.NotNil(t, result)
			} else {
				assert.Nil(t, result)
			}
		})
	}
}

func commonFilterTests(agentName, o11yID string) []filterTestCase {
	return []filterTestCase{
		{"no filter, match all", agentFilter{}, true},
		{"filter by agent name, match", agentFilter{AgentName: agentName}, true},
		{"filter by agent name, no match", agentFilter{AgentName: "other-agent"}, false},
		{"filter by cluster name, match", agentFilter{KubernetesClusterName: "test-cluster"}, true},
		{"filter by cluster name, no match", agentFilter{KubernetesClusterName: "other-cluster"}, false},
		{"filter by tenant id, match", agentFilter{TenantID: "tenant-1"}, true},
		{"filter by tenant id, no match", agentFilter{TenantID: "tenant-2"}, false},
		{"filter by o11y id, match", agentFilter{O11yID: o11yID}, true},
		{"filter by o11y id, no match", agentFilter{O11yID: "o11y-000"}, false},
		{"multiple filters, all match", agentFilter{AgentName: agentName, TenantID: "tenant-1", KubernetesClusterName: "test-cluster", O11yID: o11yID}, true},
		{"multiple filters, one fails", agentFilter{TenantID: "tenant-1", KubernetesClusterName: "other-cluster"}, false},
	}
}

func TestMatchKubernetesAgent(t *testing.T) {
	agent := agentv2.KubernetesAgent{
		ObjectMeta: metav1.ObjectMeta{
			Name: "test-agent",
			Annotations: map[string]string{
				consts.AnnotationKeyO11yID: "o11y-123",
			},
		},
		Spec: agentv2.KubernetesAgentSpec{
			KubernetesClusterName: "test-cluster",
			O11y: agentv2.O11yConfigV2{
				TenantID:      "tenant-1",
				ApplicationID: "app-1",
			},
		},
		Status: agentv2.KubernetesAgentStatus{Phase: "Normal"},
	}

	runFilterTests(t, agent, matchKubernetesAgent, commonFilterTests("test-agent", "o11y-123"))
}

func TestMatchO11yK8sAgent(t *testing.T) {
	agent := agentv2.O11yK8sAgent{
		ObjectMeta: metav1.ObjectMeta{
			Name: "test-o11yk8s-agent",
			Annotations: map[string]string{
				consts.AnnotationKeyO11yID: "o11y-456",
			},
		},
		Spec: agentv2.O11yK8sAgentSpec{
			KubernetesClusterName: "test-cluster",
			O11y: agentv2.O11yConfig{
				TenantID:      "tenant-1",
				ProjectID:     "project-1",
				ApplicationID: "app-1",
			},
		},
		Status: agentv2.O11yK8sAgentStatus{Phase: "Normal"},
	}

	runFilterTests(t, agent, matchO11yK8sAgent, commonFilterTests("test-o11yk8s-agent", "o11y-456"))
}

func TestMatchO11yTiDBAgent(t *testing.T) {
	agent := agentv2.O11yTiDBAgent{
		ObjectMeta: metav1.ObjectMeta{
			Name: "test-o11ytidb-agent",
			Annotations: map[string]string{
				consts.AnnotationKeyO11yID: "o11y-789",
			},
		},
		Spec: agentv2.O11yTiDBAgentSpec{
			KubernetesClusterName: "test-cluster",
			O11y: agentv2.O11yConfig{
				TenantID:      "tenant-1",
				ProjectID:     "project-1",
				ApplicationID: "app-1",
			},
		},
		Status: agentv2.O11yTiDBAgentStatus{Phase: "Normal"},
	}

	runFilterTests(t, agent, matchO11yTiDBAgent, commonFilterTests("test-o11ytidb-agent", "o11y-789"))
}

func TestMatchSharedPoolAgent(t *testing.T) {
	agent := agentv2.SharedPoolAgent{
		ObjectMeta: metav1.ObjectMeta{
			Name: "test-sharedpool-agent",
			Annotations: map[string]string{
				consts.AnnotationKeyO11yID: "o11y-spool",
			},
		},
		Spec: agentv2.SharedPoolAgentSpec{
			KubernetesClusterName: "test-cluster",
			SharedPoolID:          "spool-1",
			O11y: agentv2.O11yConfigV2{
				TenantID:      "tenant-1",
				ApplicationID: "app-1",
			},
		},
		Status: agentv2.SharedPoolAgentStatus{Phase: "Normal"},
	}

	runFilterTests(t, agent, matchSharedPoolAgent, commonFilterTests("test-sharedpool-agent", "o11y-spool"))
}

func TestMatchTidbGroupAgent(t *testing.T) {
	agent := agentv2.TiDBGroupAgent{
		ObjectMeta: metav1.ObjectMeta{
			Name: "test-tidbgroup-agent",
			Annotations: map[string]string{
				consts.AnnotationKeyO11yID: "o11y-tgroup",
			},
		},
		Spec: agentv2.TiDBGroupAgentSpec{
			KubernetesClusterName: "test-cluster",
			SharedPoolID:          "spool-1",
			O11y: agentv2.O11yConfigV2{
				TenantID:      "tenant-1",
				ApplicationID: "app-1",
			},
		},
		Status: agentv2.TiDBGroupAgentStatus{Phase: "Normal"},
	}

	runFilterTests(t, agent, matchTidbGroupAgent, commonFilterTests("test-tidbgroup-agent", "o11y-tgroup"))
}

func TestRouteGetAgentSpec(t *testing.T) {
	tests := []struct {
		name           string
		agentType      string
		agentName      string
		setupResources []client.Object
		expectedCode   int
		expectedSpec   AgentSpecResponse
	}{
		{
			name:      "get kubernetes agent spec with annotations",
			agentType: consts.AgentKubernetes,
			agentName: "test-agent",
			setupResources: []client.Object{
				&agentv2.KubernetesAgent{
					ObjectMeta: metav1.ObjectMeta{
						Name: "test-agent",
						Annotations: map[string]string{
							consts.AnnotationKeyO11yDesiredSpec: "large",
							consts.AnnotationKeyO11yMinSpec:     "small",
							consts.AnnotationKeyO11yMaxSpec:     "xlarge",
						},
					},
				},
			},
			expectedCode: http.StatusOK,
			expectedSpec: AgentSpecResponse{
				AgentType:   consts.AgentKubernetes,
				AgentName:   "test-agent",
				DesiredSpec: "large",
				MinSpec:     "small",
				MaxSpec:     "xlarge",
			},
		},
		{
			name:      "get agent spec with no annotations",
			agentType: consts.AgentKubernetes,
			agentName: "test-agent-no-annotations",
			setupResources: []client.Object{
				&agentv2.KubernetesAgent{
					ObjectMeta: metav1.ObjectMeta{
						Name: "test-agent-no-annotations",
					},
				},
			},
			expectedCode: http.StatusOK,
			expectedSpec: AgentSpecResponse{
				AgentType: consts.AgentKubernetes,
				AgentName: "test-agent-no-annotations",
			},
		},
		{
			name:           "agent not found",
			agentType:      consts.AgentKubernetes,
			agentName:      "non-existent",
			setupResources: []client.Object{},
			expectedCode:   http.StatusBadRequest,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := makeTestServer(tt.setupResources...)
			path := fmt.Sprintf("/agents/spec?agentType=%s&agentName=%s", tt.agentType, tt.agentName)
			req := makeGetRequest(path)

			w := httptest.NewRecorder()
			s.routeGetAgentSpec(w, req)

			assert.Equal(t, tt.expectedCode, w.Code)
			if tt.expectedCode != http.StatusOK {
				return
			}
			rsp := parseResponse(w)
			assert.Equal(t, tt.expectedSpec, parseResponseData[AgentSpecResponse](rsp))
		})
	}
}

func TestRoutePostAgentSpec(t *testing.T) {
	tests := []struct {
		name               string
		agentType          string
		agentName          string
		requestBody        AgentSpecRequest
		setupResources     []client.Object
		expectedCode       int
		expectedResponse   AgentSpecResponse
		expectedAnnotation map[string]string
	}{
		{
			name:      "update kubernetes agent spec with desiredSpec only",
			agentType: consts.AgentKubernetes,
			agentName: "test-agent",
			requestBody: AgentSpecRequest{
				AgentType:   consts.AgentKubernetes,
				AgentName:   "test-agent",
				DesiredSpec: "large",
			},
			setupResources: []client.Object{
				&agentv2.KubernetesAgent{
					ObjectMeta: metav1.ObjectMeta{Name: "test-agent"},
				},
			},
			expectedCode: http.StatusOK,
			expectedResponse: AgentSpecResponse{
				AgentType:   consts.AgentKubernetes,
				AgentName:   "test-agent",
				DesiredSpec: "large",
			},
			expectedAnnotation: map[string]string{
				consts.AnnotationKeyO11yDesiredSpec: "large",
			},
		},
		{
			name:      "update o11yk8s agent spec sets previous spec and scale status",
			agentType: consts.AgentO11yK8s,
			agentName: "test-o11yk8s-agent",
			requestBody: AgentSpecRequest{
				AgentType:   consts.AgentO11yK8s,
				AgentName:   "test-o11yk8s-agent",
				DesiredSpec: "large",
			},
			setupResources: []client.Object{
				&agentv2.O11yK8sAgent{
					ObjectMeta: metav1.ObjectMeta{
						Name: "test-o11yk8s-agent",
						Annotations: map[string]string{
							consts.AnnotationKeyO11yDesiredSpec: "medium",
						},
					},
				},
			},
			expectedCode: http.StatusOK,
			expectedResponse: AgentSpecResponse{
				AgentType:   consts.AgentO11yK8s,
				AgentName:   "test-o11yk8s-agent",
				DesiredSpec: "large",
			},
			expectedAnnotation: map[string]string{
				consts.AnnotationKeyO11yDesiredSpec:     "large",
				consts.AnnotationKeyO11yPreviousSpec:    "medium",
				consts.AnnotationKeyO11ySpecScaleStatus: consts.SpecScaleStatusPendingScale,
			},
		},
		{
			name:      "update o11yk8s agent spec with default previous spec",
			agentType: consts.AgentO11yK8s,
			agentName: "test-o11yk8s-agent-no-spec",
			requestBody: AgentSpecRequest{
				AgentType:   consts.AgentO11yK8s,
				AgentName:   "test-o11yk8s-agent-no-spec",
				DesiredSpec: "large",
			},
			setupResources: []client.Object{
				&agentv2.O11yK8sAgent{
					ObjectMeta: metav1.ObjectMeta{
						Name: "test-o11yk8s-agent-no-spec",
					},
				},
			},
			expectedCode: http.StatusOK,
			expectedResponse: AgentSpecResponse{
				AgentType:   consts.AgentO11yK8s,
				AgentName:   "test-o11yk8s-agent-no-spec",
				DesiredSpec: "large",
			},
			expectedAnnotation: map[string]string{
				consts.AnnotationKeyO11yDesiredSpec:     "large",
				consts.AnnotationKeyO11yPreviousSpec:    "default",
				consts.AnnotationKeyO11ySpecScaleStatus: consts.SpecScaleStatusPendingScale,
			},
		},
		{
			name:      "update all spec fields",
			agentType: consts.AgentKubernetes,
			agentName: "test-agent-all",
			requestBody: AgentSpecRequest{
				AgentType:   consts.AgentKubernetes,
				AgentName:   "test-agent-all",
				DesiredSpec: "large",
				MinSpec:     "small",
				MaxSpec:     "xlarge",
			},
			setupResources: []client.Object{
				&agentv2.KubernetesAgent{
					ObjectMeta: metav1.ObjectMeta{Name: "test-agent-all"},
				},
			},
			expectedCode: http.StatusOK,
			expectedResponse: AgentSpecResponse{
				AgentType:   consts.AgentKubernetes,
				AgentName:   "test-agent-all",
				DesiredSpec: "large",
				MinSpec:     "small",
				MaxSpec:     "xlarge",
			},
			expectedAnnotation: map[string]string{
				consts.AnnotationKeyO11yDesiredSpec: "large",
				consts.AnnotationKeyO11yMinSpec:     "small",
				consts.AnnotationKeyO11yMaxSpec:     "xlarge",
			},
		},
		{
			name:      "empty request body returns error",
			agentType: consts.AgentKubernetes,
			agentName: "test-agent",
			requestBody: AgentSpecRequest{
				AgentType: consts.AgentKubernetes,
				AgentName: "test-agent",
			},
			setupResources: []client.Object{
				&agentv2.KubernetesAgent{
					ObjectMeta: metav1.ObjectMeta{Name: "test-agent"},
				},
			},
			expectedCode: http.StatusBadRequest,
		},
		{
			name:      "missing agentType returns error",
			agentType: consts.AgentKubernetes,
			agentName: "test-agent",
			requestBody: AgentSpecRequest{
				AgentName:   "test-agent",
				DesiredSpec: "large",
			},
			setupResources: []client.Object{
				&agentv2.KubernetesAgent{
					ObjectMeta: metav1.ObjectMeta{Name: "test-agent"},
				},
			},
			expectedCode: http.StatusBadRequest,
		},
		{
			name:      "agent not found",
			agentType: consts.AgentKubernetes,
			agentName: "non-existent",
			requestBody: AgentSpecRequest{
				AgentType:   consts.AgentKubernetes,
				AgentName:   "non-existent",
				DesiredSpec: "large",
			},
			setupResources: []client.Object{},
			expectedCode:   http.StatusBadRequest,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := makeTestServer(tt.setupResources...)
			req := makePostRequest("/agents/spec", tt.requestBody)

			w := httptest.NewRecorder()
			s.routePostAgentSpec(w, req)

			assert.Equal(t, tt.expectedCode, w.Code)
			if tt.expectedCode != http.StatusOK {
				return
			}

			rsp := parseResponse(w)
			assert.Equal(t, tt.expectedResponse, parseResponseData[AgentSpecResponse](rsp))
			if tt.expectedAnnotation != nil {
				obj, err := s.getAgentObject(context.Background(), tt.agentType, tt.agentName)
				require.NoError(t, err)
				for key, value := range tt.expectedAnnotation {
					assert.Equal(t, value, obj.GetAnnotations()[key])
				}
			}
		})
	}
}

func TestRoutePostPauseAgent(t *testing.T) {
	tests := []struct {
		name             string
		agentType        string
		agentName        string
		setupResources   []client.Object
		expectedCode     int
		expectedResponse AgentPauseResponse
	}{
		{
			name:      "pause kubernetes agent",
			agentType: consts.AgentKubernetes,
			agentName: "test-agent",
			setupResources: []client.Object{
				&agentv2.KubernetesAgent{
					ObjectMeta: metav1.ObjectMeta{Name: "test-agent"},
				},
			},
			expectedCode: http.StatusOK,
			expectedResponse: AgentPauseResponse{
				AgentType: consts.AgentKubernetes,
				AgentName: "test-agent",
				IsPaused:  true,
			},
		},
		{
			name:           "agent not found",
			agentType:      consts.AgentKubernetes,
			agentName:      "non-existent",
			setupResources: []client.Object{},
			expectedCode:   http.StatusBadRequest,
		},
		{
			name:           "missing agentType and agentName returns error",
			agentType:      "",
			agentName:      "",
			setupResources: []client.Object{},
			expectedCode:   http.StatusBadRequest,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := makeTestServer(tt.setupResources...)
			body := AgentActionRequest{AgentType: tt.agentType, AgentName: tt.agentName}
			req := makePostRequest("/agents/pause", body)

			w := httptest.NewRecorder()
			s.routePostPauseAgent(w, req)

			assert.Equal(t, tt.expectedCode, w.Code)
			if tt.expectedCode != http.StatusOK {
				return
			}

			rsp := parseResponse(w)
			assert.Equal(t, tt.expectedResponse, parseResponseData[AgentPauseResponse](rsp))

			obj, err := s.getAgentObject(context.Background(), tt.agentType, tt.agentName)
			require.NoError(t, err)
			assert.Equal(t, "true", obj.GetAnnotations()[consts.AnnotationKeyPauseReconcile])
		})
	}
}

func TestRoutePostResumeAgent(t *testing.T) {
	tests := []struct {
		name             string
		agentType        string
		agentName        string
		setupResources   []client.Object
		expectedCode     int
		expectedResponse AgentPauseResponse
	}{
		{
			name:      "resume kubernetes agent",
			agentType: consts.AgentKubernetes,
			agentName: "test-agent",
			setupResources: []client.Object{
				&agentv2.KubernetesAgent{
					ObjectMeta: metav1.ObjectMeta{
						Name: "test-agent",
						Annotations: map[string]string{
							consts.AnnotationKeyPauseReconcile: "true",
						},
					},
				},
			},
			expectedCode: http.StatusOK,
			expectedResponse: AgentPauseResponse{
				AgentType: consts.AgentKubernetes,
				AgentName: "test-agent",
				IsPaused:  false,
			},
		},
		{
			name:           "agent not found",
			agentType:      consts.AgentKubernetes,
			agentName:      "non-existent",
			setupResources: []client.Object{},
			expectedCode:   http.StatusBadRequest,
		},
		{
			name:           "missing agentType and agentName returns error",
			agentType:      "",
			agentName:      "",
			setupResources: []client.Object{},
			expectedCode:   http.StatusBadRequest,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := makeTestServer(tt.setupResources...)
			body := AgentActionRequest{AgentType: tt.agentType, AgentName: tt.agentName}
			req := makePostRequest("/agents/resume", body)

			w := httptest.NewRecorder()
			s.routePostResumeAgent(w, req)

			assert.Equal(t, tt.expectedCode, w.Code)
			if tt.expectedCode != http.StatusOK {
				return
			}

			rsp := parseResponse(w)
			assert.Equal(t, tt.expectedResponse, parseResponseData[AgentPauseResponse](rsp))

			obj, err := s.getAgentObject(context.Background(), tt.agentType, tt.agentName)
			require.NoError(t, err)
			assert.NotContains(t, obj.GetAnnotations(), consts.AnnotationKeyPauseReconcile)
		})
	}
}

func TestRoutePostRestartAgent(t *testing.T) {
	tests := []struct {
		name             string
		agentType        string
		agentName        string
		requestBody      AgentRestartRequest
		setupResources   []client.Object
		expectedCode     int
		expectedResponse AgentRestartResponse
		expectedAnnotKey string
	}{
		{
			name:      "restart vmagent",
			agentType: consts.AgentKubernetes,
			agentName: "test-agent",
			requestBody: AgentRestartRequest{
				Component: consts.ComponentVMAgent,
			},
			setupResources: []client.Object{
				&agentv2.KubernetesAgent{
					ObjectMeta: metav1.ObjectMeta{
						Name: "test-agent",
					},
				},
			},
			expectedCode: http.StatusOK,
			expectedResponse: AgentRestartResponse{
				AgentType: consts.AgentKubernetes,
				AgentName: "test-agent",
				Component: consts.ComponentVMAgent,
			},
			expectedAnnotKey: consts.AnnotationKeyO11yVMAgentRestart,
		},
		{
			name:      "restart vector",
			agentType: consts.AgentO11yK8s,
			agentName: "test-o11yk8s",
			requestBody: AgentRestartRequest{
				Component: consts.ComponentVector,
			},
			setupResources: []client.Object{
				&agentv2.O11yK8sAgent{
					ObjectMeta: metav1.ObjectMeta{
						Name: "test-o11yk8s",
					},
				},
			},
			expectedCode: http.StatusOK,
			expectedResponse: AgentRestartResponse{
				AgentType: consts.AgentO11yK8s,
				AgentName: "test-o11yk8s",
				Component: consts.ComponentVector,
			},
			expectedAnnotKey: consts.AnnotationKeyO11yVectorRestart,
		},
		{
			name:      "invalid component",
			agentType: consts.AgentKubernetes,
			agentName: "test-agent",
			requestBody: AgentRestartRequest{
				Component: "invalid",
			},
			setupResources: []client.Object{},
			expectedCode:   http.StatusBadRequest,
		},
		{
			name:      "agent not found",
			agentType: consts.AgentKubernetes,
			agentName: "non-existent",
			requestBody: AgentRestartRequest{
				Component: consts.ComponentVMAgent,
			},
			setupResources: []client.Object{},
			expectedCode:   http.StatusBadRequest,
		},
		{
			name:      "missing agentType and agentName returns error",
			agentType: "",
			agentName: "",
			requestBody: AgentRestartRequest{
				Component: consts.ComponentVMAgent,
			},
			setupResources: []client.Object{},
			expectedCode:   http.StatusBadRequest,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := makeTestServer(tt.setupResources...)
			tt.requestBody.AgentType = tt.agentType
			tt.requestBody.AgentName = tt.agentName
			req := makePostRequest("/agents/restart", tt.requestBody)

			w := httptest.NewRecorder()
			s.routePostRestartAgent(w, req)

			assert.Equal(t, tt.expectedCode, w.Code)
			if tt.expectedCode != http.StatusOK {
				return
			}

			rsp := parseResponse(w)
			assert.Equal(t, tt.expectedResponse, parseResponseData[AgentRestartResponse](rsp))

			obj, err := s.getAgentObject(context.Background(), tt.agentType, tt.agentName)
			require.NoError(t, err)
			assert.Equal(t, "true", obj.GetAnnotations()[tt.expectedAnnotKey])
		})
	}
}

func TestRoutePostDrainNode(t *testing.T) {
	tests := []struct {
		name             string
		agentType        string
		agentName        string
		requestBody      PodNodeDrainRequest
		setupResources   []client.Object
		expectedCode     int
		expectedResponse PodNodeDrainResponse
	}{
		{
			name:      "drain node for kubernetes agent",
			agentType: consts.AgentKubernetes,
			agentName: "test-agent",
			requestBody: PodNodeDrainRequest{
				Namespace: "monitoring",
				PodName:   "vector-pod-1",
			},
			setupResources: []client.Object{
				&agentv2.KubernetesAgent{
					ObjectMeta: metav1.ObjectMeta{
						Name: "test-agent",
					},
				},
			},
			expectedCode: http.StatusOK,
			expectedResponse: PodNodeDrainResponse{
				AgentType: consts.AgentKubernetes,
				AgentName: "test-agent",
				Namespace: "monitoring",
				PodName:   "vector-pod-1",
			},
		},
		{
			name:      "missing namespace",
			agentType: consts.AgentKubernetes,
			agentName: "test-agent",
			requestBody: PodNodeDrainRequest{
				Namespace: "",
				PodName:   "pod-1",
			},
			setupResources: []client.Object{},
			expectedCode:   http.StatusBadRequest,
		},
		{
			name:      "missing podName",
			agentType: consts.AgentKubernetes,
			agentName: "test-agent",
			requestBody: PodNodeDrainRequest{
				Namespace: "monitoring",
				PodName:   "",
			},
			setupResources: []client.Object{},
			expectedCode:   http.StatusBadRequest,
		},
		{
			name:      "agent not found",
			agentType: consts.AgentKubernetes,
			agentName: "non-existent",
			requestBody: PodNodeDrainRequest{
				Namespace: "monitoring",
				PodName:   "pod-1",
			},
			setupResources: []client.Object{},
			expectedCode:   http.StatusBadRequest,
		},
		{
			name:      "missing agentType and agentName returns error",
			agentType: "",
			agentName: "",
			requestBody: PodNodeDrainRequest{
				Namespace: "monitoring",
				PodName:   "pod-1",
			},
			setupResources: []client.Object{},
			expectedCode:   http.StatusBadRequest,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := makeTestServer(tt.setupResources...)
			tt.requestBody.AgentType = tt.agentType
			tt.requestBody.AgentName = tt.agentName
			req := makePostRequest("/agents/drainnode", tt.requestBody)

			w := httptest.NewRecorder()
			s.routePostDrainNode(w, req)

			assert.Equal(t, tt.expectedCode, w.Code)
			if tt.expectedCode != http.StatusOK {
				return
			}

			rsp := parseResponse(w)
			assert.Equal(t, tt.expectedResponse, parseResponseData[PodNodeDrainResponse](rsp))

			obj, err := s.getAgentObject(context.Background(), tt.agentType, tt.agentName)
			require.NoError(t, err)

			drainValue := tt.requestBody.Namespace + "/" + tt.requestBody.PodName
			assert.Equal(t, drainValue, obj.GetAnnotations()[consts.AnnotationKeyO11yDrainNode])
		})
	}
}

func TestRouteListAgents(t *testing.T) {
	setupResources := []client.Object{
		&agentv2.KubernetesAgent{
			ObjectMeta: metav1.ObjectMeta{
				Name: "k8s-agent",
				Annotations: map[string]string{
					consts.AnnotationKeyO11yID: "o11y-a",
				},
			},
			Spec: agentv2.KubernetesAgentSpec{
				KubernetesClusterName: "cluster-1",
				O11y: agentv2.O11yConfigV2{
					TenantID:      "tenant-1",
					ApplicationID: "app-1",
				},
			},
			Status: agentv2.KubernetesAgentStatus{Phase: "Normal"},
		},
		&agentv2.O11yK8sAgent{
			ObjectMeta: metav1.ObjectMeta{
				Name: "o11y-k8s-agent",
				Annotations: map[string]string{
					consts.AnnotationKeyO11yID: "o11y-b",
				},
			},
			Spec: agentv2.O11yK8sAgentSpec{
				KubernetesClusterName: "cluster-1",
				O11y: agentv2.O11yConfig{
					TenantID:      "tenant-1",
					ProjectID:     "project-1",
					ApplicationID: "app-1",
				},
			},
			Status: agentv2.O11yK8sAgentStatus{Phase: "Normal"},
		},
		&agentv2.O11yTiDBAgent{
			ObjectMeta: metav1.ObjectMeta{
				Name: "o11y-tidb-agent",
				Annotations: map[string]string{
					consts.AnnotationKeyO11yID: "o11y-c",
				},
			},
			Spec: agentv2.O11yTiDBAgentSpec{
				KubernetesClusterName: "cluster-2",
				O11y: agentv2.O11yConfig{
					TenantID:      "tenant-2",
					ProjectID:     "project-2",
					ApplicationID: "app-2",
				},
			},
			Status: agentv2.O11yTiDBAgentStatus{Phase: "Normal"},
		},
		&agentv2.SharedPoolAgent{
			ObjectMeta: metav1.ObjectMeta{
				Name: "sharedpool-agent",
				Annotations: map[string]string{
					consts.AnnotationKeyO11yID: "o11y-d",
				},
			},
			Spec: agentv2.SharedPoolAgentSpec{
				KubernetesClusterName: "cluster-1",
				SharedPoolID:          "spool-1",
				O11y: agentv2.O11yConfigV2{
					TenantID:      "tenant-1",
					ApplicationID: "app-1",
				},
			},
			Status: agentv2.SharedPoolAgentStatus{Phase: "Normal"},
		},
		&agentv2.TiDBGroupAgent{
			ObjectMeta: metav1.ObjectMeta{
				Name: "tidbgroup-agent",
				Annotations: map[string]string{
					consts.AnnotationKeyO11yID: "o11y-e",
				},
			},
			Spec: agentv2.TiDBGroupAgentSpec{
				KubernetesClusterName: "cluster-2",
				SharedPoolID:          "spool-2",
				O11y: agentv2.O11yConfigV2{
					TenantID:      "tenant-2",
					ApplicationID: "app-2",
				},
			},
			Status: agentv2.TiDBGroupAgentStatus{Phase: "Normal"},
		},
	}

	tests := []struct {
		name             string
		query            string
		expectedCode     int
		expectedResponse []AgentInfoResponse
	}{
		{
			name:         "list all agents",
			query:        "",
			expectedCode: http.StatusOK,
			expectedResponse: []AgentInfoResponse{
				{
					Name:                  "k8s-agent",
					AgentType:             consts.AgentKubernetes,
					KubernetesClusterName: "cluster-1",
					TenantID:              "tenant-1",
					ApplicationID:         "app-1",
					O11yID:                "o11y-a",
					Phase:                 "Normal",
				},
				{
					Name:                  "o11y-k8s-agent",
					AgentType:             consts.AgentO11yK8s,
					KubernetesClusterName: "cluster-1",
					TenantID:              "tenant-1",
					ApplicationID:         "app-1",
					ProjectID:             "project-1",
					O11yID:                "o11y-b",
					Phase:                 "Normal",
				},
				{
					Name:                  "o11y-tidb-agent",
					AgentType:             consts.AgentO11yTidb,
					KubernetesClusterName: "cluster-2",
					TenantID:              "tenant-2",
					ApplicationID:         "app-2",
					ProjectID:             "project-2",
					O11yID:                "o11y-c",
					Phase:                 "Normal",
				},
				{
					Name:                  "sharedpool-agent",
					AgentType:             consts.AgentSharedPool,
					KubernetesClusterName: "cluster-1",
					TenantID:              "tenant-1",
					ApplicationID:         "app-1",
					SharedPoolID:          "spool-1",
					O11yID:                "o11y-d",
					Phase:                 "Normal",
				},
				{
					Name:                  "tidbgroup-agent",
					AgentType:             consts.AgentTidbGroup,
					KubernetesClusterName: "cluster-2",
					TenantID:              "tenant-2",
					ApplicationID:         "app-2",
					SharedPoolID:          "spool-2",
					O11yID:                "o11y-e",
					Phase:                 "Normal",
				},
			},
		},
		{
			name:         "list agents by cluster",
			query:        "?kubernetesClusterName=cluster-1",
			expectedCode: http.StatusOK,
			expectedResponse: []AgentInfoResponse{
				{
					Name:                  "k8s-agent",
					AgentType:             consts.AgentKubernetes,
					KubernetesClusterName: "cluster-1",
					TenantID:              "tenant-1",
					ApplicationID:         "app-1",
					O11yID:                "o11y-a",
					Phase:                 "Normal",
				},
				{
					Name:                  "o11y-k8s-agent",
					AgentType:             consts.AgentO11yK8s,
					KubernetesClusterName: "cluster-1",
					TenantID:              "tenant-1",
					ApplicationID:         "app-1",
					ProjectID:             "project-1",
					O11yID:                "o11y-b",
					Phase:                 "Normal",
				},
				{
					Name:                  "sharedpool-agent",
					AgentType:             consts.AgentSharedPool,
					KubernetesClusterName: "cluster-1",
					TenantID:              "tenant-1",
					ApplicationID:         "app-1",
					SharedPoolID:          "spool-1",
					O11yID:                "o11y-d",
					Phase:                 "Normal",
				},
			},
		},
		{
			name:         "list agents by type",
			query:        "?agentType=o11yk8s",
			expectedCode: http.StatusOK,
			expectedResponse: []AgentInfoResponse{
				{
					Name:                  "o11y-k8s-agent",
					AgentType:             consts.AgentO11yK8s,
					KubernetesClusterName: "cluster-1",
					TenantID:              "tenant-1",
					ApplicationID:         "app-1",
					ProjectID:             "project-1",
					O11yID:                "o11y-b",
					Phase:                 "Normal",
				},
			},
		},
		{
			name:         "list agents by type and tenant",
			query:        "?agentType=o11yk8s&tenantId=tenant-1",
			expectedCode: http.StatusOK,
			expectedResponse: []AgentInfoResponse{
				{
					Name:                  "o11y-k8s-agent",
					AgentType:             consts.AgentO11yK8s,
					KubernetesClusterName: "cluster-1",
					TenantID:              "tenant-1",
					ApplicationID:         "app-1",
					ProjectID:             "project-1",
					O11yID:                "o11y-b",
					Phase:                 "Normal",
				},
			},
		},
		{
			name:         "get single agent by type and name",
			query:        "?agentType=kubernetes&agentName=k8s-agent",
			expectedCode: http.StatusOK,
			expectedResponse: []AgentInfoResponse{
				{
					Name:                  "k8s-agent",
					AgentType:             consts.AgentKubernetes,
					KubernetesClusterName: "cluster-1",
					TenantID:              "tenant-1",
					ApplicationID:         "app-1",
					O11yID:                "o11y-a",
					Phase:                 "Normal",
				},
			},
		},
		{
			name:             "agent not found returns empty list",
			query:            "?agentType=kubernetes&agentName=missing-agent",
			expectedCode:     http.StatusOK,
			expectedResponse: []AgentInfoResponse{},
		},
		{
			name:         "invalid agent type",
			query:        "?agentType=invalid",
			expectedCode: http.StatusBadRequest,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := makeTestServer(setupResources...)
			path := fmt.Sprintf("/agents%s", tt.query)
			req := makeGetRequest(path)

			w := httptest.NewRecorder()
			s.routeListAgents(w, req)

			assert.Equal(t, tt.expectedCode, w.Code)
			if tt.expectedCode != http.StatusOK {
				return
			}

			rsp := parseResponse(w)
			assert.ElementsMatch(t, tt.expectedResponse, parseResponseData[[]AgentInfoResponse](rsp))
		})
	}
}

func TestRoutePostAgentStorage(t *testing.T) {
	tests := []struct {
		name             string
		agentType        string
		agentName        string
		requestBody      AgentStorageRequest
		setupResources   []client.Object
		expectedCode     int
		expectedResponse AgentStorageResponse
	}{
		{
			name:      "set vmagent storage for kubernetes agent",
			agentType: consts.AgentKubernetes,
			agentName: "test-agent",
			requestBody: AgentStorageRequest{
				Component: consts.ComponentVMAgent,
				Size:      "50Gi",
			},
			setupResources: []client.Object{
				&agentv2.KubernetesAgent{
					ObjectMeta: metav1.ObjectMeta{
						Name: "test-agent",
					},
				},
			},
			expectedCode: http.StatusOK,
			expectedResponse: AgentStorageResponse{
				AgentType: consts.AgentKubernetes,
				AgentName: "test-agent",
				Component: consts.ComponentVMAgent,
				Size:      "50Gi",
			},
		},
		{
			name:      "missing component",
			agentType: consts.AgentKubernetes,
			agentName: "test-agent",
			requestBody: AgentStorageRequest{
				Component: "",
				Size:      "50Gi",
			},
			setupResources: []client.Object{},
			expectedCode:   http.StatusBadRequest,
		},
		{
			name:      "missing size",
			agentType: consts.AgentKubernetes,
			agentName: "test-agent",
			requestBody: AgentStorageRequest{
				Component: consts.ComponentVMAgent,
				Size:      "",
			},
			setupResources: []client.Object{},
			expectedCode:   http.StatusBadRequest,
		},
		{
			name:      "unsupported component",
			agentType: consts.AgentKubernetes,
			agentName: "test-agent",
			requestBody: AgentStorageRequest{
				Component: "invalid",
				Size:      "50Gi",
			},
			setupResources: []client.Object{},
			expectedCode:   http.StatusBadRequest,
		},
		{
			name:      "agent not found",
			agentType: consts.AgentKubernetes,
			agentName: "non-existent",
			requestBody: AgentStorageRequest{
				Component: consts.ComponentVMAgent,
				Size:      "50Gi",
			},
			setupResources: []client.Object{},
			expectedCode:   http.StatusBadRequest,
		},
		{
			name:      "missing agentType and agentName returns error",
			agentType: "",
			agentName: "",
			requestBody: AgentStorageRequest{
				Component: consts.ComponentVMAgent,
				Size:      "50Gi",
			},
			setupResources: []client.Object{},
			expectedCode:   http.StatusBadRequest,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := makeTestServer(tt.setupResources...)
			tt.requestBody.AgentType = tt.agentType
			tt.requestBody.AgentName = tt.agentName
			req := makePostRequest("/agents/storage", tt.requestBody)

			w := httptest.NewRecorder()
			s.routePostAgentStorage(w, req)

			assert.Equal(t, tt.expectedCode, w.Code)
			if tt.expectedCode != http.StatusOK {
				return
			}

			rsp := parseResponse(w)
			assert.Equal(t, tt.expectedResponse, parseResponseData[AgentStorageResponse](rsp))

			obj, err := s.getAgentObject(context.Background(), tt.agentType, tt.agentName)
			require.NoError(t, err)
			assert.Equal(t, tt.requestBody.Size, obj.GetAnnotations()[consts.AnnotationKeyO11yVMAgentStorageSize])
		})
	}
}

func TestGetAgentObject(t *testing.T) {
	tests := []struct {
		name      string
		agent     client.Object
		agentType string
		agentName string
		wantErr   bool
	}{
		{
			name:      "finds O11yTiDBAgent",
			agent:     &agentv2.O11yTiDBAgent{ObjectMeta: metav1.ObjectMeta{Name: "tidb-1", Namespace: "default"}},
			agentType: consts.AgentO11yTidb,
			agentName: "tidb-1",
		},
		{
			name:      "finds SharedPoolAgent",
			agent:     &agentv2.SharedPoolAgent{ObjectMeta: metav1.ObjectMeta{Name: "sp-1", Namespace: "default"}},
			agentType: consts.AgentSharedPool,
			agentName: "sp-1",
		},
		{
			name:      "finds TiDBGroupAgent",
			agent:     &agentv2.TiDBGroupAgent{ObjectMeta: metav1.ObjectMeta{Name: "tg-1", Namespace: "default"}},
			agentType: consts.AgentTidbGroup,
			agentName: "tg-1",
		},
		{
			name:      "returns error when agent not found",
			agentType: consts.AgentSharedPool,
			agentName: "missing",
			wantErr:   true,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			var s *Server
			if tt.agent != nil {
				s = makeTestServer(tt.agent)
			} else {
				s = makeTestServer()
			}
			obj, err := s.getAgentObject(context.Background(), tt.agentType, tt.agentName)
			if tt.wantErr {
				assert.Error(t, err)
			} else {
				require.NoError(t, err)
				assert.Equal(t, tt.agentName, obj.GetName())
			}
		})
	}
}