package mcp

import (
	"net/http"

	"github.com/go-chi/chi/v5"
	mcpsdk "github.com/modelcontextprotocol/go-sdk/mcp"
)

// MountTransports 在 chi router 上挂载 MCP transport 端点：
//   - POST/GET /mcp  Streamable HTTP（2025-03-26+，主推）
//   - POST/GET /mcp/sse  legacy HTTP+SSE（2024-11-05，SSE handler 自行根据 ?sessionid 分派）
//
// 中间件链：Auth → RateLimit → injectTokenID → SDK handler。
func MountTransports(
	r chi.Router,
	srv *mcpsdk.Server,
	ts TokenService,
	rl *RateLimiter,
) {
	streamable := mcpsdk.NewStreamableHTTPHandler(func(*http.Request) *mcpsdk.Server {
		return srv
	}, nil)
	sse := mcpsdk.NewSSEHandler(func(*http.Request) *mcpsdk.Server {
		return srv
	}, nil)

	// 嵌套 group 让中间件作用范围限定到 /mcp/* 子树
	r.Group(func(r chi.Router) {
		r.Use(NewAuthMiddleware(ts))
		r.Use(NewRateLimitMiddleware(rl))
		r.Use(injectTokenIDMiddleware())

		// Streamable HTTP
		r.Method(http.MethodPost, "/mcp", streamable)
		r.Method(http.MethodGet, "/mcp", streamable)

		// Legacy HTTP+SSE（SSE handler 自行分派 GET=stream / POST=message via ?sessionid）
		r.Method(http.MethodGet, "/mcp/sse", sse)
		r.Method(http.MethodPost, "/mcp/sse", sse)
	})
}

// injectTokenIDMiddleware 在 auth middleware 之后跑，把 AuthSession.TokenID
// 复制到 ctx 的 mcpTokenIDCtxKey，供 AuditWrap 在 audit metadata 里加 via_mcp 标记。
func injectTokenIDMiddleware() func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			sess, ok := FromContext(r.Context())
			if !ok {
				next.ServeHTTP(w, r)
				return
			}
			ctx := withMcpTokenID(r.Context(), sess.TokenID)
			next.ServeHTTP(w, r.WithContext(ctx))
		})
	}
}