package server

import (
	
	

	
)

// EnableSampling enables sampling capabilities for the server.
// This allows the server to send sampling requests to clients that support it.
func ( *MCPServer) () {
	.capabilitiesMu.Lock()
	defer .capabilitiesMu.Unlock()

	 := true
	.capabilities.sampling = &
}

// RequestSampling sends a sampling request to the client.
// The client must have declared sampling capability during initialization.
func ( *MCPServer) ( context.Context,  mcp.CreateMessageRequest) (*mcp.CreateMessageResult, error) {
	 := ClientSessionFromContext()
	if  == nil {
		return nil, fmt.Errorf("no active session")
	}

	// Check if the session supports sampling requests
	if ,  := .(SessionWithSampling);  {
		return .RequestSampling(, )
	}

	// Check for inprocess sampling handler in context
	if  := InProcessSamplingHandlerFromContext();  != nil {
		return .CreateMessage(, )
	}

	return nil, fmt.Errorf("session does not support sampling")
}

// SessionWithSampling extends ClientSession to support sampling requests.
type SessionWithSampling interface {
	ClientSession
	RequestSampling(ctx context.Context, request mcp.CreateMessageRequest) (*mcp.CreateMessageResult, error)
}

// inProcessSamplingHandlerKey is the context key for storing inprocess sampling handler
type inProcessSamplingHandlerKey struct{}

// WithInProcessSamplingHandler adds a sampling handler to the context for inprocess clients
func ( context.Context,  SamplingHandler) context.Context {
	return context.WithValue(, inProcessSamplingHandlerKey{}, )
}

// InProcessSamplingHandlerFromContext retrieves the inprocess sampling handler from context
func ( context.Context) SamplingHandler {
	if ,  := .Value(inProcessSamplingHandlerKey{}).(SamplingHandler);  {
		return 
	}
	return nil
}