Transport.hs

-- | Low-level transport for Substrate RPC calls
--
-- Pure IO functions for WebSocket communication with the Substrate backend.
-- All calls go through '<backend>.call' for routing.
module Plexus.Transport
  ( -- * RPC Calls (collected)
    rpcCall
  , rpcCallWith

    -- * RPC Calls (streaming)
  , rpcCallStreaming
  , invokeMethodStreaming

    -- * Bidirectional Response
  , sendBidirectionalResponse

    -- * Schema Fetching
  , fetchSchemaAt
  , fetchMethodSchemaAt
  , extractSchema
  , extractSchemaResult

    -- * Method Invocation (collected)
  , invokeMethod
  , invokeRaw
  ) where

import Control.Exception (SomeException, IOException, catch, fromException)
import qualified Control.Exception as E
import Control.Concurrent.MVar (MVar, newMVar, modifyMVar)
import Data.Aeson hiding (Error)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Text (Text)
import qualified Data.Text as T
import qualified Streaming.Prelude as S
import qualified Network.Socket as NS
import System.IO.Unsafe (unsafePerformIO)

import Plexus.Client (SubstrateConfig(..), connect, disconnect, substrateRpc, defaultConfig)
import Plexus.Client.Pool (ConnectionPool, createConnectionPool, withPooledConnection, defaultPlexusPoolConfig)
import Plexus.Types (PlexusStreamItem(..), TransportError(..), Response(..), StandardResponse)
import Plexus.Schema.Recursive (PluginSchema, MethodSchema, SchemaResult(..), parsePluginSchema, parseSchemaResult)

