package mcp

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strings"
	"testing"

	"github.com/gin-gonic/gin"

	"github.com/tidbcloud/AgentTeam/desktop/internal/token"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
)

var _ *types.MCPToolResult = (*toolResult)(nil)
var _ types.MCPToolContent = toolContent{}

type receiptAckHandler struct {
	receipt  string
	owner    types.MCPToolResultDeliveryOwner
	agentID  string
	receipts []string
	outcome  types.MCPToolDeliveryOutcome
	calls    int
}

func (h *receiptAckHandler) Tools() []types.MCPToolDef       { return nil }
func (h *receiptAckHandler) SkillDoc(types.ToolNamer) string { return "" }
func (h *receiptAckHandler) SkillDocSplit(types.ToolNamer) types.SkillDocParts {
	return types.SkillDocParts{}
}
func (h *receiptAckHandler) CallTool(context.Context, string, string, json.RawMessage) *types.MCPToolResult {
	result := types.NewToolText("claimed")
	result.Meta = map[string]any{types.MCPRuntimeDeliveryReceiptsMetaKey: []string{h.receipt}}
	result.DeliveryOwner = h.owner
	return result
}
func (h *receiptAckHandler) AcknowledgeMCPToolResultDelivery(agentID string, _ string, receipts []string, outcome types.MCPToolDeliveryOutcome) {
	h.calls++
	h.agentID = agentID
	h.receipts = append([]string(nil), receipts...)
	h.outcome = outcome
}

type exactReceiptOwner struct {
	receipts [][]string
	outcomes []types.MCPToolDeliveryOutcome
}

func (o *exactReceiptOwner) AcknowledgeMCPToolResultDelivery(receipts []string, outcome types.MCPToolDeliveryOutcome) {
	o.receipts = append(o.receipts, append([]string(nil), receipts...))
	o.outcomes = append(o.outcomes, outcome)
}

type shortHTTPResponseWriter struct {
	header http.Header
	status int
}

func (w *shortHTTPResponseWriter) Header() http.Header {
	if w.header == nil {
		w.header = make(http.Header)
	}
	return w.header
}
func (w *shortHTTPResponseWriter) WriteHeader(status int) { w.status = status }
func (w *shortHTTPResponseWriter) Write(p []byte) (int, error) {
	if len(p) == 0 {
		return 0, nil
	}
	return len(p) - 1, nil
}

func rawMsg(s string) *json.RawMessage {
	r := json.RawMessage(s)
	return &r
}

func TestToolText(t *testing.T) {
	result := toolText("hello")
	if result.IsError {
		t.Error("toolText should not be error")
	}
	if len(result.Content) != 1 || result.Content[0].Text != "hello" {
		t.Errorf("toolText content = %v", result.Content)
	}
}

func TestToolError(t *testing.T) {
	result := toolError("something failed")
	if !result.IsError {
		t.Error("toolError should be error")
	}
	if result.Content[0].Text != "Error: something failed" {
		t.Errorf("toolError text = %q", result.Content[0].Text)
	}
}

