package server

import (
	
	
	
	

	
)

// ClientSession represents an active session that can be used by MCPServer to interact with client.
type ClientSession interface {
	// Initialize marks session as fully initialized and ready for notifications
	Initialize()
	// Initialized returns if session is ready to accept notifications
	Initialized() bool
	// NotificationChannel provides a channel suitable for sending notifications to client.
	NotificationChannel() chan<- mcp.JSONRPCNotification
	// SessionID is a unique identifier used to track user session.
	SessionID() string
}

// SessionWithLogging is an extension of ClientSession that can receive log message notifications and set log level
type SessionWithLogging interface {
	ClientSession
	// SetLogLevel sets the minimum log level
	SetLogLevel(level mcp.LoggingLevel)
	// GetLogLevel retrieves the minimum log level
	GetLogLevel() mcp.LoggingLevel
}

// SessionWithTools is an extension of ClientSession that can store session-specific tool data
type SessionWithTools interface {
	ClientSession
	// GetSessionTools returns the tools specific to this session, if any
	// This method must be thread-safe for concurrent access
	GetSessionTools() map[string]ServerTool
	// SetSessionTools sets tools specific to this session
	// This method must be thread-safe for concurrent access
	SetSessionTools(tools map[string]ServerTool)
}

// SessionWithResources is an extension of ClientSession that can store session-specific resource data
type SessionWithResources interface {
	ClientSession
	// GetSessionResources returns the resources specific to this session, if any
	// This method must be thread-safe for concurrent access
	GetSessionResources() map[string]ServerResource
	// SetSessionResources sets resources specific to this session
	// This method must be thread-safe for concurrent access
	SetSessionResources(resources map[string]ServerResource)
}

// SessionWithResourceTemplates is an extension of ClientSession that can store session-specific resource template data
type SessionWithResourceTemplates interface {
	ClientSession
	// GetSessionResourceTemplates returns the resource templates specific to this session, if any
	// This method must be thread-safe for concurrent access
	GetSessionResourceTemplates() map[string]ServerResourceTemplate
	// SetSessionResourceTemplates sets resource templates specific to this session
	// This method must be thread-safe for concurrent access
	SetSessionResourceTemplates(templates map[string]ServerResourceTemplate)
}

// SessionWithClientInfo is an extension of ClientSession that can store client info
type SessionWithClientInfo interface {
	ClientSession
	// GetClientInfo returns the client information for this session
	GetClientInfo() mcp.Implementation
	// SetClientInfo sets the client information for this session
	SetClientInfo(clientInfo mcp.Implementation)
	// GetClientCapabilities returns the client capabilities for this session
	GetClientCapabilities() mcp.ClientCapabilities
	// SetClientCapabilities sets the client capabilities for this session
	SetClientCapabilities(clientCapabilities mcp.ClientCapabilities)
}

// SessionWithElicitation is an extension of ClientSession that can send elicitation requests
type SessionWithElicitation interface {
	ClientSession
	// RequestElicitation sends an elicitation request to the client and waits for response
	RequestElicitation(ctx context.Context, request mcp.ElicitationRequest) (*mcp.ElicitationResult, error)
}

// SessionWithRoots is an extension of ClientSession that can send list roots requests
type SessionWithRoots interface {
	ClientSession
	// ListRoots sends an list roots request to the client and waits for response
	ListRoots(ctx context.Context, request mcp.ListRootsRequest) (*mcp.ListRootsResult, error)
}

// SessionWithStreamableHTTPConfig extends ClientSession to support streamable HTTP transport configurations
type SessionWithStreamableHTTPConfig interface {
	ClientSession
	// UpgradeToSSEWhenReceiveNotification upgrades the client-server communication to SSE stream when the server
	// sends notifications to the client
	//
	// The protocol specification:
	// - If the server response contains any JSON-RPC notifications, it MUST either:
	//   - Return Content-Type: text/event-stream to initiate an SSE stream, OR
	//   - Return Content-Type: application/json for a single JSON object
	// - The client MUST support both response types.
	//
	// Reference: https://modelcontextprotocol.io/specification/2025-03-26/basic/transports#sending-messages-to-the-server
	UpgradeToSSEWhenReceiveNotification()
}

