Connection.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StrictData #-}

-- | Multiplexed connection and channel abstractions for Hegel.
--
-- This module implements a background-reader model: a dedicated thread reads
-- packets from the input stream and dispatches them to per-channel inboxes.
-- Channels use STM 'TQueue' for incoming messages.
--
-- A 'Connection' manages a pair of 'Handle's (one for reading, one for
-- writing) with multiplexed logical channels. A 'Channel' provides
-- request\/response messaging over a single logical channel.
module Hegel.Connection
  ( -- * Types
    Connection (..)
  , Channel (..)
  , InboxItem (..)
  , RequestError (..)
  , PendingRequest

    -- * Connection lifecycle
  , createConnection
  , closeConnection

    -- * Channel management
  , newChannel
  , connectChannel
  , closeChannel

    -- * Sending
  , sendPacket
  , sendRequest
  , sendRequestCBOR

    -- * Receiving
  , receiveResponse
  , receiveRequest
  , sendResponseValue

    -- * Handshake
  , sendHandshake

    -- * Result handling
  , resultOrError

    -- * Pending requests
  , request
  , pendingGet

    -- * Accessors
  , controlChannel
  , channelId
  , serverHasExited
  , setServerExited
  , isLive
  , sendRequestAndWait
  , sendResponseRaw
  , serverCrashedMessage
  ) where

import Codec.CBOR.Term (Term (..))
import Control.Concurrent.Async (Async, async, cancel)
import Control.Concurrent.MVar (MVar, newMVar, withMVar)
import Control.Concurrent.STM
  ( STM
  , TQueue
  , TVar
  , atomically
  , check
  , newTQueueIO
  , newTVarIO
  , orElse
  , readTQueue
  , readTVar
  , readTVarIO
  , registerDelay
  , writeTQueue
  , writeTVar
  , modifyTVar'
  )
import Control.Exception (Exception, SomeException, catch, throwIO)
import Control.Monad (unless, when)
import Data.Bits (shiftL, (.|.))
import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as BS8
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import qualified Data.Text as T
import Data.Word (Word32)
import System.IO (Handle, hClose, hSetBinaryMode)

import Hegel.Protocol

-- ---------------------------------------------------------------------------
-- Constants
-- ---------------------------------------------------------------------------

-- | Handshake string sent by the client to initiate the protocol.
handshakeString :: ByteString
handshakeString = "hegel_handshake_start"

-- | Default timeout in microseconds for channel operations (30 seconds).
channelTimeoutMicros :: Int
channelTimeoutMicros = 30 * 1000000

-- | Message displayed when the server has crashed.
serverCrashedMessage :: String
serverCrashedMessage =
  "The hegel server process has exited unexpectedly. Check .hegel/server.log for details."

-- ---------------------------------------------------------------------------
-- Types
-- ---------------------------------------------------------------------------

-- | Sentinel value placed in a channel's inbox when the connection shuts down.
data InboxItem
  = Pkt Packet
  -- ^ A normal packet.
  | Shutdown
  -- ^ The connection has been closed.
  deriving (Show, Eq)

-- | Error response from the peer.
data RequestError = RequestError
  { reqErrorMessage :: String
  , reqErrorType :: String
  , reqErrorData :: [(Term, Term)]
  }
  deriving (Show)

instance Exception RequestError

-- | Multiplexed connection to a Hegel peer.
data Connection = Connection
  { connReadHandle :: Handle
  -- ^ Handle for reading packets.
  , connWriteHandle :: Handle
  -- ^ Handle for writing packets.
  , connNextChannelId :: TVar Int
  -- ^ Counter for allocating new channel IDs.
  , connChannels :: TVar (Map Word32 Channel)
  -- ^ Map from channel ID to live channel.
  , connRunning :: TVar Bool
  -- ^ Whether the connection is still active.
  , connServerExited :: TVar Bool
  -- ^ Whether the server process has exited unexpectedly.
  , connWriterLock :: MVar ()
  -- ^ Lock for serializing writes to the output handle.
  , connReaderThread :: Async ()
  -- ^ Background reader thread.
  , connDebug :: Bool
  -- ^ Whether debug logging is enabled.
  }

