Pipe.hs

-- | Internal module. Not part of the public API.
module Database.Bolty.Connection.Pipe
  ( connect
  , close
  , reset
  , ping
  , flush
  , fetch
  , sendRequest
  , receiveResponse
  , requireState
  , setState
  , getState
  , logon
  , logoff
  , sendTelemetry
  , MonadPipe
  -- * Connection accessors
  , connectionVersion
  , connectionAgent
  , connectionId
  , connectionTelemetryEnabled
  , connectionServerIdleTimeout
  , touchConnection
  , connectionLastActivity
  -- * Transaction primitives (plain IO)
  , beginTx
  , commitTx
  , rollbackTx
  , tryRollback
  -- * Plain IO wire helpers
  , requireStateIO
  , flushIO
  , fetchIO
  ) where

import           Data.Kind              (Constraint, Type)
import           Control.Exception      (Exception, SomeException, catch, throwIO, try)
import           Control.Monad          (when, unless)
import           Control.Monad.Trans    (MonadIO(..))
import           Control.Monad.Extra    (whenMaybe, whileJustM)
import           Control.Monad.Except
import           Data.Persist            (putBE, runPutLazy)
import qualified Data.Persist            as P
import           Data.IORef             (newIORef, readIORef, writeIORef)
import           GHC.Clock             (getMonotonicTimeNSec)
import           Data.Bits                      (shiftR, (.&.))
import           Data.Int                       (Int64)
import           Data.Word
import           Data.PackStream.Integer        (fromPSInteger)
import           GHC.Stack              (HasCallStack)
import           Prelude
import qualified Data.ByteString        as BS
import qualified Data.ByteString.Lazy   as BSL
import qualified Data.HashMap.Lazy      as H
import qualified Data.Text              as T
import qualified Data.PackStream        as PS
import           Data.PackStream.Ps     (Ps(..))
import qualified Data.Vector            as V
import qualified Network.Connection     as NC

import           Database.Bolty.Connection.Type
import qualified Database.Bolty.Connection.Connection as C
import           Database.Bolty.Message.Response
import           Database.Bolty.Message.Request
import           Database.Bolty.Value.Helpers (isNewVersion, versionMajor, supportsLogonLogoff, supportsTelemetry)
import           Database.Bolty.Util


-- | Constraint alias for monads that can perform BOLT pipe operations.
type MonadPipe :: (Type -> Type) -> Constraint
type MonadPipe m = (HasCallStack, MonadIO m, MonadError Error m)


exceptToThrow :: (Exception e, MonadIO m) => ExceptT e m b -> m b
exceptToThrow f = runExceptT f >>= \case
  Right y -> pure y
  Left e  -> liftIO $ throwIO e


-- | Connect to a Neo4j server, perform the BOLT handshake and authentication.
connect :: (MonadIO m, HasCallStack) => ValidatedConfig -> m Connection
connect ValidatedConfig{..} = do
  (rawConn, tout) <- C.connect use_tls host port timeout
  (server_version, agent, conn_id, telem, idleTimeout) <- exceptToThrow $ handshake rawConn tout ValidatedConfig{..}
  stateRef <- liftIO $ newIORef Ready
  now <- liftIO getMonotonicTimeNSec
  actRef <- liftIO $ newIORef now
  pure $ Connection { rawConnection = rawConn, timeout_milliseconds = tout
                    , version = server_version, server_state = stateRef
                    , server_agent = agent, connection_id = conn_id, lastActivity = actRef
                    , telemetry_enabled = telem, serverIdleTimeout = idleTimeout
                    , queryLogger = queryLogger
                    , notificationHandler = notificationHandler }


-- | Send GOODBYE (if supported) and close the connection.
close :: (MonadIO m, HasCallStack) => Connection -> m ()
close Connection{rawConnection, timeout_milliseconds, version, server_state} = exceptToThrow $ do
  _ <- whenMaybe (isNewVersion version) $ sendRequest rawConnection timeout_milliseconds RGoodbye
  liftIO $ writeIORef server_state Disconnected
  C.close rawConnection timeout_milliseconds


