package mcp

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

// 本文件用 httptest 起真的 HTTP MCP server，端到端验证 Streamable HTTP 传输：
// modern（2026-07-28 无会话）/ SSE 响应 / legacy 回退（initialize + Mcp-Session-Id）/ 401 鉴权。
// stdio 的用例在 mcp_test.go，不受本文件影响。

// httpRecorder 记录服务端看到的请求（httptest 的 handler 在另一个 goroutine，需要加锁）。
type httpRecorder struct {
	mu          sync.Mutex
	reqs        []recordedReq
	sessionSeen bool
}

type recordedReq struct {
	method string
	header http.Header
}

func (r *httpRecorder) add(method string, h http.Header) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.reqs = append(r.reqs, recordedReq{method: method, header: h.Clone()})
	if h.Get("Mcp-Session-Id") != "" {
		r.sessionSeen = true
	}
}

// every 检查是否所有请求都满足 fn。
func (r *httpRecorder) every(fn func(recordedReq) bool) bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	if len(r.reqs) == 0 {
		return false
	}
	for _, q := range r.reqs {
		if !fn(q) {
			return false
		}
	}
	return true
}

func (r *httpRecorder) count() int {
	r.mu.Lock()
	defer r.mu.Unlock()
	return len(r.reqs)
}

func (r *httpRecorder) sawSession() bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	return r.sessionSeen
}

// hasMethodName 检查某次请求是否带了正确的 Mcp-Name 头。
func (r *httpRecorder) hasMethodName(method, name string) bool {
	r.mu.Lock()
	defer r.mu.Unlock()
	for _, q := range r.reqs {
		if q.method == method {
			return q.header.Get("Mcp-Name") == name
		}
	}
	return false
}

func mustJSON(v any) string {
	b, _ := json.Marshal(v)
	return string(b)
}

// modernServer 起一个现代（无会话）MCP HTTP server。mode 决定响应形式：json 或 sse。
func modernServer(rec *httpRecorder, mode string) *httptest.Server {
	write := func(w http.ResponseWriter, resp map[string]any) {
		if mode == "sse" {
			w.Header().Set("Content-Type", "text/event-stream")
			// 先来 keep-alive 注释与一条与本次请求相关的通知，客户端应跳过它们，只取最后的响应。
			fmt.Fprintf(w, ": keep-alive\r\n\r\n")
			fmt.Fprintf(w, "data: %s\r\n\r\n", mustJSON(map[string]any{
				"jsonrpc": "2.0",
				"method":  "notifications/progress",
				"params":  map[string]any{"progress": 10},
			}))
			fmt.Fprintf(w, "data: %s\r\n\r\n", mustJSON(resp))
			if f, ok := w.(http.Flusher); ok {
				f.Flush()
			}
			return
		}
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(resp)
	}

	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var req struct {
			ID     json.RawMessage `json:"id"`
			Method string          `json:"method"`
		}
		_ = json.NewDecoder(r.Body).Decode(&req)
		rec.add(req.Method, r.Header)

		resp := map[string]any{"jsonrpc": "2.0", "id": req.ID}
		switch req.Method {
		case "server/discover":
			resp["result"] = map[string]any{
				"resultType":        "complete",
				"supportedVersions": []string{"2026-07-28"},
				"capabilities":      map[string]any{"tools": map[string]any{}},
				"_meta":             map[string]any{metaServerInfo: map[string]any{"name": "fake-http", "version": "3.0.0"}},
			}
		case "tools/list":
			resp["result"] = map[string]any{
				"resultType": "complete",
				"tools": []map[string]any{{
					"name":        "echo",
					"description": "Echo",
					"inputSchema": map[string]any{"type": "object"},
				}},
			}
		case "tools/call":
			resp["result"] = map[string]any{
				"resultType": "complete",
				"content":    []map[string]any{{"type": "text", "text": "http echo"}},
				"isError":    false,
			}
		default:
			resp["error"] = map[string]any{"code": -32601, "message": "no such method"}
		}
		write(w, resp)
	}))
}

// TestHTTPModernProtocol：现代 Streamable HTTP。
// 核心断言：每个请求都带 MCP-Protocol-Version / 与 body 一致的 Mcp-Method，
// tools/call 带 Mcp-Name=echo，Accept 同时接受两种 content type，且不用 Mcp-Session-Id。

func TestHTTPModernProtocol(t *testing.T) {
	rec := &httpRecorder{}
	srv := modernServer(rec, "json")
	defer srv.Close()

	cli, err := Start("httptest", Config{URL: srv.URL})
	if err != nil {
		t.Fatalf("start: %v", err)
	}
	defer cli.Close()

	if cli.Era() != eraModern {
		t.Fatalf("era = %q, want modern", cli.Era())
	}
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	got, err := cli.CallTool(ctx, "echo", `{"text":"x"}`)
	if err != nil {
		t.Fatalf("call: %v", err)
	}
	if got != "http echo" {
		t.Fatalf("call result: %q", got)
	}

	if !rec.every(func(q recordedReq) bool {
		return q.header.Get("MCP-Protocol-Version") == protocolVersion
	}) {
		t.Errorf("每个请求都必须带 MCP-Protocol-Version %s", protocolVersion)
	}
	if !rec.every(func(q recordedReq) bool {
		return q.header.Get("Mcp-Method") == q.method
	}) {
		t.Errorf("每个请求都必须带与 body 一致的 Mcp-Method 头（%d 个请求）", rec.count())
	}
	if !rec.every(func(q recordedReq) bool {
		a := q.header.Get("Accept")
		return strings.Contains(a, "application/json") && strings.Contains(a, "text/event-stream")
	}) {
		t.Error("Accept 必须同时列 application/json 与 text/event-stream")
	}
	if !rec.hasMethodName("tools/call", "echo") {
		t.Error("tools/call 必须带 Mcp-Name 头（值为 params.name）")
	}
	if rec.sawSession() {
		t.Error("现代传输不应使用 Mcp-Session-Id")
	}
}

