package transport

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

type admittedKey struct{}

func fixture(t *testing.T) (*Handler, *int, *int) {
	t.Helper()
	authCalls, toolCalls := 0, 0
	h, err := NewHandler(Config{PublicOrigin: "https://example.test", ResourceURL: "https://example.test/mcp", ServerName: "fixture", ServerVersion: "1", AllowedOrigin: func(origin string) bool { return origin == "https://example.test" }, Authorize: func(ctx context.Context, bearer string, _ string) (context.Context, error) {
		authCalls++
		switch bearer {
		case "valid":
			return context.WithValue(ctx, admittedKey{}, true), nil
		case "forbidden":
			return nil, Forbidden
		case "busy":
			return nil, RateLimited
		case "secret-error":
			return nil, errors.New("secret internal failure")
		}
		return nil, Unauthenticated
	}, Tools: []Tool{{Name: "read", InputSchema: json.RawMessage(`{"type":"object"}`)}}, Execute: func(ctx context.Context, name string, args json.RawMessage) (any, error) {
		toolCalls++
		if ctx.Value(admittedKey{}) != true {
			t.Error("tool received no admitted context")
		}
		if string(args) == `{"fail":true}` {
			return nil, errors.New("sensitive query")
		}
		return map[string]any{"content": []any{}}, nil
	}})
	if err != nil {
		t.Fatal(err)
	}
	return h, &authCalls, &toolCalls
}
func request(h *Handler, method, path, body, bearer string, headers map[string]string) *httptest.ResponseRecorder {
	r := httptest.NewRequest(method, "https://attacker.invalid"+path, strings.NewReader(body))
	r.Header.Set("Content-Type", "application/json")
	r.Header.Set("Accept", "application/json, text/event-stream")
	if bearer != "" {
		r.Header.Set("Authorization", "Bearer "+bearer)
	}
	for k, v := range headers {
		r.Header.Set(k, v)
	}
	w := httptest.NewRecorder()
	h.ServeHTTP(w, r)
	return w
}
func TestEveryPostInvokesRuntimeUserAdmission(t *testing.T) {
	h, auth, tools := fixture(t)
	for _, body := range []string{`{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-11-25"}}`, `{"jsonrpc":"2.0","id":2,"method":"tools/list"}`, `{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"read","arguments":{}}}`, `{"jsonrpc":"2.0","method":"notifications/initialized"}`} {
		if w := request(h, "POST", "/mcp", body, "", nil); w.Code != 401 || w.Header().Get("WWW-Authenticate") != `Bearer realm="mcp"` {
			t.Fatalf("missing bearer: %d", w.Code)
		}
		if w := request(h, "POST", "/mcp", body, "valid", nil); w.Code != 200 && w.Code != 202 {
			t.Fatalf("admitted request: %d", w.Code)
		}
	}
	if *auth != 4 || *tools != 1 {
		t.Fatalf("admission/tool calls: %d/%d", *auth, *tools)
	}
}
func TestOriginDeniedBeforeBearerAndDiscoveryIgnoresHost(t *testing.T) {
	h, auth, _ := fixture(t)
	if w := request(h, "POST", "/mcp", "{}", "valid", map[string]string{"Origin": "https://evil.test"}); w.Code != 403 {
		t.Fatal(w.Code)
	}
	if *auth != 0 {
		t.Fatal("disallowed origin reached bearer admission")
	}
	for _, path := range []string{"/.well-known/oauth-protected-resource"} {
		w := request(h, "GET", path, "", "", nil)
		if w.Code != 200 || strings.Contains(w.Body.String(), "attacker") {
			t.Fatalf("unsafe discovery: %d", w.Code)
		}
	}
}
func TestProtocolAndFailureBoundaries(t *testing.T) {
	h, _, tools := fixture(t)
	for _, tt := range []struct {
		name, method, body, bearer string
		headers                    map[string]string
		code                       int
	}{
		{name: "wrong token", method: "POST", body: "{}", bearer: "invalid", code: 401},
		{name: "permission", method: "POST", body: "{}", bearer: "forbidden", code: 403},
		{name: "quota", method: "POST", body: "{}", bearer: "busy", code: 429},
		{name: "owner failure", method: "POST", body: "{}", bearer: "secret-error", code: 503},
		{name: "SSE", method: "GET", code: 405},
		{name: "method", method: "DELETE", code: 405},
		{name: "media", method: "POST", body: "{}", bearer: "valid", headers: map[string]string{"Accept": "application/json"}, code: 406},
		{name: "charset", method: "POST", body: "{}", bearer: "valid", headers: map[string]string{"Content-Type": "application/json;charset=utf-16"}, code: 415},
		{name: "version", method: "POST", body: "{}", bearer: "valid", headers: map[string]string{"MCP-Protocol-Version": "unsupported"}, code: 400},
		{name: "oversize", method: "POST", body: strings.Repeat(" ", maxRequestBytes+1), bearer: "valid", code: 413},
		{name: "batch", method: "POST", body: `[]`, bearer: "valid", code: 400},
		{name: "bad id", method: "POST", body: `{"jsonrpc":"2.0","id":{},"method":"ping"}`, bearer: "valid", code: 400},
		{name: "utf8", method: "POST", body: string([]byte{255}), bearer: "valid", code: 400},
	} {
		t.Run(tt.name, func(t *testing.T) {
			w := request(h, tt.method, "/mcp", tt.body, tt.bearer, tt.headers)
			if w.Code != tt.code || strings.Contains(w.Body.String(), "secret") {
				t.Fatalf("got %d want %d", w.Code, tt.code)
			}
		})
	}
	if *tools != 0 {
		t.Fatal("invalid request dispatched a tool")
	}
}
func TestUnknownWritesAndToolErrors(t *testing.T) {
	h, _, tools := fixture(t)
	w := request(h, "POST", "/mcp", `{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"write","arguments":{}}}`, "valid", nil)
	if w.Code != 200 || !strings.Contains(w.Body.String(), `"code":-32602`) || *tools != 0 {
		t.Fatal("unknown tool was not rejected")
	}
	w = request(h, "POST", "/mcp", `{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"read","arguments":{"fail":true}}}`, "valid", nil)
	if !strings.Contains(w.Body.String(), `"isError":true`) || strings.Contains(w.Body.String(), "sensitive") {
		t.Fatal("unsafe tool failure")
	}
}
func TestNotificationAndNegotiation(t *testing.T) {
	h, _, _ := fixture(t)
	w := request(h, "POST", "/mcp", `{"jsonrpc":"2.0","method":"notifications/initialized"}`, "valid", nil)
	if w.Code != 202 || w.Body.Len() != 0 {
		t.Fatal("notification returned body")
	}
	w = request(h, "POST", "/mcp", `{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"future"}}`, "valid", nil)
	if w.Code != 200 || !strings.Contains(w.Body.String(), `"id":1`) || !strings.Contains(w.Body.String(), latestMCPProtocolVersion) {
		t.Fatal("version negotiation failed")
	}
}
func TestRepeatedAuthorizationHeadersRejected(t *testing.T) {
	h, auth, _ := fixture(t)
	r := httptest.NewRequest(http.MethodPost, "https://example.test/mcp", strings.NewReader(`{}`))
	r.Header.Add("Authorization", "Bearer valid")
	r.Header.Add("Authorization", "Bearer valid")
	w := httptest.NewRecorder()
	h.ServeHTTP(w, r)
	if w.Code != 401 || *auth != 0 {
		t.Fatal("ambiguous bearer admitted")
	}
}

func TestNotificationsCannotExecuteTools(t *testing.T) {
	h, _, tools := fixture(t)
	for _, body := range []string{
		`{"jsonrpc":"2.0","method":"tools/call","params":{"name":"read","arguments":{}}}`,
		`{"jsonrpc":"2.0","id":null,"method":"tools/call","params":{"name":"read","arguments":{}}}`,
		`{"jsonrpc":"2.0","id":1,"method":"notifications/initialized"}`,
	} {
		w := request(h, "POST", "/mcp", body, "valid", nil)
		if w.Code != 400 || *tools != 0 {
			t.Fatal("invalid request executed a tool")
		}
	}
}