-- | Send RESET and transition the connection back to the Ready state.
reset :: (MonadIO m, HasCallStack) => Connection -> m ()
reset Connection{rawConnection, timeout_milliseconds, server_state} = exceptToThrow $ do
  sendRequest rawConnection timeout_milliseconds RReset
  response <- receiveResponse rawConnection timeout_milliseconds
  case response of
    (RSuccess _) -> liftIO $ writeIORef server_state Ready
    _            -> throwError ResetFailed


-- | Check if a Pipe is alive by sending RESET and expecting SUCCESS.
-- Returns True if healthy, False otherwise. Catches all exceptions.
ping :: MonadIO m => Connection -> m Bool
ping Connection{rawConnection, timeout_milliseconds, server_state} = liftIO $ do
  result <- try @SomeException $ runExceptT $ do
    sendRequest rawConnection timeout_milliseconds RReset
    response <- receiveResponse rawConnection timeout_milliseconds
    case response of
      (RSuccess _) -> liftIO $ writeIORef server_state Ready
      _            -> throwError ResetFailed
  pure $ case result of
    Right (Right _) -> True
    _               -> False


-- | Check that the connection is in one of the allowed states, or throw 'InvalidState'
requireState :: MonadPipe m => Connection -> [ServerState] -> T.Text -> m ()
requireState Connection{server_state} allowed action = do
  st <- liftIO $ readIORef server_state
  unless (st `elem` allowed) $
    throwError $ InvalidState st action


-- | Update the server state
setState :: MonadIO m => Connection -> ServerState -> m ()
setState Connection{server_state} st = liftIO $ writeIORef server_state st


-- | Get the current server state
getState :: MonadIO m => Connection -> m ServerState
getState Connection{server_state} = liftIO $ readIORef server_state


handshake :: MonadPipe m => NC.Connection -> Int -> ValidatedConfig -> m (Word32, T.Text, T.Text, Bool, Maybe Int)
handshake rawConn tout ValidatedConfig{versions, user_agent, scheme = authScheme, routing} = do
    -- https://neo4j.com/docs/bolt/current/bolt/handshake/#_bolt_identification
    C.send rawConn tout (encodeStrict (0x6060B017 :: Word32))
    C.send rawConn tout $ P.runPut $ mapM_ putBE versions
    server_version :: Word32 <- C.receiveBinary rawConn tout 4
    when (not $ versionAccepted server_version versions) $ throwError $ UnsupportedServerVersion server_version
    if supportsLogonLogoff server_version then do
      -- Bolt 5.1+: HELLO without credentials (no scheme/principal/credentials fields)
      let serverMinor = fromIntegral ((server_version `shiftR` 8) .&. 0xFF) :: Word8
      let helloDict = PsDictionary $ H.fromList $
            [ ("user_agent", PS.toPs user_agent)
            ] <>
            -- Bolt 5.3+ requires bolt_agent dictionary
            ( if serverMinor >= 3 then
                [("bolt_agent", PsDictionary $ H.singleton "product" (PS.toPs user_agent))]
              else
                []
            ) <>
            ( case routing of
                NoRouting                          -> []
                Routing                            -> [("routing", PsDictionary H.empty)]
                RoutingSpec address urlQueryParams  -> [("routing", PsDictionary $ H.fromList $
                    ("address", PS.toPs address) : [(k, PS.toPs v) | (k, v) <- H.toList urlQueryParams])]
            )
      sendPs rawConn tout $ PsStructure 0x01 $ V.singleton helloDict
      helloResponse <- receiveResponse rawConn tout
      case helloResponse of
        RSuccess meta -> do
          let agent   = psToText $ H.lookupDefault (PsString "") "server" meta
          let conn_id = psToText $ H.lookupDefault (PsString "") "connection_id" meta
          let telem   = parseTelemetryHint meta
          let idleTO  = parseIdleTimeoutHint meta
          -- Now send LOGON with credentials
          sendRequest rawConn tout $ RLogon $ Logon authScheme
          logonResponse <- receiveResponse rawConn tout
          case logonResponse of
            RSuccess _ -> pure (server_version, agent, conn_id, telem, idleTO)
            _          -> throwError AuthentificationFailed
        _ -> throwError AuthentificationFailed
    else do
      -- Bolt 5.0 and below: HELLO with credentials
      -- BOLT 4.x: include patch_bolt to get UTC-based DateTime (0x49/0x69) instead of legacy format
      let needsPatch = versionMajor server_version == 4
      sendRequest rawConn tout $ RHello $ Hello user_agent authScheme routing needsPatch
      response <- receiveResponse rawConn tout
      case response of
        RSuccess meta -> do
          let agent   = psToText $ H.lookupDefault (PsString "") "server" meta
          let conn_id = psToText $ H.lookupDefault (PsString "") "connection_id" meta
          let idleTO  = parseIdleTimeoutHint meta
          pure (server_version, agent, conn_id, False, idleTO)
        _ -> throwError AuthentificationFailed
  where
    psToText (PsString t) = t
    psToText _            = ""

    -- Parse hints.telemetry.enabled from HELLO SUCCESS metadata
    parseTelemetryHint :: H.HashMap T.Text Ps -> Bool
    parseTelemetryHint meta =
      case H.lookup "hints" meta of
        Just (PsDictionary hints) ->
          case H.lookup "telemetry.enabled" hints of
            Just (PsBoolean b) -> b
            _                  -> False
        _ -> False

    -- Parse hints.connection.recv_timeout_seconds from HELLO SUCCESS metadata
    parseIdleTimeoutHint :: H.HashMap T.Text Ps -> Maybe Int
    parseIdleTimeoutHint meta =
      case H.lookup "hints" meta of
        Just (PsDictionary hints) ->
          case H.lookup "connection.recv_timeout_seconds" hints of
            Just (PsInteger n) -> fmap fromIntegral (fromPSInteger n :: Maybe Int64)
            _                  -> Nothing
        _ -> Nothing