func TestMCPToolDeliveryReceiptCommitsOnlyAfterCompleteHTTPWrite(t *testing.T) {
	receipt := types.NewRuntimeDeliveryReceipt("0123456789abcdef0123456789abcdef")
	body := `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"test_read","arguments":{}}}`

	t.Run("direct complete write is delivered", func(t *testing.T) {
		handler := &receiptAckHandler{receipt: receipt}
		router := testBatchRouter(handler)
		recorder := postBatch(router, body)
		if recorder.Code != http.StatusOK {
			t.Fatalf("status = %d: %s", recorder.Code, recorder.Body.String())
		}
		if handler.outcome != types.MCPToolDeliveryDelivered {
			t.Fatalf("outcome = %q, want delivered", handler.outcome)
		}
		if handler.agentID != "agent-1" || len(handler.receipts) != 1 || handler.receipts[0] != receipt {
			t.Fatalf("ack = agent %q receipts %#v", handler.agentID, handler.receipts)
		}
	})

	t.Run("codex first hop is deferred not delivered", func(t *testing.T) {
		handler := &receiptAckHandler{receipt: receipt}
		router := testBatchRouter(handler)
		req := httptest.NewRequest(http.MethodPost, "/mcp/server/test-server/agent-1", strings.NewReader(body))
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set(types.MCPDeferredToolDeliveryHeader, "true")
		recorder := httptest.NewRecorder()
		router.ServeHTTP(recorder, req)
		if handler.outcome != types.MCPToolDeliveryDeferred {
			t.Fatalf("outcome = %q, want deferred", handler.outcome)
		}
	})

	t.Run("exact owner bypasses mutable handler lookup", func(t *testing.T) {
		owner := &exactReceiptOwner{}
		handler := &receiptAckHandler{receipt: receipt, owner: owner}
		router := testBatchRouter(handler)
		recorder := postBatch(router, body)
		if recorder.Code != http.StatusOK {
			t.Fatalf("status = %d: %s", recorder.Code, recorder.Body.String())
		}
		if handler.calls != 0 {
			t.Fatalf("mutable handler fallback called %d times", handler.calls)
		}
		if len(owner.outcomes) != 1 || owner.outcomes[0] != types.MCPToolDeliveryDelivered {
			t.Fatalf("exact owner outcomes = %#v, want delivered", owner.outcomes)
		}
	})

	t.Run("direct short write retains exact owner without replay acknowledgement", func(t *testing.T) {
		owner := &exactReceiptOwner{}
		handler := &receiptAckHandler{receipt: receipt, owner: owner}
		router := testBatchRouter(handler)
		req := httptest.NewRequest(http.MethodPost, "/mcp/server/test-server/agent-1", strings.NewReader(body))
		req.Header.Set("Content-Type", "application/json")
		writer := &shortHTTPResponseWriter{}
		router.ServeHTTP(writer, req)
		if handler.calls != 0 || len(owner.outcomes) != 0 {
			t.Fatalf("ambiguous direct write acknowledged handler=%d owner=%#v", handler.calls, owner.outcomes)
		}
	})

	t.Run("deferred short write preserves transfer until late final hop", func(t *testing.T) {
		owner := &exactReceiptOwner{}
		handler := &receiptAckHandler{receipt: receipt, owner: owner}
		router := testBatchRouter(handler)
		req := httptest.NewRequest(http.MethodPost, "/mcp/server/test-server/agent-1", strings.NewReader(body))
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set(types.MCPDeferredToolDeliveryHeader, "true")
		writer := &shortHTTPResponseWriter{}
		router.ServeHTTP(writer, req)
		if handler.calls != 0 {
			t.Fatalf("mutable handler fallback called %d times", handler.calls)
		}
		if len(owner.outcomes) != 1 || owner.outcomes[0] != types.MCPToolDeliveryDeferred {
			t.Fatalf("exact owner outcomes = %#v, want deferred", owner.outcomes)
		}

		// Model the Codex child's later successful JSON-RPC response write. The
		// deferred ambiguous HTTP first hop must leave the exact owner usable for
		// this terminal event rather than releasing it into reconciliation.
		owner.AcknowledgeMCPToolResultDelivery([]string{receipt}, types.MCPToolDeliveryDelivered)
		if len(owner.outcomes) != 2 || owner.outcomes[1] != types.MCPToolDeliveryDelivered {
			t.Fatalf("late final-hop outcomes = %#v", owner.outcomes)
		}
	})
}

func TestMCPProtocolPrefersExactDeliveryOwnerOverMuxFallback(t *testing.T) {
	receipt := types.NewRuntimeDeliveryReceipt("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa")
	owner := &exactReceiptOwner{}
	message := &toolMuxTestHandler{}
	mux, err := newToolMux(toolChild{kind: toolChildMessage, inventory: message})
	if err != nil {
		t.Fatal(err)
	}
	acknowledgeMCPToolResultDelivery(mux, "agent-a", "runtime-a", []*rpcResponse{{
		deliveryReceipts: []string{receipt},
		deliveryOwner:    owner,
	}}, types.MCPToolDeliveryDelivered)

	if message.ackCalls != 0 {
		t.Fatalf("fallback called %d times", message.ackCalls)
	}
	if len(owner.receipts) != 1 || len(owner.receipts[0]) != 1 || owner.receipts[0][0] != receipt {
		t.Fatalf("exact owner receipts = %#v", owner.receipts)
	}
}