// clientSessionKey is the context key for storing current client notification channel.
type clientSessionKey struct{}

// ClientSessionFromContext retrieves current client notification context from context.
func ( context.Context) ClientSession {
	if ,  := .Value(clientSessionKey{}).(ClientSession);  {
		return 
	}
	return nil
}

// WithContext sets the current client session and returns the provided context
func ( *MCPServer) (
	 context.Context,
	 ClientSession,
) context.Context {
	return context.WithValue(, clientSessionKey{}, )
}

// RegisterSession saves session that should be notified in case if some server attributes changed.
func ( *MCPServer) (
	 context.Context,
	 ClientSession,
) error {
	 := .SessionID()
	if ,  := .sessions.LoadOrStore(, );  {
		return ErrSessionExists
	}
	.hooks.RegisterSession(, )
	return nil
}

func ( *MCPServer) ( mcp.LoggingMessageNotification) mcp.JSONRPCNotification {
	return mcp.JSONRPCNotification{
		JSONRPC: mcp.JSONRPC_VERSION,
		Notification: mcp.Notification{
			Method: .Method,
			Params: mcp.NotificationParams{
				AdditionalFields: map[string]any{
					"level":  .Params.Level,
					"logger": .Params.Logger,
					"data":   .Params.Data,
				},
			},
		},
	}
}

func ( *MCPServer) ( context.Context,  mcp.LoggingMessageNotification) error {
	 := ClientSessionFromContext()
	if  == nil || !.Initialized() {
		return ErrNotificationNotInitialized
	}
	,  := .(SessionWithLogging)
	if ! {
		return ErrSessionDoesNotSupportLogging
	}
	if !.Params.Level.ShouldSendTo(.GetLogLevel()) {
		return nil
	}
	return .sendNotificationCore(, , .buildLogNotification())
}

func ( *MCPServer) ( mcp.JSONRPCNotification) {
	.sessions.Range(func(,  any) bool {
		if ,  := .(ClientSession);  && .Initialized() {
			if ,  := .(SessionWithStreamableHTTPConfig);  {
				.UpgradeToSSEWhenReceiveNotification()
			}
			select {
			case .NotificationChannel() <- :
				// Successfully sent notification
			default:
				// Channel is blocked, if there's an error hook, use it
				if .hooks != nil && len(.hooks.OnError) > 0 {
					 := ErrNotificationChannelBlocked
					// Copy hooks pointer to local variable to avoid race condition
					 := .hooks
					go func( string,  *Hooks) {
						 := context.Background()
						// Use the error hook to report the blocked channel
						.onError(, nil, "notification", map[string]any{
							"method":    .Method,
							"sessionID": ,
						}, fmt.Errorf("notification channel blocked for session %s: %w", , ))
					}(.SessionID(), )
				}
			}
		}
		return true
	})
}

func ( *MCPServer) ( ClientSession,  mcp.JSONRPCNotification) error {
	// upgrades the client-server communication to SSE stream when the server sends notifications to the client
	if ,  := .(SessionWithStreamableHTTPConfig);  {
		.UpgradeToSSEWhenReceiveNotification()
	}
	select {
	case .NotificationChannel() <- :
		return nil
	default:
		// Channel is blocked, if there's an error hook, use it
		if .hooks != nil && len(.hooks.OnError) > 0 {
			 := ErrNotificationChannelBlocked
			 := context.Background()
			// Copy hooks pointer to local variable to avoid race condition
			 := .hooks
			go func( string,  *Hooks) {
				// Use the error hook to report the blocked channel
				.onError(, nil, "notification", map[string]any{
					"method":    .Method,
					"sessionID": ,
				}, fmt.Errorf("notification channel blocked for session %s: %w", , ))
			}(.SessionID(), )
		}
		return ErrNotificationChannelBlocked
	}
}

