package mcp

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"strings"

	"github.com/gin-gonic/gin"
	"github.com/tidbcloud/AgentTeam/desktop/internal/messagesend"
	"github.com/tidbcloud/AgentTeam/desktop/internal/token"
)

const mcpSessionHeader = "Mcp-Session-Id"

type requestHeadersContextKey struct{}

// requestHeader returns the inbound MCP HTTP request header value attached to
// the current JSON-RPC dispatch context.
func requestHeader(ctx context.Context, name string) string {
	headers, ok := ctx.Value(requestHeadersContextKey{}).(http.Header)
	if !ok {
		return ""
	}
	return headers.Get(name)
}

func mcpProtectedResourcePath(basePath string) string {
	base := strings.TrimSpace(basePath)
	base = strings.TrimSuffix(base, "/")
	if base == "" {
		return oauthProtectedResourcePath
	}
	return base + oauthProtectedResourcePath
}

func setProtectedResourceChallenge(c *gin.Context, basePath string) {
	c.Header("WWW-Authenticate", fmt.Sprintf(`Bearer resource_metadata=%q`, mcpProtectedResourcePath(basePath)))
}

func protectedResourceBasePathFromWildcard(c *gin.Context) string {
	path := strings.TrimSpace(c.Param("resourcePath"))
	if path == "" || path == "/" {
		return ""
	}
	if !strings.HasPrefix(path, "/") {
		path = "/" + path
	}
	return path
}

func protectedResourceBasePathFromSuffix(c *gin.Context) string {
	base := strings.TrimSuffix(strings.TrimSpace(c.Request.URL.Path), oauthProtectedResourcePath)
	base = strings.TrimSuffix(base, "/")
	if base == "" {
		return ""
	}
	return base
}

func writeProtectedResourceMetadata(c *gin.Context, basePath string) {
	base := strings.TrimSpace(basePath)
	if base == "" {
		_ = writeMCPJSON(c, http.StatusNotFound, gin.H{
			"error":   "not_found",
			"message": "oauth-protected-resource requires a concrete MCP resource path",
		})
		return
	}

	_ = writeMCPJSON(c, http.StatusOK, gin.H{
		"resource":                 base,
		"authorization_servers":    []string{},
		"bearer_methods_supported": []string{"header"},
	})
}

// ---------------------------------------------------------------------------
// AuthMiddleware
// ---------------------------------------------------------------------------

// AccessControl returns a gin middleware that validates Bearer tokens against
// the default token manager and sets the agentID in the context.
func AccessControl() gin.HandlerFunc {
	return accessControlWithManagerForTest(token.Default())
}

// accessControlWithManagerForTest returns a gin middleware that validates Bearer
// tokens against manager. Server-scoped MCP routes are checked against the
// route's server binding; generic MCP routes only require the token's agent
// binding.
func accessControlWithManagerForTest(manager *token.KeyManager) gin.HandlerFunc {
	return accessControlWithManager(manager, func(manager *token.KeyManager, agentID, serverID, tok string) (token.AgentTokenInfo, error) {
		if serverID != "" {
			return manager.ValidateAgentTokenForServerInfo(agentID, serverID, tok)
		}
		return manager.ValidateAgentTokenInfo(agentID, tok)
	})
}

// leaderAccessControlWithManagerForTest returns leader-only MCP middleware using
// manager. Memory uses this instead of generic AccessControl so an isolated
// execution token cannot inherit the leader's local Memory capability.
func leaderAccessControlWithManagerForTest(manager *token.KeyManager) gin.HandlerFunc {
	return accessControlWithManager(manager, func(manager *token.KeyManager, agentID, serverID, tok string) (token.AgentTokenInfo, error) {
		return manager.ValidateUnscopedAgentTokenInfo(agentID, tok)
	})
}

var errMissingBearerAuthorization = errors.New("missing or malformed bearer authorization")

// parseBearerAuthorization extracts the bearer token from a strict
// "Bearer <token>" Authorization header. It rejects malformed schemes such as
// "Bearerat-dev" so MCP auth and message routing enforce one credential contract.
func parseBearerAuthorization(header string) (string, bool) {
	const prefix = "Bearer "
	if !strings.HasPrefix(header, prefix) {
		return "", false
	}
	credential := header[len(prefix):]
	if credential == "" ||
		strings.TrimSpace(credential) != credential ||
		strings.ContainsAny(credential, " \t\r\n") {
		return "", false
	}
	return credential, true
}

type agentTokenValidator func(manager *token.KeyManager, agentID, serverID, tok string) (token.AgentTokenInfo, error)

func accessControlWithManager(manager *token.KeyManager, validate agentTokenValidator) gin.HandlerFunc {
	if manager == nil {
		manager = token.Default()
	}
	return func(c *gin.Context) {
		agentID := strings.TrimSpace(c.Param("agentId"))
		if agentID == "" {
			c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{
				"error": "agent id is required",
			})
			return
		}

		tok, ok := parseBearerAuthorization(c.GetHeader("Authorization"))
		if !ok {
			abortMCPAuth(c, errMissingBearerAuthorization)
			return
		}
		serverID := strings.TrimSpace(c.Param("serverId"))
		_, err := validate(manager, agentID, serverID, tok)
		if err != nil {
			abortMCPAuth(c, err)
			return
		}
		c.Set("agentID", agentID)
		c.Header(mcpSessionHeader, "agentteam-"+agentID)
		// Bind the authenticated bearer's irreversible identity into the
		// request context so downstream tool guards can scope reads to this
		// exact credential.
		c.Request = c.Request.WithContext(messagesend.WithBearerID(c.Request.Context(), tok))
		c.Next()
	}
}

func abortMCPAuth(c *gin.Context, err error) {
	status := http.StatusUnauthorized
	msg := "invalid agent bridge token"
	if errors.Is(err, token.ErrAgentMismatch) ||
		errors.Is(err, token.ErrAgentTokenServerMismatch) ||
		errors.Is(err, token.ErrAgentTokenScopeMismatch) {
		status = http.StatusForbidden
		msg = "token does not belong to this agent"
	}
	setProtectedResourceChallenge(c, c.Request.URL.Path)
	c.AbortWithStatusJSON(status, gin.H{"error": msg})
}

// RegisterDiscoveryRoutes installs public well-known endpoints so MCP clients
// probing OAuth discovery never fall through to the SPA handler.