func TestMCPProtocolDeduplicatesFallbackReceipts(t *testing.T) {
	first := types.NewRuntimeDeliveryReceipt("bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb")
	second := types.NewRuntimeDeliveryReceipt("cccccccccccccccccccccccccccccccc")
	message := &toolMuxTestHandler{}
	mux, err := newToolMux(toolChild{kind: toolChildMessage, inventory: message})
	if err != nil {
		t.Fatal(err)
	}
	acknowledgeMCPToolResultDelivery(mux, "agent-a", " runtime-a ", []*rpcResponse{
		{deliveryReceipts: []string{first, second}},
		{deliveryReceipts: []string{first}},
	}, types.MCPToolDeliveryUncertain)

	if message.ackCalls != 1 {
		t.Fatalf("fallback calls = %d, want 1", message.ackCalls)
	}
	if got := strings.Join(message.ackReceipts, ","); got != first+","+second {
		t.Fatalf("fallback receipts = %q", got)
	}
	if message.ackAgentID != "agent-a" || message.ackRuntimeSessionID != "runtime-a" {
		t.Fatalf("fallback identity = %q/%q", message.ackAgentID, message.ackRuntimeSessionID)
	}
}

func TestParseErrorResponse(t *testing.T) {
	resp := parseErrorResponse("parse error")
	if resp.JSONRPC != "2.0" {
		t.Errorf("JSONRPC = %q, want %q", resp.JSONRPC, "2.0")
	}
	if resp.Error == nil {
		t.Fatal("Error should not be nil")
	}
	if resp.Error.Code != -32700 {
		t.Errorf("Error.Code = %d, want -32700", resp.Error.Code)
	}
	if resp.Error.Message != "parse error" {
		t.Errorf("Error.Message = %q, want %q", resp.Error.Message, "parse error")
	}
	if resp.ID == nil {
		t.Error("ID should not be nil")
	}
}

func TestMCPSessionID(t *testing.T) {
	if got := mcpSessionID("agent-1"); got != "agentteam-agent-1" {
		t.Errorf("mcpSessionID() = %q, want %q", got, "agentteam-agent-1")
	}
	if got := mcpSessionID("  agent-2  "); got != "agentteam-agent-2" {
		t.Errorf("mcpSessionID() with whitespace = %q, want %q", got, "agentteam-agent-2")
	}
}

func TestDispatchRPC_Initialize(t *testing.T) {
	handler := &mockToolsListHandler{}

	body, _ := json.Marshal(rpcRequest{
		JSONRPC: "2.0",
		ID:      rawMsg(`"1"`),
		Method:  "initialize",
		Params:  json.RawMessage(`{}`),
	})

	resp, ok := dispatchRPC(context.TODO(), "agent-1", handler, body)
	if !ok {
		t.Fatal("dispatchRPC should return ok=true for initialize")
	}
	if resp.JSONRPC != "2.0" {
		t.Errorf("JSONRPC = %q", resp.JSONRPC)
	}
	if resp.Error != nil {
		t.Errorf("unexpected error: %v", resp.Error)
	}
}

func TestDispatchRPC_Ping(t *testing.T) {
	body, _ := json.Marshal(rpcRequest{
		JSONRPC: "2.0",
		ID:      rawMsg(`"2"`),
		Method:  "ping",
	})

	resp, ok := dispatchRPC(context.TODO(), "agent-1", nil, body)
	if !ok {
		t.Fatal("dispatchRPC should return ok=true for ping")
	}
	if resp.Error != nil {
		t.Errorf("unexpected error: %v", resp.Error)
	}
}

func TestDispatchRPC_MethodNotFound(t *testing.T) {
	body, _ := json.Marshal(rpcRequest{
		JSONRPC: "2.0",
		ID:      rawMsg(`"3"`),
		Method:  "nonexistent",
	})

	resp, ok := dispatchRPC(context.TODO(), "agent-1", nil, body)
	if !ok {
		t.Fatal("dispatchRPC should return ok=true for method not found")
	}
	if resp.Error == nil {
		t.Fatal("expected error for method not found")
	}
	if resp.Error.Code != -32601 {
		t.Errorf("Error.Code = %d, want -32601", resp.Error.Code)
	}
}

func TestDispatchRPC_ParseError(t *testing.T) {
	resp, ok := dispatchRPC(context.TODO(), "agent-1", nil, []byte("invalid json"))
	if !ok {
		t.Fatal("dispatchRPC should return ok=true for parse error")
	}
	if resp.Error == nil {
		t.Fatal("expected error for parse error")
	}
	if resp.Error.Code != -32700 {
		t.Errorf("Error.Code = %d, want -32700", resp.Error.Code)
	}
}