-- | Logical channel for request\/response messaging.
data Channel = Channel
  { chanId :: Word32
  -- ^ The channel ID.
  , chanConn :: Connection
  -- ^ The connection this channel belongs to.
  , chanInbox :: TQueue InboxItem
  -- ^ Inbox for incoming packets dispatched by the reader thread.
  , chanResponses :: TVar (Map Word32 ByteString)
  -- ^ Map from message ID to response payload (filled by 'processOneMessage').
  , chanRequests :: TQueue Packet
  -- ^ Queue of incoming request packets (filled by 'processOneMessage').
  , chanNextMessageId :: TVar Word32
  -- ^ Next message ID to use for outgoing requests.
  , chanClosed :: TVar Bool
  -- ^ Whether this channel has been closed.
  }

-- | A pending request handle that caches the decoded response.
data PendingRequest = PendingRequest
  { prChannel :: Channel
  , prMessageId :: Word32
  , prValue :: IORef (Maybe Term)
  }

-- ---------------------------------------------------------------------------
-- Connection lifecycle
-- ---------------------------------------------------------------------------

-- | Create a new connection using separate handles for reading and writing.
--
-- A control channel (channel 0) is automatically created and a background
-- reader thread is spawned. Both handles are set to binary mode.
createConnection :: Handle -> Handle -> Bool -> IO Connection
createConnection readH writeH debug = do
  hSetBinaryMode readH True
  hSetBinaryMode writeH True

  nextChanId <- newTVarIO 1
  channels <- newTVarIO Map.empty
  running <- newTVarIO True
  serverExited <- newTVarIO False
  writerLock <- newMVar ()

  -- Create a placeholder connection (reader thread filled in below)
  let mkConn reader = Connection
        { connReadHandle = readH
        , connWriteHandle = writeH
        , connNextChannelId = nextChanId
        , connChannels = channels
        , connRunning = running
        , connServerExited = serverExited
        , connWriterLock = writerLock
        , connReaderThread = reader
        , connDebug = debug
        }

  -- Create control channel (channel 0)
  controlInbox <- newTQueueIO
  controlResponses <- newTVarIO Map.empty
  controlRequests <- newTQueueIO
  controlNextMsgId <- newTVarIO 1
  controlClosed <- newTVarIO False

  -- We need the connection reference for the channel, but the connection
  -- needs the reader thread. Use a two-phase approach: create a dummy
  -- Async, build everything, then spawn the real reader.
  dummyAsync <- async (return ())

  let tempConn = mkConn dummyAsync

  let ctrl = Channel
        { chanId = 0
        , chanConn = tempConn
        , chanInbox = controlInbox
        , chanResponses = controlResponses
        , chanRequests = controlRequests
        , chanNextMessageId = controlNextMsgId
        , chanClosed = controlClosed
        }

  atomically $ writeTVar channels (Map.singleton 0 ctrl)

  -- Spawn background reader thread
  reader <- async (readerLoop tempConn)

  -- The connection returned uses the same TVars so the reader thread's
  -- reference via tempConn is equivalent. We return a version with the
  -- real Async handle for cleanup.
  let conn = mkConn reader

  -- Update the control channel to reference the real connection
  -- (shares the same TVars, so this is just for the Async handle)
  let ctrl' = ctrl { chanConn = conn }
  atomically $ writeTVar channels (Map.singleton 0 ctrl')

  return conn

-- | Close the connection and signal all live channels.
closeConnection :: Connection -> IO ()
closeConnection conn = do
  isRunning <- readTVarIO (connRunning conn)
  if not isRunning
    then return ()
    else do
      atomically $ writeTVar (connRunning conn) False
      catch (hClose (connReadHandle conn)) (\(_ :: SomeException) -> return ())
      -- Always attempt to close the write handle; if it is the same as
      -- the read handle, hClose will simply fail silently.
      catch (hClose (connWriteHandle conn)) (\(_ :: SomeException) -> return ())
      signalAllChannels conn
      cancel (connReaderThread conn)

-- ---------------------------------------------------------------------------
-- Background reader
-- ---------------------------------------------------------------------------

-- | Background reader thread function. Reads packets from the connection's
-- read handle and dispatches them to the appropriate channel's inbox.
-- Exits when the stream closes or an error occurs.
readerLoop :: Connection -> IO ()
readerLoop conn = do
  let loop = do
        isRunning <- readTVarIO (connRunning conn)
        when isRunning $ do
          pkt <- readPacket (connReadHandle conn)
          dispatchPacket conn pkt
          loop

  catch loop (\(_ :: SomeException) -> do
    atomically $ writeTVar (connServerExited conn) True)

  signalAllChannels conn

-- | Dispatch a single packet to the appropriate channel.
dispatchPacket :: Connection -> Packet -> IO ()
dispatchPacket conn pkt
  | packetPayload pkt == closeChannelPayload = return ()
  | otherwise = do
      mChan <- Map.lookup (packetChannelId pkt) <$> readTVarIO (connChannels conn)
      case mChan of
        Just ch -> atomically $ writeTQueue (chanInbox ch) (Pkt pkt)
        Nothing ->
          -- For non-reply messages to non-existent channels, send error back
          unless (packetIsReply pkt) $ do
            let errorMsg = "Message sent to non-existent channel "
                  ++ show (packetChannelId pkt)
                errorPayload = encodeTerm $
                  TMap [(TString "error", TString (T.pack errorMsg))]
            catch
              (sendPacket conn Packet
                { packetChannelId = packetChannelId pkt
                , packetMessageId = packetMessageId pkt
                , packetIsReply = True
                , packetPayload = errorPayload
                })
              (\(_ :: SomeException) -> return ())

-- | Push 'Shutdown' into all live channels' inboxes.
signalAllChannels :: Connection -> IO ()
signalAllChannels conn = do
  chans <- readTVarIO (connChannels conn)
  mapM_ (\ch -> atomically $ writeTQueue (chanInbox ch) Shutdown) (Map.elems chans)

-- ---------------------------------------------------------------------------
-- Channel management
-- ---------------------------------------------------------------------------

-- | Create a new logical channel on the connection.
--
-- Client channels use odd IDs: @(counter \<\< 1) | 1@.
newChannel :: Connection -> IO Channel
newChannel conn = do
  chanid <- atomically $ do
    counter <- readTVar (connNextChannelId conn)
    let cid = fromIntegral ((counter `shiftL` 1) .|. 1) :: Word32
    writeTVar (connNextChannelId conn) (counter + 1)
    return cid

  inbox <- newTQueueIO
  responses <- newTVarIO Map.empty
  requests <- newTQueueIO
  nextMsgId <- newTVarIO 1
  closed <- newTVarIO False

  let channel = Channel
        { chanId = chanid
        , chanConn = conn
        , chanInbox = inbox
        , chanResponses = responses
        , chanRequests = requests
        , chanNextMessageId = nextMsgId
        , chanClosed = closed
        }

  atomically $ modifyTVar' (connChannels conn) (Map.insert chanid channel)
  return channel

-- | Connect to a channel created by the peer (using the peer-assigned ID).
connectChannel :: Connection -> Word32 -> IO Channel
connectChannel conn cid = do
  exists <- atomically $ do
    chans <- readTVar (connChannels conn)
    return (Map.member cid chans)

  if exists
    then ioError $ userError $
      "Channel already connected as " ++ show cid ++ "."
    else do
      inbox <- newTQueueIO
      responses <- newTVarIO Map.empty
      requests <- newTQueueIO
      nextMsgId <- newTVarIO 1
      closed <- newTVarIO False

      let channel = Channel
            { chanId = cid
            , chanConn = conn
            , chanInbox = inbox
            , chanResponses = responses
            , chanRequests = requests
            , chanNextMessageId = nextMsgId
            , chanClosed = closed
            }

      atomically $ modifyTVar' (connChannels conn) (Map.insert cid channel)
      return channel

-- | Close a channel and notify the peer. Idempotent.
closeChannel :: Channel -> IO ()
closeChannel ch = do
  alreadyClosed <- readTVarIO (chanClosed ch)
  if alreadyClosed
    then return ()
    else do
      atomically $ writeTVar (chanClosed ch) True
      isRunning <- readTVarIO (connRunning (chanConn ch))
      when isRunning $
        sendPacket (chanConn ch) Packet
          { packetPayload = closeChannelPayload
          , packetMessageId = closeChannelMessageId
          , packetChannelId = chanId ch
          , packetIsReply = False
          }

-- ---------------------------------------------------------------------------
-- Sending
-- ---------------------------------------------------------------------------

-- | Send a packet on the connection, thread-safe via the writer lock.
sendPacket :: Connection -> Packet -> IO ()
sendPacket conn pkt =
  withMVar (connWriterLock conn) $ \() ->
    writePacket (connWriteHandle conn) pkt

-- | Send raw bytes as a request on a channel. Returns the message ID for
-- matching the response.
sendRequest :: Channel -> ByteString -> IO Word32
sendRequest ch payload = do
  msgId <- atomically $ do
    mid <- readTVar (chanNextMessageId ch)
    writeTVar (chanNextMessageId ch) (mid + 1)
    return mid
  sendPacket (chanConn ch) Packet
    { packetPayload = payload
    , packetChannelId = chanId ch
    , packetIsReply = False
    , packetMessageId = msgId
    }
  return msgId

-- | Send a CBOR-encoded request on a channel. Returns the message ID.
sendRequestCBOR :: Channel -> Term -> IO Word32
sendRequestCBOR ch term = sendRequest ch (encodeTerm term)

-- ---------------------------------------------------------------------------
-- Receiving
-- ---------------------------------------------------------------------------

-- | Pop an item from the channel's inbox, with a timeout in microseconds.
-- Returns 'Nothing' on timeout.
popInboxItem :: Channel -> Int -> IO InboxItem
popInboxItem ch timeoutUs = do
  timedOut <- registerDelay timeoutUs
  result <- atomically $
    (Just <$> readTQueue (chanInbox ch))
    `orElse`
    (do expired <- readTVar timedOut
        check expired
        return Nothing)
  case result of
    Just item -> return item
    Nothing -> do
      isClosed <- readTVarIO (chanClosed ch)
      if isClosed
        then ioError $ userError $ "Channel " ++ show (chanId ch) ++ " is closed"
        else do
          exited <- readTVarIO (connServerExited (chanConn ch))
          if exited
            then ioError $ userError serverCrashedMessage
            else ioError $ userError $
              "Timed out after " ++ show (timeoutUs `div` 1000000)
                ++ "s waiting for a message on channel " ++ show (chanId ch)

-- | Process one incoming message for a channel. Dispatches replies to the
-- responses map and requests to the requests queue.
processOneMessage :: Channel -> Int -> IO ()
processOneMessage ch timeoutUs = do
  item <- popInboxItem ch timeoutUs
  case item of
    Shutdown -> ioError $ userError "Connection closed"
    Pkt pkt ->
      if packetIsReply pkt
        then atomically $ modifyTVar' (chanResponses ch)
              (Map.insert (packetMessageId pkt) (packetPayload pkt))
        else atomically $ writeTQueue (chanRequests ch) pkt

-- | Wait for a raw response to a request with the given message ID.
receiveResponseRaw :: Channel -> Word32 -> Int -> IO ByteString
receiveResponseRaw ch msgId timeoutUs = do
  let waitLoop = do
        mResp <- atomically $ do
          resps <- readTVar (chanResponses ch)
          case Map.lookup msgId resps of
            Just payload -> do
              writeTVar (chanResponses ch) (Map.delete msgId resps)
              return (Just payload)
            Nothing -> return Nothing
        case mResp of
          Just payload -> return payload
          Nothing -> do
            processOneMessage ch timeoutUs
            waitLoop
  waitLoop

-- | Wait for and decode a response, extracting the result or raising
-- 'RequestError'.
receiveResponse :: Channel -> Word32 -> IO Term
receiveResponse ch msgId = do
  raw <- receiveResponseRaw ch msgId channelTimeoutMicros
  resultOrError (decodeTerm raw)

-- | Receive the next incoming request on a channel. Returns the message ID
-- and the decoded CBOR term.
receiveRequest :: Channel -> IO (Word32, Term)
receiveRequest ch = do
  let waitLoop = do
        mPkt <- atomically $ tryReadTQueue (chanRequests ch)
        case mPkt of
          Just pkt -> return (packetMessageId pkt, decodeTerm (packetPayload pkt))
          Nothing -> do
            processOneMessage ch channelTimeoutMicros
            waitLoop
  waitLoop
  where
    tryReadTQueue :: TQueue a -> STM (Maybe a)
    tryReadTQueue q = (Just <$> readTQueue q) `orElse` return Nothing

-- | Send a success response with a value wrapped as @{\"result\": value}@.
sendResponseValue :: Channel -> Word32 -> Term -> IO ()
sendResponseValue ch msgId value = do
  let payload = encodeTerm $ TMap [(TString "result", value)]
  sendPacket (chanConn ch) Packet
    { packetPayload = payload
    , packetChannelId = chanId ch
    , packetIsReply = True
    , packetMessageId = msgId
    }

-- ---------------------------------------------------------------------------
-- Result handling
-- ---------------------------------------------------------------------------

-- | Extract the @\"result\"@ field from a CBOR map, or throw 'RequestError'
-- if an @\"error\"@ field is present.
resultOrError :: Term -> IO Term
resultOrError body = do
  let pairs = extractMap body
      findText key = case lookup (TString key) pairs of
        Just v -> T.unpack (extractText v)
        Nothing -> ""
  case lookup (TString "error") pairs of
    Just _ -> do
      let msg = findText "error"
          errorType = findText "type"
          dat = filter (\(k, _) -> k /= TString "error" && k /= TString "type") pairs
      throwIO $ RequestError msg errorType dat
    Nothing ->
      case lookup (TString "result") pairs of
        Just v -> return v
        Nothing -> ioError $ userError "Response has neither 'result' nor 'error'"

-- ---------------------------------------------------------------------------
-- Pending requests
-- ---------------------------------------------------------------------------

-- | Send a CBOR request and return a 'PendingRequest' handle.
request :: Channel -> Term -> IO PendingRequest
request ch term = do
  msgId <- sendRequestCBOR ch term
  ref <- newIORef Nothing
  return PendingRequest
    { prChannel = ch
    , prMessageId = msgId
    , prValue = ref
    }

-- | Block until the response arrives and return the result.
--
-- Caches the response so subsequent calls return the same value or raise
-- the same error.
pendingGet :: PendingRequest -> IO Term
pendingGet pr = do
  cached <- readIORef (prValue pr)
  case cached of
    Just v -> resultOrError v
    Nothing -> do
      raw <- receiveResponseRaw (prChannel pr) (prMessageId pr) channelTimeoutMicros
      let v = decodeTerm raw
      writeIORef (prValue pr) (Just v)
      resultOrError v

-- ---------------------------------------------------------------------------
-- Handshake
-- ---------------------------------------------------------------------------

-- | Initiate the handshake as a client. Returns the server protocol version
-- string (e.g. @\"0.1\"@).
sendHandshake :: Connection -> IO String
sendHandshake conn = do
  -- Get control channel (channel 0)
  mCh <- atomically $ do
    chans <- readTVar (connChannels conn)
    return (Map.lookup 0 chans)
  ch <- case mCh of
    Just c -> return c
    Nothing -> ioError $ userError "Internal error: no control channel"

  msgId <- sendRequest ch handshakeString
  raw <- receiveResponseRaw ch msgId channelTimeoutMicros

  let response = BS8.unpack raw
      prefix = "Hegel/"
  if take 6 response /= prefix
    then ioError $ userError $ "Bad handshake response: " ++ show response
    else return (drop 6 response)

-- ---------------------------------------------------------------------------
-- Additional accessors
-- ---------------------------------------------------------------------------

-- | Get the control channel (channel 0).
controlChannel :: Connection -> IO Channel
controlChannel conn = do
  mCh <- atomically $ do
    chans <- readTVar (connChannels conn)
    return (Map.lookup 0 chans)
  case mCh of
    Just c -> return c
    Nothing -> ioError $ userError "Internal error: no control channel"

-- | Get the channel ID.
channelId :: Channel -> Word32
channelId = chanId

-- | Check if the server process has exited.
serverHasExited :: Connection -> IO Bool
serverHasExited conn = readTVarIO (connServerExited conn)

-- | Set the server-exited flag.
setServerExited :: Connection -> IO ()
setServerExited conn = atomically $ writeTVar (connServerExited conn) True

-- | Check if the connection is still active.
isLive :: Connection -> IO Bool
isLive conn = readTVarIO (connRunning conn)

-- | Send a CBOR request and wait for the decoded response (convenience).
sendRequestAndWait :: Channel -> Term -> IO Term
sendRequestAndWait ch msg = do
  msgId <- sendRequestCBOR ch msg
  receiveResponse ch msgId

-- | Send raw bytes as a reply.
sendResponseRaw :: Channel -> Word32 -> ByteString -> IO ()
sendResponseRaw ch msgId payload =
  sendPacket (chanConn ch) Packet
    { packetChannelId = chanId ch
    , packetMessageId = msgId
    , packetIsReply = True
    , packetPayload = payload
    }