func ( *MCPServer) ( string,  mcp.LoggingMessageNotification) error {
	,  := .sessions.Load()
	if ! {
		return ErrSessionNotFound
	}
	,  := .(ClientSession)
	if ! || !.Initialized() {
		return ErrSessionNotInitialized
	}
	,  := .(SessionWithLogging)
	if ! {
		return ErrSessionDoesNotSupportLogging
	}
	if !.Params.Level.ShouldSendTo(.GetLogLevel()) {
		return nil
	}
	return .sendNotificationToSpecificClient(, .buildLogNotification())
}

// UnregisterSession removes from storage session that is shut down.
func ( *MCPServer) (
	 context.Context,
	 string,
) {
	,  := .sessions.LoadAndDelete()
	if ! {
		return
	}
	if ,  := .(ClientSession);  {
		.hooks.UnregisterSession(, )
	}
}

// SendNotificationToAllClients sends a notification to all the currently active clients.
func ( *MCPServer) (
	 string,
	 map[string]any,
) {
	 := mcp.JSONRPCNotification{
		JSONRPC: mcp.JSONRPC_VERSION,
		Notification: mcp.Notification{
			Method: ,
			Params: mcp.NotificationParams{
				AdditionalFields: ,
			},
		},
	}
	.sendNotificationToAllClients()
}

// SendNotificationToClient sends a notification to the current client
func ( *MCPServer) (
	 context.Context,
	 ClientSession,
	 mcp.JSONRPCNotification,
) error {
	// upgrades the client-server communication to SSE stream when the server sends notifications to the client
	if ,  := .(SessionWithStreamableHTTPConfig);  {
		.UpgradeToSSEWhenReceiveNotification()
	}
	select {
	case .NotificationChannel() <- :
		return nil
	default:
		// Channel is blocked, if there's an error hook, use it
		if .hooks != nil && len(.hooks.OnError) > 0 {
			 := .Method
			 := ErrNotificationChannelBlocked
			// Copy hooks pointer to local variable to avoid race condition
			 := .hooks
			go func( string,  *Hooks) {
				// Use the error hook to report the blocked channel
				.onError(, nil, "notification", map[string]any{
					"method":    ,
					"sessionID": ,
				}, fmt.Errorf("notification channel blocked for session %s: %w", , ))
			}(.SessionID(), )
		}
		return ErrNotificationChannelBlocked
	}
}

// SendNotificationToClient sends a notification to the current client
func ( *MCPServer) (
	 context.Context,
	 string,
	 map[string]any,
) error {
	 := ClientSessionFromContext()
	if  == nil || !.Initialized() {
		return ErrNotificationNotInitialized
	}
	 := mcp.JSONRPCNotification{
		JSONRPC: mcp.JSONRPC_VERSION,
		Notification: mcp.Notification{
			Method: ,
			Params: mcp.NotificationParams{
				AdditionalFields: ,
			},
		},
	}
	return .sendNotificationCore(, , )
}

// SendNotificationToSpecificClient sends a notification to a specific client by session ID
func ( *MCPServer) (
	 string,
	 string,
	 map[string]any,
) error {
	,  := .sessions.Load()
	if ! {
		return ErrSessionNotFound
	}
	,  := .(ClientSession)
	if ! || !.Initialized() {
		return ErrSessionNotInitialized
	}
	 := mcp.JSONRPCNotification{
		JSONRPC: mcp.JSONRPC_VERSION,
		Notification: mcp.Notification{
			Method: ,
			Params: mcp.NotificationParams{
				AdditionalFields: ,
			},
		},
	}
	return .sendNotificationToSpecificClient(, )
}

// AddSessionTool adds a tool for a specific session
func ( *MCPServer) ( string,  mcp.Tool,  ToolHandlerFunc) error {
	return .AddSessionTools(, ServerTool{Tool: , Handler: })
}