-- | Check if a server version response matches any of the client's version specs.
-- The client sends compact version specs: word32 = [0, range, minor, major].
-- The server responds with: word32 = [0, 0, minor, major].
-- A server version matches a spec if the majors match and the server minor
-- falls in [spec_minor - range, spec_minor].
versionAccepted :: Word32 -> [Word32] -> Bool
versionAccepted serverVer = any matches
  where
    serverMajor = fromIntegral (serverVer .&. 0xFF) :: Word8
    serverMinor = fromIntegral ((serverVer `shiftR` 8) .&. 0xFF) :: Word8
    matches spec =
      let specMajor = fromIntegral (spec .&. 0xFF) :: Word8
          specMinor = fromIntegral ((spec `shiftR` 8) .&. 0xFF) :: Word8
          specRange = fromIntegral ((spec `shiftR` 16) .&. 0xFF) :: Word8
      in serverMajor == specMajor
         && serverMinor <= specMinor
         && serverMinor >= specMinor - specRange


-- | Receive and decode a chunked BOLT response from the wire.
receiveResponse :: MonadPipe m => NC.Connection -> Int -> m Response
receiveResponse rawConn tout = do
  bs :: BS.ByteString <- whileJustM $ do
          size :: Word16 <- C.receiveBinary rawConn tout 2
          if size == 0 then
            pure Nothing
          else
            Just <$> C.receiveBytestring rawConn tout (fromIntegral size)
  case PS.unpack' bs of
    Left e         -> throwError $ CannotReadResponse e
    Right response -> pure response


-- | Serialize and send a BOLT request as chunked data on the wire.
sendRequest :: MonadPipe m => NC.Connection -> Int -> Request -> m ()
sendRequest rawConn tout request = sendPs rawConn tout (PS.toPs request)

