package mcp_test

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

	"github.com/stretchr/testify/require"

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

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

	testCases := []struct {
		code     int64
		expected runner.ErrorCode
	}{
		{code: -32700, expected: runner.ErrorUpstreamFailure},
		{code: -32600, expected: runner.ErrorUpstreamFailure},
		{code: -32601, expected: runner.ErrorNotFound},
		{code: -32602, expected: runner.ErrorValidation},
		{code: -32603, expected: runner.ErrorUpstreamFailure},
	}
	for _, testCase := range testCases {
		testCase := testCase
		t.Run(fmt.Sprintf("code_%d", testCase.code), func(t *testing.T) {
			t.Parallel()
			rawMessage := fmt.Sprintf(
				"raw-jsonrpc-secret-%d", testCase.code)
			client, closeServer := newLegacyToolClient(
				t,
				func(
					writer http.ResponseWriter,
					request rpcRequest,
				) {
					writeRPCError(
						t,
						writer,
						request.ID,
						testCase.code,
						rawMessage,
					)
				},
			)
			defer closeServer()

			_, err := client.CallTool(
				context.Background(),
				runner.Credentials{"apiKey": "error-secret"},
				"fixed_tool",
				map[string]any{},
			)
			require.Equal(t, testCase.expected, runner.CodeOf(err))
			requireRuntimeErrorHasNoRawCause(t, err, rawMessage)
		})
	}
}

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

	testCases := []struct {
		status   int
		expected runner.ErrorCode
	}{
		{status: http.StatusUnauthorized,
			expected: runner.ErrorProviderAuthUnavailable},
		{status: http.StatusForbidden,
			expected: runner.ErrorProviderAuthUnavailable},
		{status: http.StatusNotFound,
			expected: runner.ErrorNotFound},
		{status: http.StatusTooManyRequests,
			expected: runner.ErrorRateLimited},
		{status: http.StatusUnprocessableEntity,
			expected: runner.ErrorValidation},
		{status: http.StatusInternalServerError,
			expected: runner.ErrorUpstreamFailure},
	}
	for _, testCase := range testCases {
		testCase := testCase
		t.Run(http.StatusText(testCase.status), func(t *testing.T) {
			t.Parallel()
			rawBody := fmt.Sprintf(
				"raw-http-secret-%d", testCase.status)
			client, closeServer := newLegacyToolClient(
				t,
				func(
					writer http.ResponseWriter,
					_ rpcRequest,
				) {
					writer.Header().Set(
						"Content-Type", "text/plain")
					writer.WriteHeader(testCase.status)
					_, _ = writer.Write([]byte(rawBody))
				},
			)
			defer closeServer()

			_, err := client.CallTool(
				context.Background(),
				runner.Credentials{"apiKey": "http-secret"},
				"fixed_tool",
				map[string]any{},
			)
			require.Equal(t, testCase.expected, runner.CodeOf(err))
			requireRuntimeErrorHasNoRawCause(t, err, rawBody)
		})
	}
}

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

	const rawToolError = "raw-tool-secret"
	client, closeServer := newLegacyToolClient(
		t,
		func(writer http.ResponseWriter, request rpcRequest) {
			writeRPCResult(t, writer, request.ID, map[string]any{
				"content": []map[string]any{
					{"type": "text", "text": rawToolError},
				},
				"isError": true,
			})
		},
	)
	defer closeServer()

	_, err := client.CallTool(
		context.Background(),
		runner.Credentials{"apiKey": "tool-secret"},
		"fixed_tool",
		map[string]any{},
	)
	require.Equal(t, runner.ErrorUpstreamFailure,
		runner.CodeOf(err))
	requireRuntimeErrorHasNoRawCause(t, err, rawToolError)
}

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

	const malformed = "raw-malformed-secret"
	client, closeServer := newLegacyToolClient(
		t,
		func(writer http.ResponseWriter, _ rpcRequest) {
			writer.Header().Set("Content-Type", "text/event-stream")
			_, _ = fmt.Fprintf(
				writer,
				"event: message\ndata: %s\n\n",
				malformed,
			)
		},
	)
	defer closeServer()

	_, err := client.CallTool(
		context.Background(),
		runner.Credentials{"apiKey": "sse-secret"},
		"fixed_tool",
		map[string]any{},
	)
	require.Equal(t, runner.ErrorUpstreamFailure,
		runner.CodeOf(err))
	requireRuntimeErrorHasNoRawCause(t, err, malformed)
}

func newLegacyToolClient(
	t *testing.T,
	toolHandler func(http.ResponseWriter, rpcRequest),
) (*mcp.Client, func()) {
	t.Helper()

	const sessionID = "error-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)
			writeLegacyInitialize(t, writer, rpcRequest.ID)
		case "notifications/initialized":
			writer.WriteHeader(http.StatusAccepted)
		case "tools/call":
			toolHandler(writer, rpcRequest)
		default:
			t.Errorf("unexpected method %q", rpcRequest.Method)
			writer.WriteHeader(http.StatusBadRequest)
		}
	}))

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