package mcp_test

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

	"github.com/stretchr/testify/require"

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

var testAuth = runner.AuthSpec{
	Kind:   runner.AuthHeader,
	Header: "X-MCP-Key",
	Fields: []runner.CredentialField{
		{
			Name:        "apiKey",
			DisplayName: "API key",
			Required:    true,
			Secret:      true,
		},
	},
}

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

	testCases := []struct {
		name    string
		options mcp.Options
	}{
		{
			name: "http endpoint",
			options: mcp.Options{
				Endpoint:     "http://mcp.example.com/mcp",
				Auth:         testAuth,
				AllowedTools: []string{"search_cases"},
			},
		},
		{
			name: "endpoint query",
			options: mcp.Options{
				Endpoint:     "https://mcp.example.com/mcp?token=secret",
				Auth:         testAuth,
				AllowedTools: []string{"search_cases"},
			},
		},
		{
			name: "endpoint fragment",
			options: mcp.Options{
				Endpoint:     "https://mcp.example.com/mcp#fragment",
				Auth:         testAuth,
				AllowedTools: []string{"search_cases"},
			},
		},
		{
			name: "endpoint userinfo",
			options: mcp.Options{
				Endpoint:     "https://user:mypass@mcp.example.com/mcp",
				Auth:         testAuth,
				AllowedTools: []string{"search_cases"},
			},
		},
		{
			name: "invalid auth",
			options: mcp.Options{
				Endpoint: "https://mcp.example.com/mcp",
				Auth: runner.AuthSpec{
					Kind: runner.AuthHeader,
				},
				AllowedTools: []string{"search_cases"},
			},
		},
		{
			name: "request budget above shared ceiling",
			options: mcp.Options{
				Endpoint:        "https://mcp.example.com/mcp",
				Auth:            testAuth,
				MaxRequestBytes: runner.MaxResponseBytes + 1,
				AllowedTools:    []string{"search_cases"},
			},
		},
		{
			name: "response budget above shared ceiling",
			options: mcp.Options{
				Endpoint:         "https://mcp.example.com/mcp",
				Auth:             testAuth,
				MaxResponseBytes: runner.MaxResponseBytes + 1,
				AllowedTools:     []string{"search_cases"},
			},
		},
		{
			name: "negative cleanup timeout",
			options: mcp.Options{
				Endpoint:       "https://mcp.example.com/mcp",
				Auth:           testAuth,
				CleanupTimeout: -time.Second,
				AllowedTools:   []string{"search_cases"},
			},
		},
		{
			name: "no allowed operation",
			options: mcp.Options{
				Endpoint: "https://mcp.example.com/mcp",
				Auth:     testAuth,
			},
		},
		{
			name: "probe only with operation authority",
			options: mcp.Options{
				Endpoint:     "https://mcp.example.com/mcp",
				Auth:         testAuth,
				ProbeOnly:    true,
				AllowedTools: []string{"search_cases"},
			},
		},
		{
			name: "duplicate tool",
			options: mcp.Options{
				Endpoint:     "https://mcp.example.com/mcp",
				Auth:         testAuth,
				AllowedTools: []string{"search_cases", "search_cases"},
			},
		},
		{
			name: "blank resource URI",
			options: mcp.Options{
				Endpoint:            "https://mcp.example.com/mcp",
				Auth:                testAuth,
				AllowedResourceURIs: []string{""},
			},
		},
	}

	for _, testCase := range testCases {
		testCase := testCase
		t.Run(testCase.name, func(t *testing.T) {
			t.Parallel()
			client, err := mcp.NewClient(testCase.options)
			require.Nil(t, client)
			require.Equal(t, runner.ErrorValidation, runner.CodeOf(err))
			require.NotContains(t, err.Error(), "secret")
			require.NotContains(t, err.Error(), "mypass")
		})
	}
}

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

	const sessionID = "session-legacy"
	var (
		mu      sync.Mutex
		methods []string
	)
	server := httptest.NewTLSServer(http.HandlerFunc(func(
		writer http.ResponseWriter,
		request *http.Request,
	) {
		require.Equal(t, "/mcp", request.URL.Path)
		require.Equal(t, "test-secret", request.Header.Get("X-MCP-Key"))
		require.Empty(t, request.URL.RawQuery)
		if request.Method == http.MethodDelete {
			require.Equal(t, sessionID, request.Header.Get("Mcp-Session-Id"))
			mu.Lock()
			methods = append(methods, "DELETE")
			mu.Unlock()
			writer.WriteHeader(http.StatusNoContent)
			return
		}
		require.Equal(t, http.MethodPost, request.Method)
		body, err := io.ReadAll(request.Body)
		require.NoError(t, err)
		var rpcRequest struct {
			JSONRPC string          `json:"jsonrpc"`
			ID      json.RawMessage `json:"id"`
			Method  string          `json:"method"`
			Params  json.RawMessage `json:"params"`
		}
		require.NoError(t, json.Unmarshal(body, &rpcRequest))
		require.Equal(t, "2.0", rpcRequest.JSONRPC)
		mu.Lock()
		methods = append(methods, rpcRequest.Method)
		mu.Unlock()

		switch rpcRequest.Method {
		case "server/discover":
			writeRPCError(t, writer, rpcRequest.ID, -32601, "legacy only")
		case "initialize":
			writer.Header().Set("Mcp-Session-Id", sessionID)
			writeRPCResult(t, writer, rpcRequest.ID, map[string]any{
				"protocolVersion": "2025-06-18",
				"capabilities":    map[string]any{},
				"serverInfo": map[string]any{
					"name":    "legacy-double",
					"version": "1.0.0",
				},
			})
		case "notifications/initialized":
			require.Equal(t, sessionID, request.Header.Get("Mcp-Session-Id"))
			writer.WriteHeader(http.StatusAccepted)
		case "tools/call":
			require.Equal(t, sessionID, request.Header.Get("Mcp-Session-Id"))
			var params struct {
				Name      string         `json:"name"`
				Arguments map[string]any `json:"arguments"`
			}
			require.NoError(t, json.Unmarshal(rpcRequest.Params, &params))
			require.Equal(t, "search_cases", params.Name)
			query, _ := params.Arguments["query"].(string)
			if query == "null" {
				writeRPCResult(t, writer, rpcRequest.ID, map[string]any{
					"content": []map[string]any{
						{"type": "text", "text": `{"fallback":true}`},
					},
					"structuredContent": nil,
				})
				return
			}
			require.Equal(t, "contract", query)
			writeRPCResult(t, writer, rpcRequest.ID, map[string]any{
				"content": []map[string]any{
					{"type": "text", "text": "case-result"},
				},
				"structuredContent": map[string]any{
					"count": float64(1),
				},
			})
		default:
			require.Failf(t, "unexpected MCP method", "method=%q", rpcRequest.Method)
		}
	}))
	defer server.Close()

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

	result, err := client.CallTool(
		context.Background(),
		runner.Credentials{"apiKey": "test-secret"},
		"search_cases",
		map[string]any{"query": "contract"},
	)
	require.NoError(t, err)
	require.Len(t, result.Content, 1)
	require.True(t, result.StructuredContentPresent)
	require.JSONEq(t, `{"type":"text","text":"case-result"}`, string(result.Content[0]))
	require.JSONEq(t, `{"count":1}`, string(result.StructuredContent))

	explicitNull, err := client.CallTool(
		context.Background(),
		runner.Credentials{"apiKey": "test-secret"},
		"search_cases",
		map[string]any{"query": "null"},
	)
	require.NoError(t, err)
	require.True(t, explicitNull.StructuredContentPresent)
	require.Empty(t, explicitNull.StructuredContent)
	require.Len(t, explicitNull.Content, 1)

	mu.Lock()
	defer mu.Unlock()
	require.Equal(t, []string{
		"server/discover",
		"initialize",
		"notifications/initialized",
		"tools/call",
		"DELETE",
		"server/discover",
		"initialize",
		"notifications/initialized",
		"tools/call",
		"DELETE",
	}, methods)
}

