package mcp_test

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"sync"
	"sync/atomic"
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/extension/runner"
	"github.com/tidbcloud/AgentTeam/internal/extension/runner/mcp"
)

func TestReadResourceUsesModernNegotiationAndPostSSE(t *testing.T) {
	t.Parallel()

	var (
		mu      sync.Mutex
		methods []string
	)
	server := httptest.NewTLSServer(http.HandlerFunc(func(
		writer http.ResponseWriter,
		request *http.Request,
	) {
		if request.Method != http.MethodPost {
			t.Errorf("unexpected HTTP method %q", request.Method)
			writer.WriteHeader(http.StatusMethodNotAllowed)
			return
		}
		if request.Header.Get("X-MCP-Key") != "resource-secret" {
			t.Errorf("missing invocation credential")
			writer.WriteHeader(http.StatusUnauthorized)
			return
		}
		body, err := io.ReadAll(request.Body)
		if err != nil {
			t.Errorf("read request: %v", err)
			return
		}
		var rpcRequest struct {
			ID     json.RawMessage `json:"id"`
			Method string          `json:"method"`
			Params json.RawMessage `json:"params"`
		}
		if err := json.Unmarshal(body, &rpcRequest); err != nil {
			t.Errorf("decode request: %v", err)
			return
		}
		mu.Lock()
		methods = append(methods, rpcRequest.Method)
		mu.Unlock()

		switch rpcRequest.Method {
		case "server/discover":
			var params struct {
				Meta map[string]json.RawMessage `json:"_meta"`
			}
			if err := json.Unmarshal(rpcRequest.Params, &params); err != nil {
				t.Errorf("decode discover params: %v", err)
				return
			}
			capabilities := params.Meta["io.modelcontextprotocol/clientCapabilities"]
			if string(capabilities) != "{}" {
				t.Errorf(
					"client capabilities = %s, want empty object",
					capabilities,
				)
				return
			}
			writeModernDiscover(t, writer, rpcRequest.ID)
		case "resources/read":
			var params struct {
				URI string `json:"uri"`
			}
			if err := json.Unmarshal(rpcRequest.Params, &params); err != nil {
				t.Errorf("decode resource params: %v", err)
				return
			}
			if params.URI != "law://statute/42" {
				t.Errorf("resource URI = %q", params.URI)
				return
			}
			writer.Header().Set("Content-Type", "text/event-stream")
			_, err := fmt.Fprintf(
				writer,
				"event: message\ndata: %s\n\n",
				mustJSON(t, map[string]any{
					"jsonrpc": "2.0",
					"id":      json.RawMessage(rpcRequest.ID),
					"result": map[string]any{
						"contents": []map[string]any{
							{
								"uri":      "law://statute/42",
								"mimeType": "text/plain",
								"text":     "article 42",
							},
							{
								"uri":      "law://statute/42/source",
								"mimeType": "application/octet-stream",
								"blob":     []byte("source"),
							},
						},
					},
				}),
			)
			if err != nil {
				t.Errorf("write SSE response: %v", err)
			}
		default:
			t.Errorf("unexpected MCP method %q", rpcRequest.Method)
			writer.WriteHeader(http.StatusBadRequest)
		}
	}))
	defer server.Close()

	client, err := mcp.NewClient(mcp.Options{
		Endpoint:            server.URL + "/mcp",
		Auth:                testAuth,
		AllowedResourceURIs: []string{"law://statute/42"},
		Transport:           server.Client().Transport,
	})
	require.NoError(t, err)

	result, err := client.ReadResource(
		context.Background(),
		runner.Credentials{"apiKey": "resource-secret"},
		"law://statute/42",
	)
	require.NoError(t, err)
	require.Equal(t, mcp.ResourceResult{
		Contents: []mcp.ResourceContent{
			{
				URI:      "law://statute/42",
				MIMEType: "text/plain",
				Text:     "article 42",
			},
			{
				URI:      "law://statute/42/source",
				MIMEType: "application/octet-stream",
				Blob:     []byte("source"),
			},
		},
	}, result)

	mu.Lock()
	defer mu.Unlock()
	require.Equal(t, []string{"server/discover", "resources/read"}, methods)
}

