{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeFamilies #-}

{- |
Module      : MCP.Server
Description : MCP server core types and interface
Copyright   : (C) 2025 Matthias Pall Gissurarson
License     : MIT
Maintainer  : mpg@mpg.is
Stability   : experimental
Portability : GHC

This module provides the core types and interface for MCP server implementations.
-}
module MCP.Server (
    -- * Server Interface
    MCPServer (..),
    ServerState (..),
    ServerConfig (..),
    MCPServerM,
    runMCPServer,
    initialServerState,

    -- * Utilities
    sendResponse,
    sendNotification,
    sendError,
) where

import Control.Monad.Except (ExceptT, runExceptT)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad.Reader (ReaderT, runReaderT)
import Control.Monad.State.Strict (StateT, runStateT)
import Data.Aeson (ToJSON, encode, toJSON)
import Data.ByteString.Lazy.Char8 qualified as LBSC
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Text (Text)
import System.IO (Handle, hFlush)

import MCP.Protocol
import MCP.Types

-- | Server state tracking initialization, capabilities, and subscriptions
data ServerState = ServerState
    { ServerState -> Bool
serverInitialized :: Bool
    , ServerState -> ServerCapabilities
serverCapabilities :: ServerCapabilities
    , ServerState -> Maybe ClientCapabilities
clientCapabilities :: Maybe ClientCapabilities
    , ServerState -> Maybe Implementation
serverInfo :: Maybe Implementation
    , ServerState -> Map Text ()
subscriptions :: Map Text ()
    }
    deriving (Int -> ServerState -> ShowS
[ServerState] -> ShowS
ServerState -> String
(Int -> ServerState -> ShowS)
-> (ServerState -> String)
-> ([ServerState] -> ShowS)
-> Show ServerState
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ServerState -> ShowS
showsPrec :: Int -> ServerState -> ShowS
$cshow :: ServerState -> String
show :: ServerState -> String
$cshowList :: [ServerState] -> ShowS
showList :: [ServerState] -> ShowS
Show)

-- | Configuration for running an MCP server
data ServerConfig = ServerConfig
    { ServerConfig -> Handle
configInput :: Handle
    , ServerConfig -> Handle
configOutput :: Handle
    , ServerConfig -> Implementation
configServerInfo :: Implementation
    , ServerConfig -> ServerCapabilities
configCapabilities :: ServerCapabilities
    }
    deriving (Int -> ServerConfig -> ShowS
[ServerConfig] -> ShowS
ServerConfig -> String
(Int -> ServerConfig -> ShowS)
-> (ServerConfig -> String)
-> ([ServerConfig] -> ShowS)
-> Show ServerConfig
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ServerConfig -> ShowS
showsPrec :: Int -> ServerConfig -> ShowS
$cshow :: ServerConfig -> String
show :: ServerConfig -> String
$cshowList :: [ServerConfig] -> ShowS
showList :: [ServerConfig] -> ShowS
Show)

-- | The monad stack for MCP server operations
type MCPServerM = ReaderT ServerConfig (StateT ServerState (ExceptT Text IO))

-- | Run an MCPServerM computation with the given config and initial state
runMCPServer :: ServerConfig -> ServerState -> MCPServerM a -> IO (Either Text (a, ServerState))
runMCPServer :: forall a.
ServerConfig
-> ServerState -> MCPServerM a -> IO (Either Text (a, ServerState))
runMCPServer ServerConfig
config ServerState
state MCPServerM a
action = ExceptT Text IO (a, ServerState)
-> IO (Either Text (a, ServerState))
forall e (m :: * -> *) a. ExceptT e m a -> m (Either e a)
runExceptT (ExceptT Text IO (a, ServerState)
 -> IO (Either Text (a, ServerState)))
-> ExceptT Text IO (a, ServerState)
-> IO (Either Text (a, ServerState))
forall a b. (a -> b) -> a -> b
$ StateT ServerState (ExceptT Text IO) a
-> ServerState -> ExceptT Text IO (a, ServerState)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (MCPServerM a
-> ServerConfig -> StateT ServerState (ExceptT Text IO) a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT MCPServerM a
action ServerConfig
config) ServerState
state

{- | Create the initial server state with the given capabilities
The server starts uninitialized and must receive an 'initialize' request
before it can handle other requests.
-}
initialServerState :: ServerCapabilities -> ServerState
initialServerState :: ServerCapabilities -> ServerState
initialServerState ServerCapabilities
caps =
    ServerState
        { $sel:serverInitialized:ServerState :: Bool
serverInitialized = Bool
False
        , $sel:serverCapabilities:ServerState :: ServerCapabilities
serverCapabilities = ServerCapabilities
caps
        , $sel:clientCapabilities:ServerState :: Maybe ClientCapabilities
clientCapabilities = Maybe ClientCapabilities
forall a. Maybe a
Nothing
        , $sel:serverInfo:ServerState :: Maybe Implementation
serverInfo = Maybe Implementation
forall a. Maybe a
Nothing
        , $sel:subscriptions:ServerState :: Map Text ()
subscriptions = Map Text ()
forall k a. Map k a
Map.empty
        }

-- | Type class for implementing MCP server handlers
class (Monad m) => MCPServer m where
    handleListResources :: ListResourcesParams -> m ListResourcesResult
    handleReadResource :: ReadResourceParams -> m ReadResourceResult
    handleListResourceTemplates :: ListResourceTemplatesParams -> m ListResourceTemplatesResult
    handleListPrompts :: ListPromptsParams -> m ListPromptsResult
    handleGetPrompt :: GetPromptParams -> m GetPromptResult
    handleListTools :: ListToolsParams -> m ListToolsResult
    handleCallTool :: CallToolParams -> m CallToolResult
    handleComplete :: CompleteParams -> m CompleteResult
    handleSetLevel :: SetLevelParams -> m ()

-- | Send a JSON-RPC response
sendResponse :: (MonadIO m, ToJSON a) => Handle -> RequestId -> a -> m ()
sendResponse :: forall (m :: * -> *) a.
(MonadIO m, ToJSON a) =>
Handle -> RequestId -> a -> m ()
sendResponse Handle
handle RequestId
reqId a
result = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    let response :: JSONRPCResponse
response = Text -> RequestId -> Value -> JSONRPCResponse
JSONRPCResponse Text
"2.0" RequestId
reqId (a -> Value
forall a. ToJSON a => a -> Value
toJSON a
result)
    Handle -> ByteString -> IO ()
LBSC.hPutStrLn Handle
handle (JSONRPCResponse -> ByteString
forall a. ToJSON a => a -> ByteString
encode JSONRPCResponse
response)
    Handle -> IO ()
hFlush Handle
handle

-- | Send a JSON-RPC error response
sendError :: (MonadIO m) => Handle -> RequestId -> JSONRPCErrorInfo -> m ()
sendError :: forall (m :: * -> *).
MonadIO m =>
Handle -> RequestId -> JSONRPCErrorInfo -> m ()
sendError Handle
handle RequestId
reqId JSONRPCErrorInfo
errorInfo = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    let response :: JSONRPCError
response = Text -> RequestId -> JSONRPCErrorInfo -> JSONRPCError
JSONRPCError Text
"2.0" RequestId
reqId JSONRPCErrorInfo
errorInfo
    Handle -> ByteString -> IO ()
LBSC.hPutStrLn Handle
handle (JSONRPCError -> ByteString
forall a. ToJSON a => a -> ByteString
encode JSONRPCError
response)
    Handle -> IO ()
hFlush Handle
handle

-- | Send a JSON-RPC notification
sendNotification :: (MonadIO m, ToJSON a) => Handle -> Text -> a -> m ()
sendNotification :: forall (m :: * -> *) a.
(MonadIO m, ToJSON a) =>
Handle -> Text -> a -> m ()
sendNotification Handle
handle Text
method a
params = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    let notification :: JSONRPCNotification
notification = Text -> Text -> Maybe Value -> JSONRPCNotification
JSONRPCNotification Text
"2.0" Text
method (Value -> Maybe Value
forall a. a -> Maybe a
Just (a -> Value
forall a. ToJSON a => a -> Value
toJSON a
params))
    Handle -> ByteString -> IO ()
LBSC.hPutStrLn Handle
handle (JSONRPCNotification -> ByteString
forall a. ToJSON a => a -> ByteString
encode JSONRPCNotification
notification)
    Handle -> IO ()
hFlush Handle
handle