func TestCallToolBindsStructuredContentPresenceToMatchingSSEResponse(
	t *testing.T,
) {
	t.Parallel()

	for _, testCase := range []struct {
		name        string
		mode        string
		wantPresent bool
		ambiguous   bool
	}{
		{
			name:        "decoy absent actual null",
			mode:        "decoy-absent",
			wantPresent: true,
		},
		{name: "decoy present actual absent", mode: "decoy-present"},
		{name: "malformed event", mode: "malformed", ambiguous: true},
		{name: "multiple matching", mode: "multiple", ambiguous: true},
	} {
		testCase := testCase
		t.Run(testCase.name, func(t *testing.T) {
			t.Parallel()

			server := httptest.NewTLSServer(http.HandlerFunc(func(
				writer http.ResponseWriter,
				request *http.Request,
			) {
				rpcRequest, ok := decodeRPCRequest(t, request)
				require.True(t, ok)
				switch rpcRequest.Method {
				case "server/discover":
					writeModernDiscover(t, writer, rpcRequest.ID)
				case "tools/call":
					writer.Header().Set("Content-Type", "text/event-stream")
					actual := map[string]any{
						"jsonrpc": "2.0",
						"id":      json.RawMessage(rpcRequest.ID),
						"result": map[string]any{
							"content": []map[string]any{{
								"type": "text",
								"text": `{"rows":[]}`,
							}},
						},
					}
					decoy := map[string]any{
						"jsonrpc": "2.0",
						"id":      "decoy",
						"result": map[string]any{
							"content": []map[string]any{},
						},
					}
					switch testCase.mode {
					case "decoy-absent":
						actual["result"].(map[string]any)["structuredContent"] = nil
					case "decoy-present":
						decoy["result"].(map[string]any)["structuredContent"] = nil
					case "malformed":
						_, err := fmt.Fprint(
							writer,
							"event: message\ndata: {malformed\n\n",
						)
						require.NoError(t, err)
					case "multiple":
						_, err := fmt.Fprintf(
							writer,
							"event: message\ndata: %s\n\n",
							mustJSON(t, actual),
						)
						require.NoError(t, err)
					}
					if strings.HasPrefix(testCase.mode, "decoy-") {
						_, err := fmt.Fprintf(
							writer,
							"event: message\ndata: %s\n\n",
							mustJSON(t, decoy),
						)
						require.NoError(t, err)
					}
					_, err := fmt.Fprintf(
						writer,
						"event: message\ndata: %s\n\n",
						mustJSON(t, actual),
					)
					require.NoError(t, err)
				default:
					require.Failf(
						t,
						"unexpected MCP method",
						"method=%q",
						rpcRequest.Method,
					)
				}
			}))
			defer server.Close()

			client, err := mcp.NewClient(mcp.Options{
				Endpoint:     server.URL + "/mcp",
				Auth:         testAuth,
				AllowedTools: []string{"search_cases"},
				Transport:    server.Client().Transport,
			})
			require.NoError(t, err)
			result, err := client.CallTool(
				context.Background(),
				runner.Credentials{"apiKey": "test-secret"},
				"search_cases",
				map[string]any{"query": "contract"},
			)
			if testCase.ambiguous && err != nil {
				require.Equal(
					t,
					runner.ErrorUpstreamFailure,
					runner.CodeOf(err),
				)
				return
			}
			require.NoError(t, err)
			require.Equal(
				t,
				testCase.wantPresent || testCase.ambiguous,
				result.StructuredContentPresent,
			)
		})
	}
}

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

	const sessionID = "probe-session"
	var (
		mu      sync.Mutex
		methods []string
	)
	server := httptest.NewTLSServer(http.HandlerFunc(func(
		writer http.ResponseWriter,
		request *http.Request,
	) {
		require.Equal(t, "probe-secret", request.Header.Get("X-MCP-Key"))
		if request.Method == http.MethodDelete {
			require.Equal(t, sessionID, request.Header.Get("Mcp-Session-Id"))
			mu.Lock()
			methods = append(methods, "DELETE")
			mu.Unlock()
			writer.WriteHeader(http.StatusNoContent)
			return
		}
		rpcRequest, ok := decodeRPCRequest(t, request)
		require.True(t, ok)
		mu.Lock()
		methods = append(methods, rpcRequest.Method)
		mu.Unlock()
		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": "2025-06-18",
				"capabilities":    map[string]any{},
				"serverInfo": map[string]any{
					"name": "probe-double", "version": "1.0.0",
				},
			})
		case "notifications/initialized":
			writer.WriteHeader(http.StatusAccepted)
		default:
			require.Failf(
				t,
				"probe issued a capability operation",
				"method=%q",
				rpcRequest.Method,
			)
		}
	}))
	defer server.Close()

	client, err := mcp.NewClient(mcp.Options{
		Endpoint:  server.URL + "/mcp",
		Auth:      testAuth,
		ProbeOnly: true,
		Transport: server.Client().Transport,
	})
	require.NoError(t, err)
	require.NoError(t, client.Probe(
		context.Background(),
		runner.Credentials{"apiKey": "probe-secret"},
	))

	mu.Lock()
	defer mu.Unlock()
	require.Equal(t, []string{
		"server/discover",
		"initialize",
		"notifications/initialized",
		"DELETE",
	}, methods)
}

func writeRPCResult(
	t *testing.T,
	writer http.ResponseWriter,
	id json.RawMessage,
	result any,
) {
	t.Helper()
	writer.Header().Set("Content-Type", "application/json")
	require.NoError(t, json.NewEncoder(writer).Encode(map[string]any{
		"jsonrpc": "2.0",
		"id":      json.RawMessage(id),
		"result":  result,
	}))
}

func writeRPCError(
	t *testing.T,
	writer http.ResponseWriter,
	id json.RawMessage,
	code int64,
	message string,
) {
	t.Helper()
	writer.Header().Set("Content-Type", "application/json")
	require.NoError(t, json.NewEncoder(writer).Encode(map[string]any{
		"jsonrpc": "2.0",
		"id":      json.RawMessage(id),
		"error": map[string]any{
			"code":    code,
			"message": message,
		},
	}))
}