{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE OverloadedStrings #-}

module Language.LSP.Test.Decoding where

import Control.Exception
import Control.Lens
import Data.Aeson
import Data.Aeson.Types
import Data.ByteString.Lazy.Char8 qualified as B
import Data.Foldable
import Data.Functor.Product
import Data.Maybe
import Language.LSP.Protocol.Lens qualified as L
import Language.LSP.Protocol.Message
import Language.LSP.Test.Exceptions
import System.IO
import System.IO.Error
import Prelude hiding (id)

import Data.IxMap
import Data.Kind

{- | Fetches the next message bytes based on
 the Content-Length header
-}
getNextMessage :: Handle -> IO B.ByteString
getNextMessage :: Handle -> IO ByteString
getNextMessage Handle
h = do
  [([Char], [Char])]
headers <- Handle -> IO [([Char], [Char])]
getHeaders Handle
h
  case [Char] -> Int
forall a. Read a => [Char] -> a
read ([Char] -> Int) -> ([Char] -> [Char]) -> [Char] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> [Char]
forall a. HasCallStack => [a] -> [a]
init ([Char] -> Int) -> Maybe [Char] -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Char] -> [([Char], [Char])] -> Maybe [Char]
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup [Char]
"Content-Length" [([Char], [Char])]
headers of
    Maybe Int
Nothing -> SessionException -> IO ByteString
forall a e. Exception e => e -> a
throw SessionException
NoContentLengthHeader
    Just Int
size -> Handle -> Int -> IO ByteString
B.hGet Handle
h Int
size

addHeader :: B.ByteString -> B.ByteString
addHeader :: ByteString -> ByteString
addHeader ByteString
content =
  [ByteString] -> ByteString
B.concat
    [ ByteString
"Content-Length: "
    , [Char] -> ByteString
B.pack ([Char] -> ByteString) -> [Char] -> ByteString
forall a b. (a -> b) -> a -> b
$ Int64 -> [Char]
forall a. Show a => a -> [Char]
show (Int64 -> [Char]) -> Int64 -> [Char]
forall a b. (a -> b) -> a -> b
$ ByteString -> Int64
B.length ByteString
content
    , ByteString
"\r\n"
    , ByteString
"\r\n"
    , ByteString
content
    ]

getHeaders :: Handle -> IO [(String, String)]
getHeaders :: Handle -> IO [([Char], [Char])]
getHeaders Handle
h = do
  [Char]
l <- IO [Char] -> (IOError -> IO [Char]) -> IO [Char]
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
catch (Handle -> IO [Char]
hGetLine Handle
h) IOError -> IO [Char]
forall {a}. IOError -> a
eofHandler
  let ([Char]
name, [Char]
val) = (Char -> Bool) -> [Char] -> ([Char], [Char])
forall a. (a -> Bool) -> [a] -> ([a], [a])
span (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/= Char
':') [Char]
l
  if [Char] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Char]
val then [([Char], [Char])] -> IO [([Char], [Char])]
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return [] else (([Char]
name, Int -> [Char] -> [Char]
forall a. Int -> [a] -> [a]
drop Int
2 [Char]
val) :) ([([Char], [Char])] -> [([Char], [Char])])
-> IO [([Char], [Char])] -> IO [([Char], [Char])]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Handle -> IO [([Char], [Char])]
getHeaders Handle
h
 where
  eofHandler :: IOError -> a
eofHandler IOError
e
    | IOError -> Bool
isEOFError IOError
e = SessionException -> a
forall a e. Exception e => e -> a
throw SessionException
UnexpectedServerTermination
    | Bool
otherwise = IOError -> a
forall a e. Exception e => e -> a
throw IOError
e

type RequestMap = IxMap LspId (SMethod :: Method ClientToServer Request -> Type)

newRequestMap :: RequestMap
newRequestMap :: RequestMap
newRequestMap = RequestMap
forall {a} (k :: a -> *) (f :: a -> *). IxMap k f
emptyIxMap

updateRequestMap :: RequestMap -> LspId m -> SClientMethod m -> Maybe RequestMap
updateRequestMap :: forall (m :: Method 'ClientToServer 'Request).
RequestMap -> LspId m -> SClientMethod m -> Maybe RequestMap
updateRequestMap RequestMap
reqMap LspId m
id SClientMethod m
method = LspId m -> SClientMethod m -> RequestMap -> Maybe RequestMap
forall {a} (k :: a -> *) (m :: a) (f :: a -> *).
IxOrd k =>
k m -> f m -> IxMap k f -> Maybe (IxMap k f)
insertIxMap LspId m
id SClientMethod m
method RequestMap
reqMap

