4 files changed, 51 insertions, 3 deletions

diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index a1934014..e1104cb9 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -72,7 +72,6 @@ import Data.Either
 import Data.List as L
 import Data.Monoid
 import Data.Text as T
-import Data.BEncode (BValue)
 import qualified Data.Set as Set
          ;import Data.Set (Set)
 import Network
@@ -85,13 +84,20 @@ import Network.KRPC hiding (Options, def)
 import Network.KRPC.Message (ReflectedIP(..))
 import Network.KRPC.Manager (QueryFailure(..))
 import Data.Torrent
-import Network.BitTorrent.Address
 import Network.BitTorrent.DHT.Message
 import Network.BitTorrent.DHT.Routing as R
 import Network.BitTorrent.DHT.Session
 import Control.Concurrent.STM
 import qualified Network.BitTorrent.DHT.Search as Search
+#ifdef VERSION_bencoding
+import Network.BitTorrent.Address
+import Data.BEncode (BValue)
 import Network.DHT.Mainline
+#else
+import Network.BitTorrent.Address hiding (NodeId)
+import Data.ByteString (ByteString)
+import Data.Tox
+#endif
 
 {-----------------------------------------------------------------------
 --  Handlers
@@ -99,12 +105,23 @@ import Network.DHT.Mainline
 
 nodeHandler :: ( Address ip
                , KRPC (Query a) (Response b)
+#ifdef VERSION_bencoding
                , Envelope (Query a) (Response b) ~ BValue )
+#else
+               , Envelope (Query a) (Response b) ~ ByteString )
+#endif
            => (NodeAddr ip -> a -> DHT ip b) -> NodeHandler ip
+#ifdef VERSION_bencoding
 nodeHandler action = handler mainline $ \ sockAddr qry -> do
   let remoteId  = queringNodeId qry
       read_only = queryIsReadOnly qry
       q         = queryParams qry
+#else
+nodeHandler action = handler (error "TODO TOX Messaging") $ \ sockAddr qry -> do
+  let remoteId  = msgClient qry
+      read_only = False
+      q         = msgPayload qry
+#endif
   case fromSockAddr sockAddr of
     Nothing    -> throwIO BadAddress
     Just naddr -> do
diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index db8e7cff..20dba595 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -107,7 +107,9 @@ import Data.Torrent as Torrent
 import Network.KRPC as KRPC hiding (Options, def)
 import qualified Network.KRPC as KRPC (def)
 import Network.KRPC.Message (KMessageOf)
+#ifdef VERSION_bencoding
 import Data.BEncode (BValue)
+#endif
 import Network.BitTorrent.Address
 import Network.BitTorrent.DHT.ContactInfo (PeerStore)
 import qualified Network.BitTorrent.DHT.ContactInfo as P
@@ -314,7 +316,11 @@ instance MonadLogger (DHT ip) where
     logger <- asks loggerFun
     liftIO $ logger loc src lvl (toLogStr msg)
 
+#ifdef VERSION_bencoding
 type NodeHandler ip = Handler (DHT ip) KMessageOf BValue
+#else
+type NodeHandler ip = Handler (DHT ip) KMessageOf ByteString
+#endif
 
 -- | Run DHT session. You /must/ properly close session using
 -- 'closeNode' function, otherwise socket or other scarce resources may
diff --git a/src/Network/KRPC/Manager.hs b/src/Network/KRPC/Manager.hs
index b1e93101..58ac7674 100644
--- a/src/Network/KRPC/Manager.hs
+++ b/src/Network/KRPC/Manager.hs
@@ -151,7 +151,11 @@ data Manager h = Manager
   , listenerThread     :: !(MVar ThreadId)
   , transactionCounter :: {-# UNPACK #-} !TransactionCounter
   , pendingCalls       :: {-# UNPACK #-} !PendingCalls
+#ifdef VERSION_bencoding
   , handlers           :: [Handler h KMessageOf BValue]
+#else
+  , handlers           :: [Handler h KMessageOf BC.ByteString]
+#endif
   }
 
 -- | A monad which can perform or handle queries.
@@ -185,7 +189,11 @@ sockAddrFamily (SockAddrCan _        ) = AF_CAN
 -- run 'listen'.
 newManager :: Options                       -- ^ various protocol options;
            -> SockAddr                      -- ^ address to listen on;
+#ifdef VERSION_bencoding
            -> [Handler h KMessageOf BValue] -- ^ handlers to run on incoming queries.
+#else
+           -> [Handler h KMessageOf BC.ByteString] -- ^ handlers to run on incoming queries.
+#endif
            -> IO (Manager h)                -- ^ new rpc manager.
 newManager opts @ Options {..} servAddr handlers = do
     validateOptions opts
@@ -218,7 +226,11 @@ isActive Manager {..} = liftIO $ isBound sock
 
 -- | Normally you should use Control.Monad.Trans.Resource.allocate
 -- function.
+#ifdef VERSION_bencoding
 withManager :: Options -> SockAddr -> [Handler h KMessageOf BValue]
+#else
+withManager :: Options -> SockAddr -> [Handler h KMessageOf BC.ByteString]
+#endif
             -> (Manager h -> IO a) -> IO a
 withManager opts addr hs = bracket (newManager opts addr hs) closeManager
 
@@ -420,14 +432,22 @@ handler msging body = (name, wrapper)
         Right a -> Right . seal <$> body addr a
 
 runHandler :: MonadKRPC h m
+#ifdef VERSION_bencoding
            => HandlerBody h KMessageOf BValue -> SockAddr -> KQuery -> m KResult
+#else
+           => HandlerBody h KMessageOf BC.ByteString -> SockAddr -> KQuery -> m KResult
+#endif
 runHandler h addr m = Lifted.catches wrapper failbacks
   where
     signature = querySignature (queryMethod m) (queryId m) addr
 
     wrapper = do
       $(logDebugS) "handler.quered" signature
+#ifdef VERSION_bencoding
       result <- liftHandler (h addr (Q m))
+#else
+      result <- liftHandler (h addr m)
+#endif
 
       case result of
         Left msg -> do
diff --git a/src/Network/KRPC/Message.hs b/src/Network/KRPC/Message.hs
index 19f9fc9e..2f5f6729 100644
--- a/src/Network/KRPC/Message.hs
+++ b/src/Network/KRPC/Message.hs
@@ -34,7 +34,9 @@ module Network.KRPC.Message
        , unknownMessage
 
          -- * Query
+#ifdef VERSION_bencoding
        , KQueryOf(..)
+#endif
        , KQuery
 #ifndef VERSION_bencoding
        , queryArgs
@@ -44,7 +46,9 @@ module Network.KRPC.Message
        , MethodName
 
          -- * Response
+#ifdef VERSION_bencoding
        , KResponseOf(..)
+#endif
        , KResponse
        , ReflectedIP(..)
 
@@ -357,5 +361,6 @@ instance BEncode KMessage where
     <|> E <$> fromBEncode b
     <|> decodingError "KMessage: unknown message or message tag"
 #else
-type KMessage = Tox.Message
+type KMessageOf = Tox.Message
+type KMessage = KMessageOf B.ByteString
 #endif