// AddSessionTools adds tools for a specific session
func ( *MCPServer) ( string,  ...ServerTool) error {
	,  := .sessions.Load()
	if ! {
		return ErrSessionNotFound
	}

	,  := .(SessionWithTools)
	if ! {
		return ErrSessionDoesNotSupportTools
	}

	.implicitlyRegisterToolCapabilities()

	// Get existing tools (this should return a thread-safe copy)
	 := .GetSessionTools()

	// Create a new map to avoid concurrent modification issues
	 := make(map[string]ServerTool, len()+len())

	// Copy existing tools
	maps.Copy(, )

	// Add new tools
	for ,  := range  {
		[.Tool.Name] = 
	}

	// Set the tools (this should be thread-safe)
	.SetSessionTools()

	// It only makes sense to send tool notifications to initialized sessions --
	// if we're not initialized yet the client can't possibly have sent their
	// initial tools/list message.
	//
	// For initialized sessions, honor tools.listChanged, which is specifically
	// about whether notifications will be sent or not.
	// see <https://modelcontextprotocol.io/specification/2025-03-26/server/tools#capabilities>
	if .Initialized() && .capabilities.tools != nil && .capabilities.tools.listChanged {
		// Send notification only to this session
		if  := .SendNotificationToSpecificClient(, "notifications/tools/list_changed", nil);  != nil {
			// Log the error but don't fail the operation
			// The tools were successfully added, but notification failed
			if .hooks != nil && len(.hooks.OnError) > 0 {
				 := .hooks
				go func( string,  *Hooks) {
					 := context.Background()
					.onError(, nil, "notification", map[string]any{
						"method":    "notifications/tools/list_changed",
						"sessionID": ,
					}, fmt.Errorf("failed to send notification after adding tools: %w", ))
				}(, )
			}
		}
	}

	return nil
}

// DeleteSessionTools removes tools from a specific session
func ( *MCPServer) ( string,  ...string) error {
	,  := .sessions.Load()
	if ! {
		return ErrSessionNotFound
	}

	,  := .(SessionWithTools)
	if ! {
		return ErrSessionDoesNotSupportTools
	}

	// Get existing tools (this should return a thread-safe copy)
	 := .GetSessionTools()
	if  == nil {
		return nil
	}

	// Create a new map to avoid concurrent modification issues
	 := make(map[string]ServerTool, len())

	// Copy existing tools except those being deleted
	maps.Copy(, )

	// Remove specified tools
	for ,  := range  {
		delete(, )
	}

	// Set the tools (this should be thread-safe)
	.SetSessionTools()

	// It only makes sense to send tool notifications to initialized sessions --
	// if we're not initialized yet the client can't possibly have sent their
	// initial tools/list message.
	//
	// For initialized sessions, honor tools.listChanged, which is specifically
	// about whether notifications will be sent or not.
	// see <https://modelcontextprotocol.io/specification/2025-03-26/server/tools#capabilities>
	if .Initialized() && .capabilities.tools != nil && .capabilities.tools.listChanged {
		// Send notification only to this session
		if  := .SendNotificationToSpecificClient(, "notifications/tools/list_changed", nil);  != nil {
			// Log the error but don't fail the operation
			// The tools were successfully deleted, but notification failed
			if .hooks != nil && len(.hooks.OnError) > 0 {
				 := .hooks
				go func( string,  *Hooks) {
					 := context.Background()
					.onError(, nil, "notification", map[string]any{
						"method":    "notifications/tools/list_changed",
						"sessionID": ,
					}, fmt.Errorf("failed to send notification after deleting tools: %w", ))
				}(, )
			}
		}
	}

	return nil
}

// AddSessionResource adds a resource for a specific session
func ( *MCPServer) ( string,  mcp.Resource,  ResourceHandlerFunc) error {
	return .AddSessionResources(, ServerResource{Resource: , Handler: })
}