sendPs :: MonadPipe m => NC.Connection -> Int -> Ps -> m ()
sendPs rawConn tout ps = do
  let bs = PS.pack ps
  let chunks = chunksOfBSL 65_535 bs
  let encoded = BSL.toStrict $ mconcat (map addChunkHeader chunks) <> BSL.pack [0, 0]
  C.send rawConn tout encoded
  where addChunkHeader :: BSL.ByteString -> BSL.ByteString
        addChunkHeader chunk = let size = fromIntegral (BSL.length chunk) :: Word16
                                in runPutLazy (putBE size) <> chunk


-- | Send a request through a 'Connection'.
flush :: MonadPipe m => Connection -> Request -> m ()
flush Connection{rawConnection, timeout_milliseconds} request =
  sendRequest rawConnection timeout_milliseconds request


-- | Receive the next response from a 'Connection'.
fetch :: MonadPipe m => Connection -> m Response
fetch Connection{rawConnection, timeout_milliseconds} =
  receiveResponse rawConnection timeout_milliseconds


-- | Send LOGOFF, expect SUCCESS, transition to Authentication state.
-- Only valid on Bolt 5.1+ connections in Ready state.
logoff :: (MonadIO m, HasCallStack) => Connection -> m ()
logoff Connection{rawConnection, timeout_milliseconds, server_state} = exceptToThrow $ do
  sendRequest rawConnection timeout_milliseconds RLogoff
  response <- receiveResponse rawConnection timeout_milliseconds
  case response of
    RSuccess _ -> liftIO $ writeIORef server_state Authentication
    _          -> throwError $ WrongMessageFormat "Expected SUCCESS after LOGOFF"


-- | Send LOGON with credentials, expect SUCCESS, transition to Ready state.
-- Only valid on Bolt 5.1+ connections in Authentication state.
logon :: (MonadIO m, HasCallStack) => Connection -> Scheme -> m ()
logon Connection{rawConnection, timeout_milliseconds, server_state} authScheme = exceptToThrow $ do
  sendRequest rawConnection timeout_milliseconds $ RLogon $ Logon authScheme
  response <- receiveResponse rawConnection timeout_milliseconds
  case response of
    RSuccess _ -> liftIO $ writeIORef server_state Ready
    _          -> throwError AuthentificationFailed


-- | Send a TELEMETRY message if the server supports it.
-- No-op if telemetry is not enabled or version < 5.4.
sendTelemetry :: (MonadIO m, HasCallStack) => Connection -> TelemetryApi -> m ()
sendTelemetry Connection{telemetry_enabled, version, rawConnection, timeout_milliseconds} api
  | telemetry_enabled && supportsTelemetry version = exceptToThrow $ do
      sendRequest rawConnection timeout_milliseconds $ RTelemetry api
      response <- receiveResponse rawConnection timeout_milliseconds
      case response of
        RSuccess _ -> pure ()
        _          -> throwError $ WrongMessageFormat "Expected SUCCESS after TELEMETRY"
  | otherwise = pure ()


-- ---------------------------------------------------------------------------
-- Connection accessors
-- ---------------------------------------------------------------------------

-- | Get the negotiated BOLT protocol version.
connectionVersion :: Connection -> Word32
connectionVersion Connection{version} = version

-- | Get the server agent string.
connectionAgent :: Connection -> T.Text
connectionAgent Connection{server_agent} = server_agent

-- | Get the server-assigned connection ID.
connectionId :: Connection -> T.Text
connectionId Connection{connection_id} = connection_id

-- | Get the server-advertised idle timeout in seconds, if any.
connectionServerIdleTimeout :: Connection -> Maybe Int
connectionServerIdleTimeout Connection{serverIdleTimeout} = serverIdleTimeout

-- | Check whether the server supports telemetry.
connectionTelemetryEnabled :: Connection -> Bool
connectionTelemetryEnabled Connection{telemetry_enabled} = telemetry_enabled

-- | Update the connection's last-activity timestamp to now.
touchConnection :: MonadIO m => Connection -> m ()
touchConnection Connection{lastActivity} = liftIO $ do
  now <- getMonotonicTimeNSec
  writeIORef lastActivity now