func TestDispatchRPC_NoID(t *testing.T) {
	body, _ := json.Marshal(rpcRequest{
		JSONRPC: "2.0",
		Method:  "ping",
	})

	_, ok := dispatchRPC(context.TODO(), "agent-1", nil, body)
	if ok {
		t.Error("dispatchRPC should return ok=false for notification (no ID)")
	}
}

func TestWriteMCPStatus_SetsJSONContentType(t *testing.T) {
	gin.SetMode(gin.TestMode)
	recorder := httptest.NewRecorder()
	ctx, _ := gin.CreateTestContext(recorder)

	writeMCPStatus(ctx, http.StatusAccepted)

	if recorder.Code != http.StatusAccepted {
		t.Fatalf("status = %d, want %d", recorder.Code, http.StatusAccepted)
	}
	if recorder.Header().Get("Content-Type") != "application/json" {
		t.Fatalf("Content-Type = %q, want application/json", recorder.Header().Get("Content-Type"))
	}
}

func TestMCPProtectedResourcePath(t *testing.T) {
	if got := mcpProtectedResourcePath("/mcp/server/ch/agent"); got != "/mcp/server/ch/agent/.well-known/oauth-protected-resource" {
		t.Fatalf("protected resource path = %q", got)
	}
	if got := mcpProtectedResourcePath(""); got != oauthProtectedResourcePath {
		t.Fatalf("empty protected resource path = %q", got)
	}
}

func TestSetProtectedResourceChallenge(t *testing.T) {
	gin.SetMode(gin.TestMode)
	recorder := httptest.NewRecorder()
	ctx, _ := gin.CreateTestContext(recorder)

	setProtectedResourceChallenge(ctx, "/mcp/server/ch/agent")

	want := `Bearer resource_metadata="/mcp/server/ch/agent/.well-known/oauth-protected-resource"`
	if got := recorder.Header().Get("WWW-Authenticate"); got != want {
		t.Fatalf("WWW-Authenticate = %q, want %q", got, want)
	}
}

func TestAccessControlAllowsActiveServerSessionToken(t *testing.T) {
	manager := newMCPTestTokenManager(t)
	router := testAccessControlRouter(manager)
	sessionKey := "sk1:session"
	tokenValue, err := manager.GenerateAgentTokenForServer(sessionKey, "server-a")
	if err != nil {
		t.Fatalf("GenerateAgentTokenForServer() error = %v", err)
	}

	recorder := postAuthenticatedPing(router, "/mcp/server/server-a/"+sessionKey, tokenValue)

	if recorder.Code != http.StatusOK {
		t.Fatalf("status = %d, want %d: %s", recorder.Code, http.StatusOK, recorder.Body.String())
	}
}

func TestLeaderAccessControlAllowsActiveSessionToken(t *testing.T) {
	manager := newMCPTestTokenManager(t)
	router := testleaderAccessControlRouter(manager)
	sessionKey := "sk1:session"
	tokenValue, err := manager.GenerateAgentTokenForServer(sessionKey, "server-a")
	if err != nil {
		t.Fatalf("GenerateAgentTokenForServer() error = %v", err)
	}

	recorder := postAuthenticatedPing(router, "/mcp/agentteam/"+sessionKey, tokenValue)

	if recorder.Code != http.StatusOK {
		t.Fatalf("status = %d, want %d: %s", recorder.Code, http.StatusOK, recorder.Body.String())
	}
}

