{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeOperators #-}

module Language.LSP.Server.Control
  (
  -- * Running
    runServer
  , runServerWith
  , runServerWithHandles
  ) where

import           Control.Concurrent
import           Control.Concurrent.STM.TChan
import           Control.Monad
import           Control.Monad.STM
import qualified Data.Aeson as J
import qualified Data.Attoparsec.ByteString as Attoparsec
import Data.Attoparsec.ByteString.Char8
import qualified Data.ByteString as BS
import Data.ByteString.Builder.Extra (defaultChunkSize)
import qualified Data.ByteString.Lazy as BSL
import qualified Data.Text.Lazy as TL
import qualified Data.Text.Lazy.Encoding as TL
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import           Data.List
import           Language.LSP.Server.Core
import           Language.LSP.Server.Processing
import           Language.LSP.Types
import           Language.LSP.VFS
import           System.IO
import           System.Log.Logger


-- ---------------------------------------------------------------------

-- | Convenience function for 'runServerWithHandles stdin stdout'.
runServer :: ServerDefinition config
                -- ^ function to be called once initialize has
                -- been received from the client. Further message
                -- processing will start only after this returns.
    -> IO Int
runServer :: ServerDefinition config -> IO Int
runServer = Handle -> Handle -> ServerDefinition config -> IO Int
forall config.
Handle -> Handle -> ServerDefinition config -> IO Int
runServerWithHandles Handle
stdin Handle
stdout

-- | Starts a language server over the specified handles. 
-- This function will return once the @exit@ notification is received.
runServerWithHandles ::
       Handle
    -- ^ Handle to read client input from.
    -> Handle
    -- ^ Handle to write output to.
    -> ServerDefinition config
    -> IO Int         -- exit code
runServerWithHandles :: Handle -> Handle -> ServerDefinition config -> IO Int
runServerWithHandles Handle
hin Handle
hout ServerDefinition config
serverDefinition = do

  Handle -> BufferMode -> IO ()
hSetBuffering Handle
hin BufferMode
NoBuffering
  Handle -> TextEncoding -> IO ()
hSetEncoding  Handle
hin TextEncoding
utf8

  Handle -> BufferMode -> IO ()
hSetBuffering Handle
hout BufferMode
NoBuffering
  Handle -> TextEncoding -> IO ()
hSetEncoding  Handle
hout TextEncoding
utf8

  let
    clientIn :: IO ByteString
clientIn = Handle -> Int -> IO ByteString
BS.hGetSome Handle
hin Int
defaultChunkSize

    clientOut :: ByteString -> IO ()
clientOut ByteString
out = do
      Handle -> ByteString -> IO ()
BSL.hPut Handle
hout ByteString
out
      Handle -> IO ()
hFlush Handle
hout

  IO ByteString
-> (ByteString -> IO ()) -> ServerDefinition config -> IO Int
forall config.
IO ByteString
-> (ByteString -> IO ()) -> ServerDefinition config -> IO Int
runServerWith IO ByteString
clientIn ByteString -> IO ()
clientOut ServerDefinition config
serverDefinition

-- | Starts listening and sending requests and responses
-- using the specified I/O.
runServerWith ::
       IO BS.ByteString
    -- ^ Client input.
    -> (BSL.ByteString -> IO ())
    -- ^ Function to provide output to.
    -> ServerDefinition config
    -> IO Int         -- exit code
runServerWith :: IO ByteString
-> (ByteString -> IO ()) -> ServerDefinition config -> IO Int
runServerWith IO ByteString
clientIn ByteString -> IO ()
clientOut ServerDefinition config
serverDefinition = do

  String -> String -> IO ()
infoM String
"lsp.runWith" String
"\n\n\n\n\nlsp:Starting up server ..."

  TChan Value
cout <- STM (TChan Value) -> IO (TChan Value)
forall a. STM a -> IO a
atomically STM (TChan Value)
forall a. STM (TChan a)
newTChan :: IO (TChan J.Value)
  ThreadId
_rhpid <- IO () -> IO ThreadId
forkIO (IO () -> IO ThreadId) -> IO () -> IO ThreadId
forall a b. (a -> b) -> a -> b
$ TChan Value -> (ByteString -> IO ()) -> IO ()
sendServer TChan Value
cout ByteString -> IO ()
clientOut

  let sendMsg :: a -> IO ()
sendMsg a
msg = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TChan Value -> Value -> STM ()
forall a. TChan a -> a -> STM ()
writeTChan TChan Value
cout (Value -> STM ()) -> Value -> STM ()
forall a b. (a -> b) -> a -> b
$ a -> Value
forall a. ToJSON a => a -> Value
J.toJSON a
msg

  (VFS -> IO ()) -> IO ()
forall r. (VFS -> IO r) -> IO r
initVFS ((VFS -> IO ()) -> IO ()) -> (VFS -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \VFS
vfs -> do
    IO ByteString
-> ServerDefinition config
-> VFS
-> (FromServerMessage -> IO ())
-> IO ()
forall config.
IO ByteString
-> ServerDefinition config
-> VFS
-> (FromServerMessage -> IO ())
-> IO ()
ioLoop IO ByteString
clientIn ServerDefinition config
serverDefinition VFS
vfs FromServerMessage -> IO ()
forall a. ToJSON a => a -> IO ()
sendMsg

  Int -> IO Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
1

-- ---------------------------------------------------------------------

ioLoop ::
     IO BS.ByteString
  -> ServerDefinition config
  -> VFS
  -> (FromServerMessage -> IO ())
  -> IO ()
ioLoop :: IO ByteString
-> ServerDefinition config
-> VFS
-> (FromServerMessage -> IO ())
-> IO ()
ioLoop IO ByteString
clientIn ServerDefinition config
serverDefinition VFS
vfs FromServerMessage -> IO ()
sendMsg = do
  Maybe (ByteString, ByteString)
minitialize <- Result ByteString -> IO (Maybe (ByteString, ByteString))
parseOne (Parser ByteString -> ByteString -> Result ByteString
forall a. Parser a -> ByteString -> Result a
parse Parser ByteString
parser ByteString
"")
  case Maybe (ByteString, ByteString)
minitialize of
    Maybe (ByteString, ByteString)
Nothing -> () -> IO ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    Just (ByteString
msg,ByteString
remainder) -> do
      case ByteString
-> Either String (RequestMessage @'FromClient 'Initialize)
forall a. FromJSON a => ByteString -> Either String a
J.eitherDecode (ByteString
 -> Either String (RequestMessage @'FromClient 'Initialize))
-> ByteString
-> Either String (RequestMessage @'FromClient 'Initialize)
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
BSL.fromStrict ByteString
msg of
        Left String
err ->
          String -> String -> IO ()
errorM String
"lsp.ioLoop" (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$
            String
"Got error while decoding initialize:\n" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
err String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"\n exiting 1 ...\n"
        Right RequestMessage @'FromClient 'Initialize
initialize -> do
          Maybe (LanguageContextEnv config)
mInitResp <- ServerDefinition config
-> VFS
-> (FromServerMessage -> IO ())
-> Message @'FromClient @'Request 'Initialize
-> IO (Maybe (LanguageContextEnv config))
forall config.
ServerDefinition config
-> VFS
-> (FromServerMessage -> IO ())
-> Message @'FromClient @'Request 'Initialize
-> IO (Maybe (LanguageContextEnv config))
initializeRequestHandler ServerDefinition config
serverDefinition VFS
vfs FromServerMessage -> IO ()
sendMsg Message @'FromClient @'Request 'Initialize
RequestMessage @'FromClient 'Initialize
initialize
          case Maybe (LanguageContextEnv config)
mInitResp of
            Maybe (LanguageContextEnv config)
Nothing -> () -> IO ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            Just LanguageContextEnv config
env -> LanguageContextEnv config -> Result ByteString -> IO ()
forall config.
LanguageContextEnv config -> Result ByteString -> IO ()
loop LanguageContextEnv config
env (Parser ByteString -> ByteString -> Result ByteString
forall a. Parser a -> ByteString -> Result a
parse Parser ByteString
parser ByteString
remainder)
  where

    parseOne :: Result BS.ByteString -> IO (Maybe (BS.ByteString,BS.ByteString))
    parseOne :: Result ByteString -> IO (Maybe (ByteString, ByteString))
parseOne (Fail ByteString
_ [String]
ctxs String
err) = do
      String -> String -> IO ()
errorM String
"lsp.parseOne" (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$
        String
"Failed to parse message header:\n" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
" > " [String]
ctxs String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
": " String -> String -> String
forall a. Semigroup a => a -> a -> a
<>
        String
err String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"\n exiting 1 ...\n"
      Maybe (ByteString, ByteString)
-> IO (Maybe (ByteString, ByteString))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (ByteString, ByteString)
forall a. Maybe a
Nothing
    parseOne (Partial ByteString -> Result ByteString
c) = do
      ByteString
bs <- IO ByteString
clientIn
      if ByteString -> Bool
BS.null ByteString
bs
        then do
          String -> String -> IO ()
errorM String
"lsp.parseON" String
"lsp:Got EOF, exiting 1 ...\n"
          Maybe (ByteString, ByteString)
-> IO (Maybe (ByteString, ByteString))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (ByteString, ByteString)
forall a. Maybe a
Nothing
        else Result ByteString -> IO (Maybe (ByteString, ByteString))
parseOne (ByteString -> Result ByteString
c ByteString
bs)
    parseOne (Done ByteString
remainder ByteString
msg) = do
      String -> String -> IO ()
debugM String
"lsp.parseOne" (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"---> " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Text -> String
T.unpack (ByteString -> Text
T.decodeUtf8 ByteString
msg)
      Maybe (ByteString, ByteString)
-> IO (Maybe (ByteString, ByteString))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (ByteString, ByteString)
 -> IO (Maybe (ByteString, ByteString)))
-> Maybe (ByteString, ByteString)
-> IO (Maybe (ByteString, ByteString))
forall a b. (a -> b) -> a -> b
$ (ByteString, ByteString) -> Maybe (ByteString, ByteString)
forall a. a -> Maybe a
Just (ByteString
msg,ByteString
remainder)

    loop :: LanguageContextEnv config -> Result ByteString -> IO ()
loop LanguageContextEnv config
env = Result ByteString -> IO ()
go
      where
        go :: Result ByteString -> IO ()
go Result ByteString
r = do
          Maybe (ByteString, ByteString)
res <- Result ByteString -> IO (Maybe (ByteString, ByteString))
parseOne Result ByteString
r
          case Maybe (ByteString, ByteString)
res of
            Maybe (ByteString, ByteString)
Nothing -> () -> IO ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            Just (ByteString
msg,ByteString
remainder) -> do
              LanguageContextEnv config -> LspT config IO () -> IO ()
forall config (m :: * -> *) a.
LanguageContextEnv config -> LspT config m a -> m a
runLspT LanguageContextEnv config
env (LspT config IO () -> IO ()) -> LspT config IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ ByteString -> LspT config IO ()
forall config. ByteString -> LspM config ()
processMessage (ByteString -> LspT config IO ())
-> ByteString -> LspT config IO ()
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
BSL.fromStrict ByteString
msg
              Result ByteString -> IO ()
go (Parser ByteString -> ByteString -> Result ByteString
forall a. Parser a -> ByteString -> Result a
parse Parser ByteString
parser ByteString
remainder)

    parser :: Parser ByteString
parser = do
      ByteString
_ <- ByteString -> Parser ByteString
string ByteString
"Content-Length: "
      Int
len <- Parser Int
forall a. Integral a => Parser a
decimal
      ByteString
_ <- ByteString -> Parser ByteString
string ByteString
_TWO_CRLF
      Int -> Parser ByteString
Attoparsec.take Int
len

-- ---------------------------------------------------------------------

-- | Simple server to make sure all output is serialised
sendServer :: TChan J.Value -> (BSL.ByteString -> IO ()) -> IO ()
sendServer :: TChan Value -> (ByteString -> IO ()) -> IO ()
sendServer TChan Value
msgChan ByteString -> IO ()
clientOut = do
  IO () -> IO ()
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    Value
msg <- STM Value -> IO Value
forall a. STM a -> IO a
atomically (STM Value -> IO Value) -> STM Value -> IO Value
forall a b. (a -> b) -> a -> b
$ TChan Value -> STM Value
forall a. TChan a -> STM a
readTChan TChan Value
msgChan

    -- We need to make sure we only send over the content of the message,
    -- and no other tags/wrapper stuff
    let str :: ByteString
str = Value -> ByteString
forall a. ToJSON a => a -> ByteString
J.encode Value
msg

    let out :: ByteString
out = [ByteString] -> ByteString
BSL.concat
                [ Text -> ByteString
TL.encodeUtf8 (Text -> ByteString) -> Text -> ByteString
forall a b. (a -> b) -> a -> b
$ String -> Text
TL.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ String
"Content-Length: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int64 -> String
forall a. Show a => a -> String
show (ByteString -> Int64
BSL.length ByteString
str)
                , ByteString -> ByteString
BSL.fromStrict ByteString
_TWO_CRLF
                , ByteString
str ]

    ByteString -> IO ()
clientOut ByteString
out
    String -> String -> IO ()
debugM String
"lsp.sendServer" (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"<--2--" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Text -> String
TL.unpack (ByteString -> Text
TL.decodeUtf8 ByteString
str)

-- |
--
--
_TWO_CRLF :: BS.ByteString
_TWO_CRLF :: ByteString
_TWO_CRLF = ByteString
"\r\n\r\n"