// AddSessionResources adds resources for a specific session
func ( *MCPServer) ( string,  ...ServerResource) error {
	,  := .sessions.Load()
	if ! {
		return ErrSessionNotFound
	}

	,  := .(SessionWithResources)
	if ! {
		return ErrSessionDoesNotSupportResources
	}

	// For session resources, we want listChanged enabled by default
	.implicitlyRegisterCapabilities(
		func() bool { return .capabilities.resources != nil },
		func() { .capabilities.resources = &resourceCapabilities{listChanged: true} },
	)

	// Get existing resources (this should return a thread-safe copy)
	 := .GetSessionResources()

	// Create a new map to avoid concurrent modification issues
	 := make(map[string]ServerResource, len()+len())

	// Copy existing resources
	maps.Copy(, )

	// Add new resources with validation
	for ,  := range  {
		// Validate that URI is non-empty
		if .Resource.URI == "" {
			return fmt.Errorf("resource URI cannot be empty")
		}

		// Validate that URI conforms to RFC 3986
		if ,  := url.ParseRequestURI(.Resource.URI);  != nil {
			return fmt.Errorf("invalid resource URI: %w", )
		}

		[.Resource.URI] = 
	}

	// Set the resources (this should be thread-safe)
	.SetSessionResources()

	// It only makes sense to send resource notifications to initialized sessions --
	// if we're not initialized yet the client can't possibly have sent their
	// initial resources/list message.
	//
	// For initialized sessions, honor resources.listChanged, which is specifically
	// about whether notifications will be sent or not.
	// see <https://modelcontextprotocol.io/specification/2025-03-26/server/resources#capabilities>
	if .Initialized() && .capabilities.resources != nil && .capabilities.resources.listChanged {
		// Send notification only to this session
		if  := .SendNotificationToSpecificClient(, "notifications/resources/list_changed", nil);  != nil {
			// Log the error but don't fail the operation
			// The resources were successfully added, but notification failed
			if .hooks != nil && len(.hooks.OnError) > 0 {
				 := .hooks
				go func( string,  *Hooks) {
					 := context.Background()
					.onError(, nil, "notification", map[string]any{
						"method":    "notifications/resources/list_changed",
						"sessionID": ,
					}, fmt.Errorf("failed to send notification after adding resources: %w", ))
				}(, )
			}
		}
	}

	return nil
}

// DeleteSessionResources removes resources from a specific session
func ( *MCPServer) ( string,  ...string) error {
	,  := .sessions.Load()
	if ! {
		return ErrSessionNotFound
	}

	,  := .(SessionWithResources)
	if ! {
		return ErrSessionDoesNotSupportResources
	}

	// Get existing resources (this should return a thread-safe copy)
	 := .GetSessionResources()
	if  == nil {
		return nil
	}

	// Create a new map to avoid concurrent modification issues
	 := make(map[string]ServerResource, len())

	// Copy existing resources except those being deleted
	maps.Copy(, )

	// Remove specified resources and track if anything was actually deleted
	 := false
	for ,  := range  {
		if ,  := [];  {
			delete(, )
			 = true
		}
	}

	// Skip no-op write if nothing was actually deleted
	if ! {
		return nil
	}

	// Set the resources (this should be thread-safe)
	.SetSessionResources()

	// It only makes sense to send resource notifications to initialized sessions --
	// if we're not initialized yet the client can't possibly have sent their
	// initial resources/list message.
	//
	// For initialized sessions, honor resources.listChanged, which is specifically
	// about whether notifications will be sent or not.
	// see <https://modelcontextprotocol.io/specification/2025-03-26/server/resources#capabilities>
	// Only send notification if something was actually deleted
	if  && .Initialized() && .capabilities.resources != nil && .capabilities.resources.listChanged {
		// Send notification only to this session
		if  := .SendNotificationToSpecificClient(, "notifications/resources/list_changed", nil);  != nil {
			// Log the error but don't fail the operation
			// The resources were successfully deleted, but notification failed
			if .hooks != nil && len(.hooks.OnError) > 0 {
				 := .hooks
				go func( string,  *Hooks) {
					 := context.Background()
					.onError(, nil, "notification", map[string]any{
						"method":    "notifications/resources/list_changed",
						"sessionID": ,
					}, fmt.Errorf("failed to send notification after deleting resources: %w", ))
				}(, )
			}
		}
	}

	return nil
}