// TestHTTPSSEResponse：服务端用 text/event-stream 回：流上先来注释与通知，最后才是响应。
func TestHTTPSSEResponse(t *testing.T) {
	rec := &httpRecorder{}
	srv := modernServer(rec, "sse")
	defer srv.Close()

	cli, err := Start("httptest", Config{URL: srv.URL})
	if err != nil {
		t.Fatalf("start: %v", err)
	}
	defer cli.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	got, err := cli.CallTool(ctx, "echo", `{"text":"x"}`)
	if err != nil {
		t.Fatalf("call: %v", err)
	}
	if got != "http echo" {
		t.Fatalf("call result: %q", got)
	}
}

// TestHTTPLegacyFallback：服务端不认识 server/discover（404 + 非现代 body）→ 回退
// initialize，之后每个请求都要带上服务端下发的 Mcp-Session-Id。
func TestHTTPLegacyFallback(t *testing.T) {
	rec := &httpRecorder{}
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var req struct {
			ID     json.RawMessage `json:"id"`
			Method string          `json:"method"`
		}
		_ = json.NewDecoder(r.Body).Decode(&req)
		rec.add(req.Method, r.Header)

		if req.Method == "server/discover" {
			// 旧服务端：404 且 body 不是 JSON-RPC 错误 → 客户端应判为 legacy。
			w.WriteHeader(http.StatusNotFound)
			_, _ = w.Write([]byte("not found"))
			return
		}
		w.Header().Set("Content-Type", "application/json")
		resp := map[string]any{"jsonrpc": "2.0", "id": req.ID}
		switch req.Method {
		case "initialize":
			// 旧版语义：initialize 响应里下发会话 id。
			w.Header().Set("Mcp-Session-Id", "sess-abc")
			resp["result"] = map[string]any{
				"protocolVersion": legacyVersion,
				"capabilities":    map[string]any{"tools": map[string]any{}},
				"serverInfo":      map[string]any{"name": "fake-http-legacy", "version": "1.2.3"},
			}
		case "tools/list":
			resp["result"] = map[string]any{
				"tools": []map[string]any{{"name": "echo", "inputSchema": map[string]any{"type": "object"}}},
			}
		case "tools/call":
			resp["result"] = map[string]any{
				"content": []map[string]any{{"type": "text", "text": "legacy http echo"}},
			}
		default:
			resp["error"] = map[string]any{"code": -32601, "message": "no such method"}
		}
		_ = json.NewEncoder(w).Encode(resp)
	}))
	defer srv.Close()

	cli, err := Start("httptest", Config{URL: srv.URL})
	if err != nil {
		t.Fatalf("start: %v", err)
	}
	defer cli.Close()

	if cli.Era() != eraLegacy {
		t.Fatalf("era = %q, want legacy（discover 404 应回退握手）", cli.Era())
	}
	if cli.Version() != legacyVersion {
		t.Errorf("version = %q, want %q", cli.Version(), legacyVersion)
	}
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	got, err := cli.CallTool(ctx, "echo", `{"text":"x"}`)
	if err != nil {
		t.Fatalf("call: %v", err)
	}
	if got != "legacy http echo" {
		t.Fatalf("call result: %q", got)
	}
	if !rec.sawSession() {
		t.Error("legacy 传输的后续请求必须带上 Mcp-Session-Id")
	}
}

// TestHTTPAuthRequired：远程 MCP 要求鉴权时给出可操作的错误（用户的 westock-mcp 正是这种情况）。
func TestHTTPAuthRequired(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("WWW-Authenticate", "Bearer resource_metadata=\"https://example.com/.well-known/oauth-protected-resource\"")
		w.WriteHeader(http.StatusUnauthorized)
		_, _ = w.Write([]byte("{\"error\":\"unauthorized\"}"))
	}))
	defer srv.Close()

	_, err := Start("httptest", Config{URL: srv.URL})
	if err == nil {
		t.Fatal("401 应报错")
	}
	if !strings.Contains(err.Error(), "鉴权") {
		t.Errorf("错误信息应说明是鉴权问题：%v", err)
	}
}

// TestHTTPHeadersPassedThrough：配置里的 headers（例如 Authorization）要真的发出去。
func TestHTTPHeadersPassedThrough(t *testing.T) {
	rec := &httpRecorder{}
	srv := modernServer(rec, "json")
	defer srv.Close()

	cli, err := Start("httptest", Config{URL: srv.URL, Headers: map[string]string{"Authorization": "Bearer tok-123"}})
	if err != nil {
		t.Fatalf("start: %v", err)
	}
	defer cli.Close()

	// 配置里的 headers 要随每个请求发出（远程 MCP 常用 Authorization 传 token）。
	if !rec.every(func(q recordedReq) bool {
		return q.header.Get("Authorization") == "Bearer tok-123"
	}) {
		t.Errorf("配置的 headers 必须随每个请求发出（%d 个请求）", rec.count())
	}
}