func TestAccessControlRejectsMalformedBearerAuthorization(t *testing.T) {
	manager := newMCPTestTokenManager(t)
	gin.SetMode(gin.TestMode)
	router := gin.New()
	group := router.Group("/api/internal/mcp/servers/:serverId/agents/:agentId", accessControlWithManagerForTest(manager))
	RegisterHandlerRoutes(group, func(string) (types.MCPHandler, bool) {
		return &mockToolsListHandler{}, true
	})
	tokenValue, err := manager.GenerateAgentTokenForServer("agent-1", "server-a")
	if err != nil {
		t.Fatalf("GenerateAgentTokenForServer() error = %v", err)
	}

	for _, authorization := range []string{
		"Bearer" + tokenValue,
		"Bearer " + tokenValue + " ",
		" Bearer " + tokenValue,
		"bearer " + tokenValue,
	} {
		req := httptest.NewRequest(
			http.MethodPost,
			"/api/internal/mcp/servers/server-a/agents/agent-1",
			strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"ping"}`),
		)
		req.Header.Set("Authorization", authorization)
		recorder := httptest.NewRecorder()
		router.ServeHTTP(recorder, req)
		if recorder.Code != http.StatusUnauthorized {
			t.Fatalf("authorization %q status = %d, want %d", authorization, recorder.Code, http.StatusUnauthorized)
		}
	}
}

func TestLeaderAccessControlRejectsStoppedSessionToken(t *testing.T) {
	manager := newMCPTestTokenManager(t)
	router := testleaderAccessControlRouter(manager)
	sessionKey := "sk1:session"
	tokenValue, err := manager.GenerateAgentTokenForServer(sessionKey, "server-a")
	if err != nil {
		t.Fatalf("GenerateAgentTokenForServer() error = %v", err)
	}
	manager.RevokeAgentToken(tokenValue)

	recorder := postAuthenticatedPing(router, "/mcp/agentteam/"+sessionKey, tokenValue)

	if recorder.Code != http.StatusUnauthorized {
		t.Fatalf("status = %d, want %d: %s", recorder.Code, http.StatusUnauthorized, recorder.Body.String())
	}
}

func TestLeaderAccessControlRejectsAgentMismatch(t *testing.T) {
	manager := newMCPTestTokenManager(t)
	router := testleaderAccessControlRouter(manager)
	tokenValue, err := manager.GenerateAgentTokenForServer("agent-a", "server-a")
	if err != nil {
		t.Fatalf("GenerateAgentTokenForServer() error = %v", err)
	}

	recorder := postAuthenticatedPing(router, "/mcp/agentteam/agent-b", tokenValue)

	if recorder.Code != http.StatusForbidden {
		t.Fatalf("status = %d, want %d: %s", recorder.Code, http.StatusForbidden, recorder.Body.String())
	}
}

func TestAccessControlRejectsStoppedSessionToken(t *testing.T) {
	manager := newMCPTestTokenManager(t)
	router := testAccessControlRouter(manager)
	sessionKey := "sk1:session"
	tokenValue, err := manager.GenerateAgentTokenForServer(sessionKey, "server-a")
	if err != nil {
		t.Fatalf("GenerateAgentTokenForServer() error = %v", err)
	}
	manager.RevokeAgentToken(tokenValue)

	recorder := postAuthenticatedPing(router, "/mcp/server/server-a/"+sessionKey, tokenValue)

	if recorder.Code != http.StatusUnauthorized {
		t.Fatalf("status = %d, want %d: %s", recorder.Code, http.StatusUnauthorized, recorder.Body.String())
	}
}

func TestAccessControlRejectsServerRouteMismatch(t *testing.T) {
	manager := newMCPTestTokenManager(t)
	router := testAccessControlRouter(manager)
	sessionKey := "sk1:session"
	tokenValue, err := manager.GenerateAgentTokenForServer(sessionKey, "server-a")
	if err != nil {
		t.Fatalf("GenerateAgentTokenForServer() error = %v", err)
	}

	recorder := postAuthenticatedPing(router, "/mcp/server/server-b/"+sessionKey, tokenValue)

	if recorder.Code != http.StatusForbidden {
		t.Fatalf("status = %d, want %d: %s", recorder.Code, http.StatusForbidden, recorder.Body.String())
	}
}

func TestRegisterDiscoveryRoutes_ProtectedResourceInsertionRoute(t *testing.T) {
	gin.SetMode(gin.TestMode)
	router := gin.New()
	RegisterDiscoveryRoutes(router)

	req, _ := http.NewRequest("GET", "/.well-known/oauth-protected-resource/mcp/server/test-server/agent-1", nil)
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status = %d, want 200: %s", w.Code, w.Body.String())
	}
	if got := w.Header().Get("Content-Type"); !strings.HasPrefix(got, "application/json") {
		t.Fatalf("Content-Type = %q, want application/json", got)
	}

	var body map[string]any
	if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil {
		t.Fatalf("unmarshal response: %v", err)
	}
	if body["resource"] != "/mcp/server/test-server/agent-1" {
		t.Fatalf("resource = %v, want /mcp/server/test-server/agent-1", body["resource"])
	}
}

func TestRegisterDiscoveryRoutes_ProtectedResourceAppendedRoute(t *testing.T) {
	gin.SetMode(gin.TestMode)
	router := gin.New()
	RegisterDiscoveryRoutes(router)

	req, _ := http.NewRequest("GET", "/mcp/server/test-server/agent-1/.well-known/oauth-protected-resource", nil)
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)

	if w.Code != http.StatusOK {
		t.Fatalf("status = %d, want 200: %s", w.Code, w.Body.String())
	}
	if got := w.Header().Get("Content-Type"); !strings.HasPrefix(got, "application/json") {
		t.Fatalf("Content-Type = %q, want application/json", got)
	}

	var body map[string]any
	if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil {
		t.Fatalf("unmarshal response: %v", err)
	}
	if body["resource"] != "/mcp/server/test-server/agent-1" {
		t.Fatalf("resource = %v, want /mcp/server/test-server/agent-1", body["resource"])
	}
}

func TestRegisterDiscoveryRoutes_AgentTeamWellKnownSuffixes(t *testing.T) {
	gin.SetMode(gin.TestMode)
	router := gin.New()
	RegisterDiscoveryRoutes(router)

	tests := []struct {
		name         string
		path         string
		wantStatus   int
		wantResource string
	}{
		{
			name:         "protected resource",
			path:         "/api/internal/mcp/servers/agentteam-main/agents/agent-1/.well-known/oauth-protected-resource",
			wantStatus:   http.StatusOK,
			wantResource: "/api/internal/mcp/servers/agentteam-main/agents/agent-1",
		},
		{
			name:       "oauth authorization server",
			path:       "/api/internal/mcp/servers/agentteam-main/agents/agent-1/.well-known/oauth-authorization-server",
			wantStatus: http.StatusNotFound,
		},
		{
			name:       "openid configuration",
			path:       "/api/internal/mcp/servers/agentteam-main/agents/agent-1/.well-known/openid-configuration",
			wantStatus: http.StatusNotFound,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			req, _ := http.NewRequest("GET", tt.path, nil)
			w := httptest.NewRecorder()
			router.ServeHTTP(w, req)

			if w.Code != tt.wantStatus {
				t.Fatalf("status = %d, want %d: %s", w.Code, tt.wantStatus, w.Body.String())
			}
			if got := w.Header().Get("Content-Type"); !strings.HasPrefix(got, "application/json") {
				t.Fatalf("Content-Type = %q, want application/json", got)
			}
			if tt.wantResource == "" {
				return
			}
			var body map[string]any
			if err := json.Unmarshal(w.Body.Bytes(), &body); err != nil {
				t.Fatalf("unmarshal response: %v", err)
			}
			if body["resource"] != tt.wantResource {
				t.Fatalf("resource = %v, want %s", body["resource"], tt.wantResource)
			}
		})
	}
}

func TestRegisterDiscoveryRoutes_OAuthMetadataReturnsJSON404(t *testing.T) {
	gin.SetMode(gin.TestMode)
	router := gin.New()
	RegisterDiscoveryRoutes(router)

	req, _ := http.NewRequest("GET", "/.well-known/oauth-authorization-server/mcp/server/test-server/agent-1", nil)
	w := httptest.NewRecorder()
	router.ServeHTTP(w, req)

	if w.Code != http.StatusNotFound {
		t.Fatalf("status = %d, want 404: %s", w.Code, w.Body.String())
	}
	if got := w.Header().Get("Content-Type"); !strings.HasPrefix(got, "application/json") {
		t.Fatalf("Content-Type = %q, want application/json", got)
	}
}

func TestDispatchRPC_ToolsList(t *testing.T) {
	handler := &mockToolsListHandler{
		tools: []types.MCPToolDef{
			types.MCPToolDef{Name: "tool1", Description: "Tool 1", InputSchema: types.MCPInputSchema{Type: "object", Properties: map[string]types.MCPProperty{}}},
		},
	}

	body, _ := json.Marshal(rpcRequest{
		JSONRPC: "2.0",
		ID:      rawMsg(`"4"`),
		Method:  "tools/list",
	})

	resp, ok := dispatchRPC(context.TODO(), "agent-1", handler, body)
	if !ok {
		t.Fatal("dispatchRPC should return ok=true for tools/list")
	}
	if resp.Error != nil {
		t.Errorf("unexpected error: %v", resp.Error)
	}
}

func TestDispatchRPC_ToolsCall(t *testing.T) {
	handler := &mockToolsListHandler{}

	body, _ := json.Marshal(rpcRequest{
		JSONRPC: "2.0",
		ID:      rawMsg(`"5"`),
		Method:  "tools/call",
		Params:  json.RawMessage(`{"name":"tool1","arguments":{}}`),
	})

	resp, ok := dispatchRPC(context.TODO(), "agent-1", handler, body)
	if !ok {
		t.Fatal("dispatchRPC should return ok=true for tools/call")
	}
	if resp.Error != nil {
		t.Errorf("unexpected error: %v", resp.Error)
	}
}

func TestDispatchRPC_ToolsCallAppliesResultBudget(t *testing.T) {
	body, _ := json.Marshal(rpcRequest{
		JSONRPC: "2.0",
		ID:      rawMsg(`"budget"`),
		Method:  "tools/call",
		Params:  json.RawMessage(`{"name":"tool1","arguments":{}}`),
	})

	resp, ok := dispatchRPC(context.TODO(), "agent-1", largeToolResultHandler{}, body)
	if !ok {
		t.Fatal("dispatchRPC should return ok=true for tools/call")
	}
	result, ok := resp.Result.(*types.MCPToolResult)
	if !ok {
		t.Fatalf("result = %T, want *types.MCPToolResult", resp.Result)
	}
	if got := len([]rune(result.Content[0].Text)); got > 24*1024 {
		t.Fatalf("text runes = %d, want within result budget", got)
	}
	structured, ok := result.StructuredContent.(map[string]any)
	if !ok {
		t.Fatalf("structuredContent = %T, want map", result.StructuredContent)
	}
	if structured["truncated"] != true {
		t.Fatalf("truncation metadata missing: %#v", structured)
	}
}

func TestDynamicRoutesApplyResultBudgetForChannelAndAgentTeam(t *testing.T) {
	for _, tc := range []struct {
		name string
		path string
	}{
		{name: "server", path: "/mcp/server/test-server/agent-1"},
		{name: "agentteam", path: "/api/internal/mcp/servers/test-server/agents/agent-1"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			router := budgetRouteTestRouter(largeToolResultHandler{})
			recorder := postPath(router, tc.path, `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"tool","arguments":{}}}`)
			if recorder.Code != http.StatusOK {
				t.Fatalf("status = %d, want 200: %s", recorder.Code, recorder.Body.String())
			}
			var resp struct {
				Result types.MCPToolResult `json:"result"`
			}
			if err := json.Unmarshal(recorder.Body.Bytes(), &resp); err != nil {
				t.Fatalf("unmarshal response: %v", err)
			}
			if got := len([]rune(resp.Result.Content[0].Text)); got > 24*1024 {
				t.Fatalf("text runes = %d, want within result budget", got)
			}
			structured, ok := resp.Result.StructuredContent.(map[string]any)
			if !ok || structured["truncated"] != true {
				t.Fatalf("truncation metadata missing: %#v", resp.Result.StructuredContent)
			}
		})
	}
}

func TestDispatchRPC_ToolsCall_InvalidParams(t *testing.T) {
	handler := &mockToolsListHandler{}

	body, _ := json.Marshal(rpcRequest{
		JSONRPC: "2.0",
		ID:      rawMsg(`"6"`),
		Method:  "tools/call",
		Params:  json.RawMessage(`{"name":"tool1"}`),
	})

	resp, ok := dispatchRPC(context.TODO(), "agent-1", handler, body)
	if !ok {
		t.Fatal("dispatchRPC should return ok=true")
	}
	if resp.Error != nil {
		t.Errorf("unexpected error: %v", resp.Error)
	}
}

func TestDispatchRPCBatchRejectsOverLimitTrailingBytes(t *testing.T) {
	router := testBatchRouter(&mockToolsListHandler{})
	validPrefix := `{"jsonrpc":"2.0","id":1,"method":"ping"}`
	body := validPrefix + strings.Repeat(" ", maxMCPRequestBodyBytes-len(validPrefix)) + "x"

	recorder := postBatch(router, body)

	if recorder.Code != http.StatusRequestEntityTooLarge {
		t.Fatalf("status = %d, want %d", recorder.Code, http.StatusRequestEntityTooLarge)
	}
}

func TestDispatchRPCBatchAcceptsBodyExactlyAtLimit(t *testing.T) {
	router := testBatchRouter(&mockToolsListHandler{})
	validPrefix := `{"jsonrpc":"2.0","id":1,"method":"ping"}`
	body := validPrefix + strings.Repeat(" ", maxMCPRequestBodyBytes-len(validPrefix))

	recorder := postBatch(router, body)

	if recorder.Code != http.StatusOK {
		t.Fatalf("status = %d, want %d: %s", recorder.Code, http.StatusOK, recorder.Body.String())
	}
}

func TestDispatchRPCBatchEmptyBatchReturnsInvalidRequest(t *testing.T) {
	router := testBatchRouter(&mockToolsListHandler{})

	recorder := postBatch(router, `[]`)

	if recorder.Code != http.StatusBadRequest {
		t.Fatalf("status = %d, want %d: %s", recorder.Code, http.StatusBadRequest, recorder.Body.String())
	}
	var resp rpcResponse
	if err := json.Unmarshal(recorder.Body.Bytes(), &resp); err != nil {
		t.Fatalf("unmarshal response: %v", err)
	}
	if resp.Error == nil || resp.Error.Code != -32600 {
		t.Fatalf("error = %#v, want code -32600", resp.Error)
	}
}

func TestDispatchRPCBatchNotificationOnlyBatchReturnsAccepted(t *testing.T) {
	router := testBatchRouter(&mockToolsListHandler{})
	body := `[{"jsonrpc":"2.0","method":"notifications/initialized"},{"jsonrpc":"2.0","method":"notifications/cancelled"}]`

	recorder := postBatch(router, body)

	if recorder.Code != http.StatusAccepted {
		t.Fatalf("status = %d, want %d: %s", recorder.Code, http.StatusAccepted, recorder.Body.String())
	}
}

func testBatchRouter(handler types.MCPHandler) *gin.Engine {
	gin.SetMode(gin.TestMode)
	router := gin.New()
	group := router.Group("/mcp/server/:serverId/:agentId", func(c *gin.Context) {
		c.Set("agentID", c.Param("agentId"))
		c.Next()
	})
	RegisterHandlerRoutes(group, func(string) (types.MCPHandler, bool) {
		return handler, true
	})
	return router
}

func budgetRouteTestRouter(handler types.MCPHandler) *gin.Engine {
	gin.SetMode(gin.TestMode)
	router := gin.New()
	serverGroup := router.Group("/mcp/server/:serverId/:agentId", func(c *gin.Context) {
		c.Set("agentID", c.Param("agentId"))
		c.Next()
	})
	RegisterHandlerRoutes(serverGroup, func(string) (types.MCPHandler, bool) { return handler, true })
	daemonGroup := router.Group("/api/internal/mcp/servers/:serverId/agents/:agentId", func(c *gin.Context) {
		c.Set("agentID", c.Param("agentId"))
		c.Next()
	})
	RegisterHandlerRoutes(daemonGroup, func(string) (types.MCPHandler, bool) { return handler, true })
	return router
}

func testAccessControlRouter(manager *token.KeyManager) *gin.Engine {
	gin.SetMode(gin.TestMode)
	router := gin.New()
	group := router.Group("/mcp/server/:serverId/:agentId", accessControlWithManagerForTest(manager))
	RegisterHandlerRoutes(group, func(string) (types.MCPHandler, bool) {
		return &mockToolsListHandler{}, true
	})
	return router
}

func testleaderAccessControlRouter(manager *token.KeyManager) *gin.Engine {
	gin.SetMode(gin.TestMode)
	router := gin.New()
	group := router.Group("/mcp/agentteam/:agentId", leaderAccessControlWithManagerForTest(manager))
	RegisterHandlerRoutes(group, func(string) (types.MCPHandler, bool) {
		return &mockToolsListHandler{}, true
	})
	return router
}

func newMCPTestTokenManager(t *testing.T) *token.KeyManager {
	t.Helper()
	manager, err := token.NewKeyManager(filepath.Join(t.TempDir(), "config.yaml"))
	if err != nil {
		t.Fatalf("NewKeyManager() error = %v", err)
	}
	return manager
}

func postBatch(router *gin.Engine, body string) *httptest.ResponseRecorder {
	req := httptest.NewRequest(http.MethodPost, "/mcp/server/test-server/agent-1", strings.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	recorder := httptest.NewRecorder()
	router.ServeHTTP(recorder, req)
	return recorder
}

func postPath(router *gin.Engine, path, body string) *httptest.ResponseRecorder {
	req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	recorder := httptest.NewRecorder()
	router.ServeHTTP(recorder, req)
	return recorder
}

func postAuthenticatedPing(router *gin.Engine, path, tokenValue string) *httptest.ResponseRecorder {
	req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{"jsonrpc":"2.0","id":1,"method":"ping"}`))
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+tokenValue)
	recorder := httptest.NewRecorder()
	router.ServeHTTP(recorder, req)
	return recorder
}