{-# LANGUAGE OverloadedStrings #-}

module Network.MCP.Client.Request where

import System.IO (hFlush)
import Control.Concurrent.MVar
import Control.Exception (throw)
import Data.Aeson
import Network.MCP.Client.Types
import qualified Data.Aeson.KeyMap as KeyMap
import qualified Data.ByteString.Char8 as C8
import qualified Data.ByteString.Lazy.Char8 as BL
import qualified Data.Text as T

-- Validate JSON-RPC request structure
validateJsonRpcRequest :: Value -> Either McpClientError ()
validateJsonRpcRequest :: Value -> Either McpClientError ()
validateJsonRpcRequest Value
v = case Value
v of
    Object Object
obj ->
        let hasMethod :: Bool
hasMethod = Key
"method" Key -> [Key] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` Object -> [Key]
forall {b}. KeyMap b -> [Key]
keys Object
obj
            hasJsonRpc :: Bool
hasJsonRpc = Key -> Object -> Maybe Value
forall v. Key -> KeyMap v -> Maybe v
KeyMap.lookup Key
"jsonrpc" Object
obj Maybe Value -> Maybe Value -> Bool
forall a. Eq a => a -> a -> Bool
== Value -> Maybe Value
forall a. a -> Maybe a
Just (Text -> Value
String Text
"2.0")
        in if Bool
hasMethod Bool -> Bool -> Bool
&& Bool
hasJsonRpc
           then () -> Either McpClientError ()
forall a b. b -> Either a b
Right ()
           else McpClientError -> Either McpClientError ()
forall a b. a -> Either a b
Left (McpClientError -> Either McpClientError ())
-> McpClientError -> Either McpClientError ()
forall a b. (a -> b) -> a -> b
$ Text -> McpClientError
ProtocolError Text
"Invalid JSON-RPC request structure"
    Value
_ -> McpClientError -> Either McpClientError ()
forall a b. a -> Either a b
Left (McpClientError -> Either McpClientError ())
-> McpClientError -> Either McpClientError ()
forall a b. (a -> b) -> a -> b
$ Text -> McpClientError
ProtocolError Text
"Request must be a JSON object"
  where
    keys :: KeyMap b -> [Key]
keys = ((Key, b) -> Key) -> [(Key, b)] -> [Key]
forall a b. (a -> b) -> [a] -> [b]
map (Key, b) -> Key
forall a b. (a, b) -> a
fst ([(Key, b)] -> [Key])
-> (KeyMap b -> [(Key, b)]) -> KeyMap b -> [Key]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KeyMap b -> [(Key, b)]
forall v. KeyMap v -> [(Key, v)]
KeyMap.toList

-- Robust request sending and response parsing
sendClientRequest
    :: Client      -- stdin handle
    -> Value       -- request message
    -> IO Value    -- parsed result
sendClientRequest :: Client -> Value -> IO Value
sendClientRequest Client
client Value
message = do
    -- Validate request structure before sending
    case Value -> Either McpClientError ()
validateJsonRpcRequest Value
message of
        Left McpClientError
err -> McpClientError -> IO ()
forall a e. Exception e => e -> a
throw McpClientError
err
        Right () -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

    Handle
hstdin <- MVar (Maybe Handle) -> IO (Maybe Handle)
forall a. MVar a -> IO a
readMVar (Client -> MVar (Maybe Handle)
clientStdin Client
client) IO (Maybe Handle) -> (Maybe Handle -> IO Handle) -> IO Handle
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Just Handle
inp -> do Handle -> IO Handle
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
return Handle
inp
        Maybe Handle
Nothing -> McpClientError -> IO Handle
forall a e. Exception e => e -> a
throw (McpClientError -> IO Handle) -> McpClientError -> IO Handle
forall a b. (a -> b) -> a -> b
$ Text -> McpClientError
ConnectionError Text
"Stdin not available"

    Handle
hstdout <- MVar (Maybe Handle) -> IO (Maybe Handle)
forall a. MVar a -> IO a
readMVar (Client -> MVar (Maybe Handle)
clientStdout Client
client) IO (Maybe Handle) -> (Maybe Handle -> IO Handle) -> IO Handle
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Just Handle
outp -> do Handle -> IO Handle
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
return Handle
outp
        Maybe Handle
Nothing -> McpClientError -> IO Handle
forall a e. Exception e => e -> a
throw (McpClientError -> IO Handle) -> McpClientError -> IO Handle
forall a b. (a -> b) -> a -> b
$ Text -> McpClientError
ConnectionError Text
"Stdout not available"

    -- Encode and send message
    let a :: ByteString
a = ByteString -> ByteString
BL.toStrict (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ Value -> ByteString
forall a. ToJSON a => a -> ByteString
encode Value
message
    Handle -> ByteString -> IO ()
C8.hPut Handle
hstdin ByteString
a

    Handle -> IO ()
hFlush Handle
hstdin

    -- Read and parse response
    ByteString
responseStr <- Handle -> IO ByteString
C8.hGetLine Handle
hstdout
    let responseBS :: ByteString
responseBS = ByteString -> ByteString
BL.fromStrict ByteString
responseStr

    -- Comprehensive response parsing
    case ByteString -> Either String Value
forall a. FromJSON a => ByteString -> Either String a
eitherDecode ByteString
responseBS of
        Left String
decodeErr ->
            McpClientError -> IO Value
forall a e. Exception e => e -> a
throw (McpClientError -> IO Value) -> McpClientError -> IO Value
forall a b. (a -> b) -> a -> b
$ Text -> McpClientError
ProtocolError (Text -> McpClientError) -> Text -> McpClientError
forall a b. (a -> b) -> a -> b
$ String -> Text
T.pack String
decodeErr
        Right (Object Object
resp) ->
            Object -> IO Value
parseResponse Object
resp
        Right Value
_ ->
            McpClientError -> IO Value
forall a e. Exception e => e -> a
throw (McpClientError -> IO Value) -> McpClientError -> IO Value
forall a b. (a -> b) -> a -> b
$ Text -> McpClientError
ProtocolError Text
"Response must be a JSON object"

-- Parse the response, handling both success and error cases
parseResponse :: KeyMap.KeyMap Value -> IO Value
parseResponse :: Object -> IO Value
parseResponse Object
resp =
    case (Key -> Object -> Maybe Value
forall v. Key -> KeyMap v -> Maybe v
KeyMap.lookup Key
"result" Object
resp, Key -> Object -> Maybe Value
forall v. Key -> KeyMap v -> Maybe v
KeyMap.lookup Key
"error" Object
resp) of
        (Just Value
result, Maybe Value
Nothing) ->
            Value -> IO Value
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Value
result
        (Maybe Value
Nothing, Just (Object Object
errorObj)) -> do
            let code :: Maybe Int
code = case Key -> Object -> Maybe Value
forall v. Key -> KeyMap v -> Maybe v
KeyMap.lookup Key
"code" Object
errorObj of
                         Just (Number Scientific
n) -> Int -> Maybe Int
forall a. a -> Maybe a
Just (Scientific -> Int
forall b. Integral b => Scientific -> b
forall a b. (RealFrac a, Integral b) => a -> b
round Scientific
n)
                         Maybe Value
_ -> Maybe Int
forall a. Maybe a
Nothing
                message :: Text
message = case Key -> Object -> Maybe Value
forall v. Key -> KeyMap v -> Maybe v
KeyMap.lookup Key
"message" Object
errorObj of
                            Just (String Text
msg) -> Text
msg
                            Maybe Value
_ -> Text
"Unknown server error"
                details :: Maybe Value
details = Key -> Object -> Maybe Value
forall v. Key -> KeyMap v -> Maybe v
KeyMap.lookup Key
"data" Object
errorObj

            McpClientError -> IO Value
forall a e. Exception e => e -> a
throw (McpClientError -> IO Value) -> McpClientError -> IO Value
forall a b. (a -> b) -> a -> b
$ ServerError
                { $sel:serverErrorCode:ConnectionError :: Maybe Int
serverErrorCode = Maybe Int
code
                , $sel:serverErrorMessage:ConnectionError :: Text
serverErrorMessage = Text
message
                , $sel:serverErrorData:ConnectionError :: Maybe Value
serverErrorData = Maybe Value
details
                }
        (Maybe Value
Nothing, Maybe Value
Nothing) ->
            McpClientError -> IO Value
forall a e. Exception e => e -> a
throw (McpClientError -> IO Value) -> McpClientError -> IO Value
forall a b. (a -> b) -> a -> b
$ Text -> McpClientError
ProtocolError Text
"Response lacks both result and error"
        (Maybe Value, Maybe Value)
_ ->
            McpClientError -> IO Value
forall a e. Exception e => e -> a
throw (McpClientError -> IO Value) -> McpClientError -> IO Value
forall a b. (a -> b) -> a -> b
$ Text -> McpClientError
ProtocolError Text
"Invalid response structure"