-- | Global connection pool cache
-- One pool per unique SubstrateConfig to enable connection reuse across calls
{-# NOINLINE poolCache #-}
poolCache :: MVar (Map SubstrateConfig ConnectionPool)
poolCache = unsafePerformIO (newMVar Map.empty)

-- | Get or create a connection pool for the given config
getOrCreatePool :: SubstrateConfig -> IO ConnectionPool
getOrCreatePool cfg = modifyMVar poolCache $ \pools ->
  case Map.lookup cfg pools of
    Just pool -> pure (pools, pool)
    Nothing -> do
      pool <- createConnectionPool cfg defaultPlexusPoolConfig
      pure (Map.insert cfg pool pools, pool)

-- | Low-level RPC call with default localhost config
rpcCall :: Text -> Text -> Value -> IO (Either TransportError [PlexusStreamItem])
rpcCall backend = rpcCallWith (defaultConfig backend)

-- | Low-level RPC call with custom config
rpcCallWith :: SubstrateConfig -> Text -> Value -> IO (Either TransportError [PlexusStreamItem])
rpcCallWith cfg method params = do
  result <- (Right <$> doCallInner cfg method params)
    `catch` categorizeException cfg
  pure result

-- | Categorize exceptions into typed TransportError
categorizeException :: SubstrateConfig -> SomeException -> IO (Either TransportError a)
categorizeException cfg e
  -- Check for connection refused (most common)
  | Just ioErr <- fromException e :: Maybe IOException = do
      let errMsg = show ioErr
      let host = T.pack $ substrateHost cfg
      let port = substratePort cfg
      pure $ Left $ if "refused" `isInfixOf` errMsg || "ECONNREFUSED" `isInfixOf` errMsg
        then ConnectionRefused host port
        else if "timeout" `isInfixOf` errMsg || "ETIMEDOUT" `isInfixOf` errMsg
          then ConnectionTimeout host port
          else NetworkError (T.pack errMsg)

  -- Catch protocol/parse errors
  | otherwise =
      let errMsg = show e
          host = T.pack $ substrateHost cfg
          port = substratePort cfg
      in pure $ Left $ if "protocol" `isInfixOf` errMsg || "parse" `isInfixOf` errMsg
        then ProtocolError (T.pack errMsg)
        else NetworkError (T.pack errMsg)

  where
    isInfixOf = T.isInfixOf `on` (T.toLower . T.pack)
    on f g x y = f (g x) (g y)

doCallInner :: SubstrateConfig -> Text -> Value -> IO [PlexusStreamItem]
doCallInner cfg method params = do
  pool <- getOrCreatePool cfg
  withPooledConnection pool $ \conn ->
    S.toList_ $ substrateRpc conn method params

-- | Streaming RPC call - invokes callback for each item as it arrives
rpcCallStreaming :: SubstrateConfig -> Text -> Value -> (PlexusStreamItem -> IO ()) -> IO (Either TransportError ())
rpcCallStreaming cfg method params onItem = do
  result <- (Right <$> doCallStreaming cfg method params onItem)
    `catch` categorizeException cfg
  pure result

doCallStreaming :: SubstrateConfig -> Text -> Value -> (PlexusStreamItem -> IO ()) -> IO ()
doCallStreaming cfg method params onItem = do
  pool <- getOrCreatePool cfg
  withPooledConnection pool $ \conn ->
    S.mapM_ onItem $ substrateRpc conn method params

-- | Streaming method invocation
invokeMethodStreaming :: SubstrateConfig -> [Text] -> Text -> Value -> (PlexusStreamItem -> IO ()) -> IO (Either TransportError ())
invokeMethodStreaming cfg namespacePath method params onItem = do
  let backend = substrateBackend cfg
  let fullPath = if null namespacePath then [backend] else namespacePath
  let dotPath = T.intercalate "." (fullPath ++ [method])
  let callParams = object ["method" .= dotPath, "params" .= params]
  rpcCallStreaming cfg (backend <> ".call") callParams onItem

-- | Fetch schema at a specific path
-- Empty path = root (<backend>.schema)
-- Non-empty path = child schema (e.g., ["solar", "earth"] -> solar.earth.schema)
fetchSchemaAt :: SubstrateConfig -> [Text] -> IO (Either TransportError PluginSchema)
fetchSchemaAt cfg path = do
  let backend = substrateBackend cfg
  let schemaMethod = if null path
        then backend <> ".schema"
        else T.intercalate "." path <> ".schema"
  result <- rpcCallWith cfg (backend <> ".call") (object ["method" .= schemaMethod])
  case result of
    Left transportErr -> pure $ Left transportErr
    Right items -> case extractSchema items of
      Left parseErr -> pure $ Left $ ProtocolError parseErr  -- Parse errors are protocol errors
      Right schema -> pure $ Right schema

-- | Extract PluginSchema from stream items
-- Returns Either Text for application-level errors (not transport)
extractSchema :: [PlexusStreamItem] -> Either Text PluginSchema
extractSchema items =
  case [dat | StreamData _ _ ct dat <- items, ".schema" `T.isSuffixOf` ct] of
    (dat:_) -> parsePluginSchema dat
    [] -> case [err | StreamError _ _ err _ <- items] of
      (err:_) -> Left err
      [] -> Left "No schema in response"

-- | Fetch a specific method's schema
-- Uses the parameter-based query: plugin.schema with {"method": "name"}
fetchMethodSchemaAt :: SubstrateConfig -> [Text] -> Text -> IO (Either TransportError MethodSchema)
fetchMethodSchemaAt cfg path methodName = do
  let backend = substrateBackend cfg
  let schemaMethod = if null path
        then backend <> ".schema"
        else T.intercalate "." path <> ".schema"
  result <- rpcCallWith cfg (backend <> ".call") (object
    [ "method" .= schemaMethod
    , "params" .= object ["method" .= methodName]
    ])
  case result of
    Left transportErr -> pure $ Left transportErr
    Right items -> case extractSchemaResult items of
      Left parseErr -> pure $ Left $ ProtocolError parseErr
      Right (SchemaMethod m) -> pure $ Right m
      Right (SchemaPlugin _) -> pure $ Left $ ProtocolError "Expected method schema, got plugin schema"

-- | Extract SchemaResult (plugin or method) from stream items
extractSchemaResult :: [PlexusStreamItem] -> Either Text SchemaResult
extractSchemaResult items =
  case [dat | StreamData _ _ ct dat <- items, ".schema" `T.isSuffixOf` ct] of
    (dat:_) -> parseSchemaResult dat
    [] -> case [err | StreamError _ _ err _ <- items] of
      (err:_) -> Left err
      [] -> Left "No schema in response"

-- | Invoke a method and return stream items
invokeMethod :: SubstrateConfig -> [Text] -> Text -> Value -> IO (Either TransportError [PlexusStreamItem])
invokeMethod cfg namespacePath method params = do
  let backend = substrateBackend cfg
  let fullPath = if null namespacePath then [backend] else namespacePath
  let dotPath = T.intercalate "." (fullPath ++ [method])
  let callParams = object ["method" .= dotPath, "params" .= params]
  rpcCallWith cfg (backend <> ".call") callParams

-- | Invoke with raw method path
invokeRaw :: SubstrateConfig -> Text -> Value -> IO (Either TransportError [PlexusStreamItem])
invokeRaw cfg method params = do
  let backend = substrateBackend cfg
  let callParams = object ["method" .= method, "params" .= params]
  rpcCallWith cfg (backend <> ".call") callParams

-- | Send a response back to the server for a bidirectional request
-- Uses {backend}.respond RPC method
sendBidirectionalResponse :: SubstrateConfig -> Text -> Response Value -> IO (Either TransportError ())
sendBidirectionalResponse cfg requestId response = do
  let backend = substrateBackend cfg
  let respondParams = object
        [ "request_id" .= requestId
        , "response_data" .= response
        ]
  result <- rpcCallWith cfg (backend <> ".respond") respondParams
  case result of
    Left transportErr -> pure $ Left transportErr
    Right items -> case [err | StreamError _ _ err _ <- items] of
      (err:_) -> pure $ Left $ ProtocolError err
      [] -> pure $ Right ()