package middleware

import (
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	ghcontext "github.com/github/github-mcp-server/pkg/context"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestWithMCPParse(t *testing.T) {
	tests := []struct {
		name           string
		method         string
		path           string
		body           string
		expectInfo     bool
		expectedMethod string
		expectedItem   string
		expectedOwner  string
		expectedRepo   string
		expectedArgs   map[string]any
	}{
		{
			name:       "health check path is skipped",
			method:     http.MethodPost,
			path:       "/_ping",
			body:       `{"jsonrpc":"2.0","method":"tools/list"}`,
			expectInfo: false,
		},
		{
			name:       "GET request is skipped",
			method:     http.MethodGet,
			path:       "/mcp",
			body:       `{"jsonrpc":"2.0","method":"tools/list"}`,
			expectInfo: false,
		},
		{
			name:       "empty body is skipped",
			method:     http.MethodPost,
			path:       "/mcp",
			body:       "",
			expectInfo: false,
		},
		{
			name:       "invalid JSON is skipped",
			method:     http.MethodPost,
			path:       "/mcp",
			body:       "not valid json",
			expectInfo: false,
		},
		{
			name:       "non-JSON-RPC 2.0 is skipped",
			method:     http.MethodPost,
			path:       "/mcp",
			body:       `{"jsonrpc":"1.0","method":"tools/list"}`,
			expectInfo: false,
		},
		{
			name:       "empty method is skipped",
			method:     http.MethodPost,
			path:       "/mcp",
			body:       `{"jsonrpc":"2.0","method":""}`,
			expectInfo: false,
		},
		{
			name:           "tools/list parses method only",
			method:         http.MethodPost,
			path:           "/mcp",
			body:           `{"jsonrpc":"2.0","method":"tools/list"}`,
			expectInfo:     true,
			expectedMethod: "tools/list",
		},
		{
			name:           "tools/call parses name",
			method:         http.MethodPost,
			path:           "/mcp",
			body:           `{"jsonrpc":"2.0","method":"tools/call","params":{"name":"get_file_contents"}}`,
			expectInfo:     true,
			expectedMethod: "tools/call",
			expectedItem:   "get_file_contents",
		},
		{
			name:           "tools/call parses owner and repo from arguments",
			method:         http.MethodPost,
			path:           "/mcp",
			body:           `{"jsonrpc":"2.0","method":"tools/call","params":{"name":"get_file_contents","arguments":{"owner":"github","repo":"github-mcp-server","path":"README.md"}}}`,
			expectInfo:     true,
			expectedMethod: "tools/call",
			expectedItem:   "get_file_contents",
			expectedOwner:  "github",
			expectedRepo:   "github-mcp-server",
			expectedArgs:   map[string]any{"owner": "github", "repo": "github-mcp-server", "path": "README.md"},
		},
		{
			name:           "tools/call with invalid arguments JSON continues without args",
			method:         http.MethodPost,
			path:           "/mcp",
			body:           `{"jsonrpc":"2.0","method":"tools/call","params":{"name":"get_file_contents","arguments":"not an object"}}`,
			expectInfo:     true,
			expectedMethod: "tools/call",
			expectedItem:   "get_file_contents",
		},
		{
			name:           "prompts/get parses name",
			method:         http.MethodPost,
			path:           "/mcp",
			body:           `{"jsonrpc":"2.0","method":"prompts/get","params":{"name":"my_prompt"}}`,
			expectInfo:     true,
			expectedMethod: "prompts/get",
			expectedItem:   "my_prompt",
		},
		{
			name:           "resources/read parses URI as item name",
			method:         http.MethodPost,
			path:           "/mcp",
			body:           `{"jsonrpc":"2.0","method":"resources/read","params":{"uri":"repo://github/github-mcp-server"}}`,
			expectInfo:     true,
			expectedMethod: "resources/read",
			expectedItem:   "repo://github/github-mcp-server",
		},
		{
			name:           "initialize method parses correctly",
			method:         http.MethodPost,
			path:           "/mcp",
			body:           `{"jsonrpc":"2.0","method":"initialize","params":{"capabilities":{}}}`,
			expectInfo:     true,
			expectedMethod: "initialize",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			var capturedInfo *ghcontext.MCPMethodInfo
			var infoCaptured bool

			// Create a handler that captures the MCPMethodInfo from context
			nextHandler := http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
				capturedInfo, infoCaptured = ghcontext.MCPMethod(r.Context())
			})

			middleware := WithMCPParse()
			handler := middleware(nextHandler)

			req := httptest.NewRequest(tt.method, tt.path, strings.NewReader(tt.body))
			rr := httptest.NewRecorder()

			handler.ServeHTTP(rr, req)

			if tt.expectInfo {
				require.True(t, infoCaptured, "MCPMethodInfo should be present in context")
				require.NotNil(t, capturedInfo)
				assert.Equal(t, tt.expectedMethod, capturedInfo.Method)
				assert.Equal(t, tt.expectedItem, capturedInfo.ItemName)
				assert.Equal(t, tt.expectedOwner, capturedInfo.Owner)
				assert.Equal(t, tt.expectedRepo, capturedInfo.Repo)
				if tt.expectedArgs != nil {
					assert.Equal(t, tt.expectedArgs, capturedInfo.Arguments)
				}
			} else {
				assert.False(t, infoCaptured, "MCPMethodInfo should not be present in context")
			}
		})
	}
}

func TestWithMCPParse_BodyRestoration(t *testing.T) {
	originalBody := `{"jsonrpc":"2.0","method":"tools/call","params":{"name":"test_tool"}}`

	var capturedBody string

	nextHandler := http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
		body, err := io.ReadAll(r.Body)
		require.NoError(t, err)
		capturedBody = string(body)
	})

	middleware := WithMCPParse()
	handler := middleware(nextHandler)

	req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(originalBody))
	rr := httptest.NewRecorder()

	handler.ServeHTTP(rr, req)

	assert.Equal(t, originalBody, capturedBody, "body should be restored for downstream handlers")
}