// AddSessionResourceTemplate adds a resource template for a specific session
func ( *MCPServer) ( string,  mcp.ResourceTemplate,  ResourceTemplateHandlerFunc) error {
	return .AddSessionResourceTemplates(, ServerResourceTemplate{
		Template: ,
		Handler:  ,
	})
}

// AddSessionResourceTemplates adds resource templates for a specific session
func ( *MCPServer) ( string,  ...ServerResourceTemplate) error {
	,  := .sessions.Load()
	if ! {
		return ErrSessionNotFound
	}

	,  := .(SessionWithResourceTemplates)
	if ! {
		return ErrSessionDoesNotSupportResourceTemplates
	}

	// For session resource templates, enable listChanged by default
	// This is the same behavior as session resources
	.implicitlyRegisterCapabilities(
		func() bool { return .capabilities.resources != nil },
		func() { .capabilities.resources = &resourceCapabilities{listChanged: true} },
	)

	// Get existing templates (this returns a thread-safe copy)
	 := .GetSessionResourceTemplates()

	// Create a new map to avoid modifying the returned copy
	 := make(map[string]ServerResourceTemplate, len()+len())

	// Copy existing templates
	maps.Copy(, )

	// Validate and add new templates
	for ,  := range  {
		if .Template.URITemplate == nil {
			return fmt.Errorf("resource template URITemplate cannot be nil")
		}
		 := .Template.URITemplate.Raw()
		if  == "" {
			return fmt.Errorf("resource template URITemplate cannot be empty")
		}
		if .Template.Name == "" {
			return fmt.Errorf("resource template name cannot be empty")
		}
		[] = 
	}

	// Set the new templates (this method must handle thread-safety)
	.SetSessionResourceTemplates()

	// Send notification if the session is initialized and listChanged is enabled
	if .Initialized() && .capabilities.resources != nil && .capabilities.resources.listChanged {
		if  := .SendNotificationToSpecificClient(, "notifications/resources/list_changed", nil);  != nil {
			// Log the error but don't fail the operation
			if .hooks != nil && len(.hooks.OnError) > 0 {
				 := .hooks
				go func( string,  *Hooks) {
					 := context.Background()
					.onError(, nil, "notification", map[string]any{
						"method":    "notifications/resources/list_changed",
						"sessionID": ,
					}, fmt.Errorf("failed to send notification after adding resource templates: %w", ))
				}(, )
			}
		}
	}

	return nil
}

// DeleteSessionResourceTemplates removes resource templates from a specific session
func ( *MCPServer) ( string,  ...string) error {
	,  := .sessions.Load()
	if ! {
		return ErrSessionNotFound
	}

	,  := .(SessionWithResourceTemplates)
	if ! {
		return ErrSessionDoesNotSupportResourceTemplates
	}

	// Get existing templates (this returns a thread-safe copy)
	 := .GetSessionResourceTemplates()

	// Track if any were actually deleted
	 := false

	// Create a new map without the deleted templates
	 := make(map[string]ServerResourceTemplate, len())
	maps.Copy(, )

	// Delete specified templates
	for ,  := range  {
		if ,  := [];  {
			delete(, )
			 = true
		}
	}

	// Only update if something was actually deleted
	if  {
		// Set the new templates (this method must handle thread-safety)
		.SetSessionResourceTemplates()

		// Send notification if the session is initialized and listChanged is enabled
		if .Initialized() && .capabilities.resources != nil && .capabilities.resources.listChanged {
			if  := .SendNotificationToSpecificClient(, "notifications/resources/list_changed", nil);  != nil {
				// Log the error but don't fail the operation
				if .hooks != nil && len(.hooks.OnError) > 0 {
					 := .hooks
					go func( string,  *Hooks) {
						 := context.Background()
						.onError(, nil, "notification", map[string]any{
							"method":    "notifications/resources/list_changed",
							"sessionID": ,
						}, fmt.Errorf("failed to send notification after deleting resource templates: %w", ))
					}(, )
				}
			}
		}
	}

	return nil
}