getRequestMap :: [FromClientMessage] -> RequestMap
getRequestMap :: [FromClientMessage] -> RequestMap
getRequestMap = (RequestMap -> FromClientMessage -> RequestMap)
-> RequestMap -> [FromClientMessage] -> RequestMap
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' RequestMap -> FromClientMessage -> RequestMap
helper RequestMap
forall {a} (k :: a -> *) (f :: a -> *). IxMap k f
emptyIxMap
 where
  helper :: RequestMap -> FromClientMessage -> RequestMap
  helper :: RequestMap -> FromClientMessage -> RequestMap
helper RequestMap
acc FromClientMessage
msg = case FromClientMessage
msg of
    FromClientMess SMethod m
m TMessage m
mess -> case SMethod m -> ClientNotOrReq m
forall {t :: MessageKind} (m :: Method 'ClientToServer t).
SClientMethod m -> ClientNotOrReq m
splitClientMethod SMethod m
m of
      ClientNotOrReq m
IsClientNot -> RequestMap
acc
      ClientNotOrReq m
IsClientReq -> Maybe RequestMap -> RequestMap
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe RequestMap -> RequestMap) -> Maybe RequestMap -> RequestMap
forall a b. (a -> b) -> a -> b
$ RequestMap -> LspId m -> SClientMethod m -> Maybe RequestMap
forall (m :: Method 'ClientToServer 'Request).
RequestMap -> LspId m -> SClientMethod m -> Maybe RequestMap
updateRequestMap RequestMap
acc (TMessage m
TRequestMessage m
mess TRequestMessage m
-> Getting (LspId m) (TRequestMessage m) (LspId m) -> LspId m
forall s a. s -> Getting a s a -> a
^. Getting (LspId m) (TRequestMessage m) (LspId m)
forall s a. HasId s a => Lens' s a
Lens' (TRequestMessage m) (LspId m)
L.id) SMethod m
SClientMethod m
m
      ClientNotOrReq m