-- | Get the monotonic timestamp (in nanoseconds) of the last connection activity.
connectionLastActivity :: MonadIO m => Connection -> m Word64
connectionLastActivity Connection{lastActivity} = liftIO $ readIORef lastActivity


-- ---------------------------------------------------------------------------
-- Transaction primitives (plain IO)
-- ---------------------------------------------------------------------------

-- | Begin an explicit transaction with bookmarks and access mode.
beginTx :: HasCallStack => Connection -> Begin -> IO ()
beginTx conn@Connection{rawConnection, timeout_milliseconds, server_state} params = exceptToThrow $ do
  requireState conn [Ready] "BEGIN"
  sendRequest rawConnection timeout_milliseconds $ RBegin params
  response <- receiveResponse rawConnection timeout_milliseconds
  case response of
    RSuccess _ -> liftIO $ writeIORef server_state TXready
    RFailure Failure{code, message} -> do
      liftIO $ writeIORef server_state Failed
      liftIO $ reset conn
      throwError $ ResponseErrorFailure code message
    _ -> do
      liftIO $ reset conn
      throwError $ WrongMessageFormat "Unexpected response to BEGIN"

-- | Commit the current transaction and return the bookmark (if any).
commitTx :: HasCallStack => Connection -> IO (Maybe T.Text)
commitTx conn@Connection{rawConnection, timeout_milliseconds, server_state} = exceptToThrow $ do
  requireState conn [TXready] "COMMIT"
  sendRequest rawConnection timeout_milliseconds RCommit
  response <- receiveResponse rawConnection timeout_milliseconds
  case response of
    RSuccess meta -> do
      liftIO $ writeIORef server_state Ready
      pure $ extractBookmark meta
    RFailure Failure{code, message} -> do
      liftIO $ writeIORef server_state Failed
      liftIO $ reset conn
      throwError $ ResponseErrorFailure code message
    _ -> do
      liftIO $ reset conn
      throwError $ WrongMessageFormat "Unexpected response to COMMIT"

-- | Rollback the current transaction.
rollbackTx :: HasCallStack => Connection -> IO ()
rollbackTx conn@Connection{rawConnection, timeout_milliseconds, server_state} = exceptToThrow $ do
  requireState conn [TXready] "ROLLBACK"
  sendRequest rawConnection timeout_milliseconds RRollback
  response <- receiveResponse rawConnection timeout_milliseconds
  case response of
    RSuccess _ -> liftIO $ writeIORef server_state Ready
    RFailure Failure{code, message} -> do
      liftIO $ writeIORef server_state Failed
      liftIO $ reset conn
      throwError $ ResponseErrorFailure code message
    _ -> do
      liftIO $ reset conn
      throwError $ WrongMessageFormat "Unexpected response to ROLLBACK"


-- | Try to rollback the current transaction, ignoring errors.
-- Used as cleanup in onException handlers.
tryRollback :: Connection -> IO ()
tryRollback conn = do
  st <- getState conn
  when (st == TXready || st == TXstreaming) $
    catch (rollbackTx conn) handler
  where
    handler :: Error -> IO ()
    handler _ = pure ()


-- ---------------------------------------------------------------------------
-- Plain IO wire helpers
-- ---------------------------------------------------------------------------

-- | Check the connection state, throw InvalidState if wrong.
requireStateIO :: HasCallStack => Connection -> [ServerState] -> T.Text -> IO ()
requireStateIO conn allowed action = do
  st <- getState conn
  unless (st `elem` allowed) $ throwIO $ InvalidState st action

-- | Send a request, throwing on wire error.
flushIO :: HasCallStack => Connection -> Request -> IO ()
flushIO Connection{rawConnection, timeout_milliseconds} req =
  exceptToThrow $ sendRequest rawConnection timeout_milliseconds req

-- | Receive a response, throwing on wire error.
fetchIO :: HasCallStack => Connection -> IO Response
fetchIO Connection{rawConnection, timeout_milliseconds} =
  exceptToThrow $ receiveResponse rawConnection timeout_milliseconds