func TestLegacyProtocolVersionAllowlist(t *testing.T) {
	t.Parallel()

	for _, version := range []string{
		"2024-11-05",
		"2025-03-26",
		"2025-06-18",
		"2025-11-25",
	} {
		version := version
		t.Run(version, func(t *testing.T) {
			t.Parallel()
			const sessionID = "legacy-version-session"
			server := httptest.NewTLSServer(http.HandlerFunc(func(
				writer http.ResponseWriter,
				request *http.Request,
			) {
				if request.Method == http.MethodDelete {
					writer.WriteHeader(http.StatusNoContent)
					return
				}
				rpcRequest, ok := decodeRPCRequest(t, request)
				if !ok {
					writer.WriteHeader(http.StatusBadRequest)
					return
				}
				switch rpcRequest.Method {
				case "server/discover":
					writeRPCError(
						t, writer, rpcRequest.ID, -32601, "legacy")
				case "initialize":
					writer.Header().Set("Mcp-Session-Id", sessionID)
					writeRPCResult(t, writer, rpcRequest.ID, map[string]any{
						"protocolVersion": version,
						"capabilities":    map[string]any{},
						"serverInfo": map[string]any{
							"name":    "version-double",
							"version": "1.0.0",
						},
					})
				case "notifications/initialized":
					writer.WriteHeader(http.StatusAccepted)
				case "tools/call":
					writeRPCResult(t, writer, rpcRequest.ID, map[string]any{
						"content": []map[string]any{
							{"type": "text", "text": version},
						},
					})
				default:
					t.Errorf("unexpected method %q", rpcRequest.Method)
					writer.WriteHeader(http.StatusBadRequest)
				}
			}))
			defer server.Close()

			client, err := mcp.NewClient(mcp.Options{
				Endpoint:     server.URL + "/mcp",
				Auth:         testAuth,
				AllowedTools: []string{"fixed_tool"},
				Transport:    server.Client().Transport,
			})
			require.NoError(t, err)
			result, err := client.CallTool(
				context.Background(),
				runner.Credentials{
					"apiKey": "legacy-version-secret",
				},
				"fixed_tool",
				map[string]any{},
			)
			require.NoError(t, err)
			require.Len(t, result.Content, 1)
			require.Contains(t, string(result.Content[0]), version)
		})
	}
}

func TestClientDisablesAutomaticMultiRoundTrip(t *testing.T) {
	t.Parallel()

	var callCount atomic.Int64
	server := httptest.NewTLSServer(http.HandlerFunc(func(
		writer http.ResponseWriter,
		request *http.Request,
	) {
		rpcRequest, ok := decodeRPCRequest(t, request)
		if !ok {
			writer.WriteHeader(http.StatusBadRequest)
			return
		}
		switch rpcRequest.Method {
		case "server/discover":
			writeModernDiscover(t, writer, rpcRequest.ID)
		case "tools/call":
			callCount.Add(1)
			writeRPCResult(t, writer, rpcRequest.ID, map[string]any{
				"resultType": "input_required",
				"inputRequests": map[string]any{
					"confirm": map[string]any{
						"method": "elicitation/create",
						"params": map[string]any{
							"message": "confirm?",
						},
					},
				},
				"requestState": "opaque-state",
			})
		default:
			t.Errorf("unexpected method %q", rpcRequest.Method)
			writer.WriteHeader(http.StatusBadRequest)
		}
	}))
	defer server.Close()

	client, err := mcp.NewClient(mcp.Options{
		Endpoint:     server.URL + "/mcp",
		Auth:         testAuth,
		AllowedTools: []string{"fixed_tool"},
		Transport:    server.Client().Transport,
	})
	require.NoError(t, err)

	_, err = client.CallTool(
		context.Background(),
		runner.Credentials{"apiKey": "mrtr-secret"},
		"fixed_tool",
		map[string]any{},
	)
	require.Equal(t, runner.ErrorUpstreamFailure,
		runner.CodeOf(err))
	require.Equal(t, int64(1), callCount.Load(),
		"automatic multi-round-trip retried the tool call")
}

func writeModernDiscover(
	t *testing.T,
	writer http.ResponseWriter,
	id json.RawMessage,
) {
	t.Helper()
	writeRPCResult(t, writer, id, map[string]any{
		"supportedVersions": []string{"2026-07-28"},
		"capabilities":      map[string]any{},
		"_meta": map[string]any{
			"io.modelcontextprotocol/serverInfo": map[string]any{
				"name":    "modern-double",
				"version": "1.0.0",
			},
		},
	})
}

func mustJSON(t *testing.T, value any) []byte {
	t.Helper()
	encoded, err := json.Marshal(value)
	require.NoError(t, err)
	return encoded
}