IsClientEither -> case TMessage m
mess of
        NotMess TNotificationMessage ('Method_CustomMethod s)
_ -> RequestMap
acc
        ReqMess TRequestMessage ('Method_CustomMethod s)
msg -> Maybe RequestMap -> RequestMap
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe RequestMap -> RequestMap) -> Maybe RequestMap -> RequestMap
forall a b. (a -> b) -> a -> b
$ RequestMap
-> LspId ('Method_CustomMethod s)
-> SClientMethod ('Method_CustomMethod s)
-> Maybe RequestMap
forall (m :: Method 'ClientToServer 'Request).
RequestMap -> LspId m -> SClientMethod m -> Maybe RequestMap
updateRequestMap RequestMap
acc (TRequestMessage ('Method_CustomMethod s)
msg TRequestMessage ('Method_CustomMethod s)
-> Getting
     (LspId ('Method_CustomMethod s))
     (TRequestMessage ('Method_CustomMethod s))
     (LspId ('Method_CustomMethod s))
-> LspId ('Method_CustomMethod s)
forall s a. s -> Getting a s a -> a
^. Getting
  (LspId ('Method_CustomMethod s))
  (TRequestMessage ('Method_CustomMethod s))
  (LspId ('Method_CustomMethod s))
forall s a. HasId s a => Lens' s a
Lens'
  (TRequestMessage ('Method_CustomMethod s))
  (LspId ('Method_CustomMethod s))
L.id) SMethod m
SClientMethod ('Method_CustomMethod s)
m
    FromClientMessage
_ -> RequestMap
acc

decodeFromServerMsg :: RequestMap -> B.ByteString -> (RequestMap, FromServerMessage)
decodeFromServerMsg :: RequestMap -> ByteString -> (RequestMap, FromServerMessage)
decodeFromServerMsg RequestMap
reqMap ByteString
bytes = Result (FromServerMessage' (Product SMethod (Const RequestMap)))
-> (RequestMap, FromServerMessage)
unP (Result (FromServerMessage' (Product SMethod (Const RequestMap)))
 -> (RequestMap, FromServerMessage))
-> Result (FromServerMessage' (Product SMethod (Const RequestMap)))
-> (RequestMap, FromServerMessage)
forall a b. (a -> b) -> a -> b
$ (Value
 -> Parser
      (FromServerMessage' (Product SMethod (Const RequestMap))))
-> Value
-> Result (FromServerMessage' (Product SMethod (Const RequestMap)))
forall a b. (a -> Parser b) -> a -> Result b
parse Value
-> Parser (FromServerMessage' (Product SMethod (Const RequestMap)))
p Value
obj
 where
  obj :: Value
obj = Maybe Value -> Value
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Value -> Value) -> Maybe Value -> Value
forall a b. (a -> b) -> a -> b
$ ByteString -> Maybe Value
forall a. FromJSON a => ByteString -> Maybe a
decode ByteString
bytes :: Value
  p :: Value
-> Parser (FromServerMessage' (Product SMethod (Const RequestMap)))
p = LookupFunc 'ClientToServer (Product SMethod (Const RequestMap))
-> Value
-> Parser (FromServerMessage' (Product SMethod (Const RequestMap)))
forall (a :: Method 'ClientToServer 'Request -> *).
LookupFunc 'ClientToServer a
-> Value -> Parser (FromServerMessage' a)
parseServerMessage (LookupFunc 'ClientToServer (Product SMethod (Const RequestMap))
 -> Value
 -> Parser
      (FromServerMessage' (Product SMethod (Const RequestMap))))
-> LookupFunc 'ClientToServer (Product SMethod (Const RequestMap))
-> Value
-> Parser (FromServerMessage' (Product SMethod (Const RequestMap)))
forall a b. (a -> b) -> a -> b
$ \LspId m
lid ->
    let (Maybe (SMethod m)
mm, RequestMap
newMap) = LspId m -> RequestMap -> (Maybe (SMethod m), RequestMap)
forall {a} (k :: a -> *) (m :: a) (f :: a -> *).
IxOrd k =>
k m -> IxMap k f -> (Maybe (f m), IxMap k f)
pickFromIxMap LspId m
lid RequestMap
reqMap
     in case Maybe (SMethod m)
mm of
          Maybe (SMethod m)
Nothing -> Maybe (SMethod m, Product SMethod (Const RequestMap) m)
forall a. Maybe a
Nothing
          Just SMethod m
m -> (SMethod m, Product SMethod (Const RequestMap) m)
-> Maybe (SMethod m, Product SMethod (Const RequestMap) m)
forall a. a -> Maybe a
Just (SMethod m
m, SMethod m
-> Const RequestMap m -> Product SMethod (Const RequestMap) m
forall {k} (f :: k -> *) (g :: k -> *) (a :: k).
f a -> g a -> Product f g a
Pair SMethod m
m (RequestMap -> Const RequestMap m
forall {k} a (b :: k). a -> Const a b
Const RequestMap
newMap))
  unP :: Result (FromServerMessage' (Product SMethod (Const RequestMap)))
-> (RequestMap, FromServerMessage)
unP (Success (FromServerMess SMethod m
m TMessage m
msg)) = (RequestMap
reqMap, SMethod m -> TMessage m -> FromServerMessage
forall (t :: MessageKind) (m :: Method 'ServerToClient t)
       (a :: Method 'ClientToServer 'Request -> *).
SMethod m -> TMessage m -> FromServerMessage' a
FromServerMess SMethod m
m TMessage m
msg)
  unP (Success (FromServerRsp (Pair SMethod m
m (Const RequestMap
newMap)) TResponseMessage m
msg)) = (RequestMap
newMap, SMethod m -> TResponseMessage m -> FromServerMessage
forall (m :: Method 'ClientToServer 'Request)
       (a :: Method 'ClientToServer 'Request -> *).
a m -> TResponseMessage m -> FromServerMessage' a
FromServerRsp SMethod m
m TResponseMessage m
msg)
  unP (Error [Char]
e) = [Char] -> (RequestMap, FromServerMessage)
forall a. HasCallStack => [Char] -> a
error ([Char] -> (RequestMap, FromServerMessage))
-> [Char] -> (RequestMap, FromServerMessage)
forall a b. (a -> b) -> a -> b
$ [Char]
"Error decoding " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Value -> [Char]
forall a. Show a => a -> [Char]
show Value
obj [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
" :" [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
e

{-
  WorkspaceWorkspaceFolders      -> error "ReqWorkspaceFolders not supported yet"
  WorkspaceConfiguration         -> error "ReqWorkspaceConfiguration not supported yet"
  CustomServerMethod _
      | "id" `HM.member` obj && "method" `HM.member` obj -> ReqCustomServer $ fromJust $ decode bytes
      | "id" `HM.member` obj -> RspCustomServer $ fromJust $ decode bytes
      | otherwise -> NotCustomServer $ fromJust $ decode bytes

Error e -> error e
-}