content: {-# LANGUAGE RecordWildCards, OverloadedStrings, DeriveDataTypeable, NamedFieldPuns, PatternGuards #-}
  
  -- TODO: the code sends an extra query after getting SUCCESS_ATOM when doing e.g. (expr 1)
  
  module Database.RethinkDB.Network (
    RethinkDBHandle(..),
    connect,
    close,
    use,
    runQLQuery,
    Cursor(..),
    makeCursor,
    next,
    nextBatch,
    collect,
    collect',
    nextResponse,
    Response(..),
    ErrorCode(..),
    RethinkDBError(..),
    RethinkDBConnectionError(..),
    More,
    noReplyWait,
    each,
    serverInfo
    ) where
  
  import Control.Monad (when, forever, forM_)
  import Data.Typeable (Typeable)
  import Network (HostName)
  import Network.Socket (
    socket, Family(AF_INET, AF_INET6), SocketType(Stream), setSocketOption,
    SocketOption(NoDelay, KeepAlive), Socket, AddrInfo(AddrInfo, addrAddress, addrFamily))
  import qualified Network.Socket as Socket
  import Network.BSD (getProtocolNumber)
  import Network.Socket.ByteString.Lazy (sendAll)
  import Network.Socket.ByteString (recv)
  import Data.ByteString.Lazy (ByteString)
  import qualified Data.ByteString.Lazy as B
  import qualified Data.ByteString.UTF8 as BS (fromString)
  import qualified Data.ByteString as BS
  import Control.Concurrent (
    writeChan, MVar, Chan, modifyMVar, takeMVar, forkIO, readChan,
    myThreadId, newMVar, ThreadId, newChan, killThread,
    newEmptyMVar, putMVar, mkWeakMVar)
  import Control.Exception (catch, Exception, throwIO, SomeException(..), bracketOnError)
  import Data.IORef (IORef, newIORef, atomicModifyIORef', readIORef, writeIORef)
  import Data.Map (Map)
  import qualified Data.Map as M
  import Data.Maybe (fromMaybe, listToMaybe, isNothing)
  import Control.Monad.Fix (fix)
  import System.IO.Unsafe (unsafeInterleaveIO)
  import System.Mem.Weak (finalize)
  import Data.Binary.Get (runGet, getWord32le, getWord64le)
  import Data.Binary.Put (runPut, putWord32le, putWord64le, putLazyByteString)
  import Data.Word (Word64, Word32)
  import qualified Data.HashMap.Strict as HM
  
  import Database.RethinkDB.Wire
  import Database.RethinkDB.Wire.Response as Response
  import Database.RethinkDB.Wire.Query as Query
  import Database.RethinkDB.Wire.VersionDummy as Protocol
  import Database.RethinkDB.Types
  import Database.RethinkDB.Datum
  import Database.RethinkDB.ReQL (
    Term, Backtrace, convertBacktrace, WireQuery(..),
    WireBacktrace(..), Term(..), Frame(..),
    TermAttribute(..))
  import Data.Foldable (toList)
  
  -- $setup
  --
  -- Get the doctests ready
  --
  -- >>> import qualified Database.RethinkDB as R
  -- >>> import Database.RethinkDB.NoClash
  -- >>> h' <- unsafeInterleaveIO $ connect "localhost" 28015 def
  -- >>> let h = use "doctests" h'
  
  type Token = Word64
  
  -- | A connection to the database server
  data RethinkDBHandle = RethinkDBHandle {
    rdbSocket :: Socket,
    rdbWriteLock :: MVar (Maybe SomeException),
    rdbToken :: IORef Token, -- ^ The next token to use
    rdbDatabase :: Database,  -- ^ The default database
    rdbWait :: IORef (Map Token (Chan Response, Term, IO ())),
    rdbThread :: ThreadId
    }
  
  data Cursor a = Cursor {
    cursorMBox :: MVar Response,
    cursorBuffer :: MVar (Either RethinkDBError ([Datum], Bool)),
    cursorMap :: Datum -> IO a }
  
  instance Functor Cursor where
    fmap f Cursor{ .. } = Cursor { cursorMap = fmap f . cursorMap, .. }
  
  instance Show RethinkDBHandle where
    show RethinkDBHandle{ rdbSocket } = "RethinkDB Connection " ++ show rdbSocket
  
  newToken :: RethinkDBHandle -> IO Token
  newToken RethinkDBHandle{rdbToken} =
    atomicModifyIORef' rdbToken $ \x -> (x+1, x)
  
  data RethinkDBConnectionError =
    RethinkDBConnectionError String
    deriving (Show, Typeable)
  instance Exception RethinkDBConnectionError
  
  getAddrFamily :: AddrInfo -> Family
  getAddrFamily addrInfo = case addrInfo of
      AddrInfo { addrFamily = AF_INET6 } -> AF_INET6
      _                                  -> AF_INET
  
  connectTo :: HostName -> Integer -> IO Socket
  connectTo host port = do
    h <- Socket.getAddrInfo Nothing (Just host) (Just $ show port)
    let addrI = head h
    let addrF = getAddrFamily addrI
    proto <- getProtocolNumber "tcp"
    bracketOnError (socket addrF Stream proto) Socket.close $ \sock -> do
      Socket.connect sock (addrAddress addrI)
      setSocketOption sock NoDelay 1
      setSocketOption sock KeepAlive 1
      return sock
  
  -- | Create a new connection to the database server
  --
  -- /Example:/ connect using the default port with no passphrase (/note:/ IPv4 and IPv6 supported)
  --
  -- >>> h <- connect "localhost" 28015 Nothing
  --
  -- > >>> h <- connect "::1" 28015 Nothing
  
  connect :: HostName -> Integer -> Maybe String -> IO RethinkDBHandle
  connect host port mauth = do
    let auth = B.fromChunks . return . BS.fromString $ fromMaybe "" mauth
    s <- connectTo host port
    sendAll s $ runPut $ do
      putWord32le magicNumber
      putWord32le (fromIntegral $ B.length auth)
      putLazyByteString auth
      putWord32le $ fromIntegral $ toWire Protocol.JSON
    res <- sGetNullTerminatedString s
    when (res /= "SUCCESS") $ throwIO (RethinkDBConnectionError $ show res)
    r <- newIORef 1
    let db' = Database "test"
    wlock <- newMVar Nothing
    waits <- newIORef M.empty
    let rdb = RethinkDBHandle s wlock r db' waits
    tid <- forkIO $ readResponses rdb
    return $ rdb tid
  
  recvAll :: Socket -> Int -> IO ByteString
  recvAll s n_ = go [] n_ where
    go acc 0 = return $ B.fromChunks $ reverse acc
    go acc n = do
      d <- recv s n
      if BS.null d
        then throwIO $ RethinkDBConnectionError "Connection closed unexpectedly"
        else go (d : acc) (n - BS.length d)
  
  sGetNullTerminatedString :: Socket -> IO ByteString
  sGetNullTerminatedString s = go [] where
    go acc = do
      c <- recv s 1
      if BS.null c || c == BS.pack [0]
        then return (B.fromChunks (reverse acc))
        else go (c : acc)
  
  magicNumber :: Word32
  magicNumber = fromIntegral $ toWire V0_4
  
  withSocket :: RethinkDBHandle -> (Socket -> IO a) -> IO a
  withSocket RethinkDBHandle{ rdbSocket, rdbWriteLock } f =
    modifyMVar rdbWriteLock $ \mex ->
    case mex of
      Nothing -> do
        a <- f rdbSocket
        return (Nothing, a)
      Just ex -> throwIO ex
  
  data RethinkDBError = RethinkDBError {
    errorCode :: ErrorCode,
    errorTerm :: Term,
    errorMessage :: String,
    errorBacktrace :: Backtrace
    } deriving (Typeable)
  
  instance Exception RethinkDBError
  
  instance Show RethinkDBError where
    show (RethinkDBError code term message backtrace) =
      show code ++ ": " ++ message ++
      if term == Datum Null
         then ""
         else "\n" ++ indent ("in " ++ show (annotate backtrace term))
      where
        indent = (\x -> case x of [] -> []; _ -> init x) . unlines . map ("  "++) . lines 
        annotate :: Backtrace -> Term -> Term
        annotate (x : xs) t | Just new <- inside x t (annotate xs) = new
        annotate _ t = Note "HERE" t
        inside (FramePos n) (Term tt a oa) f
          | n < length a = Just $ Term tt (take n a ++ [f (a!!n)] ++ drop (n+1) a) oa
        inside (FrameOpt k) (Term tt a oa) f
          | Just (before, v, after) <- extract k oa =
            Just $ Term tt a $ before ++ [TermAttribute k (f v)] ++ after
        inside _ _ _ = Nothing
        extract _ [] = Nothing
        extract k (TermAttribute kk v : xs) | k == kk = Just ([], v, xs)
        extract k (x:xs) =
          case extract k xs of
            Nothing -> Nothing
            Just (a,b,c) -> Just (x:a,b,c)
  
  -- | The response to a query
  data Response =
    ResponseError RethinkDBError |
    ResponseSingle Datum |
    ResponseBatch (Maybe More) [Datum]
  
  data More = More {
    _moreFeed :: Bool,
    _moreHandle :: RethinkDBHandle, 
    _moreToken :: Token
    }
  
  data ErrorCode =
    ErrorBrokenClient |
    ErrorBadQuery |
    ErrorRuntime |
    ErrorUnexpectedResponse
  
  instance Show ErrorCode where
    show ErrorBrokenClient = "RethinkDB: Broken client error"
    show ErrorBadQuery = "RethinkDB: Malformed query error"
    show ErrorRuntime = "RethinkDB: Runtime error"
    show ErrorUnexpectedResponse = "RethinkDB: Unexpected response"
  
  instance Show Response where
    show (ResponseError RethinkDBError {..}) =
      show errorCode ++ ": " ++
      show errorMessage ++ " (" ++
      show errorBacktrace ++ ")"
    show (ResponseSingle datum) = show datum
    show (ResponseBatch _more batch) = show batch
  
  newtype WireResponse = WireResponse { _responseDatum :: Datum }
  
  convertResponse :: RethinkDBHandle -> Term -> Token -> WireResponse -> Response
  convertResponse h q t (WireResponse (Object o)) = let
    type_    = o .? "t" >>= fromWire
    results :: Maybe [Datum]
    results  = o .? "r"
    bt       = o .? "b" --> maybe [] (convertBacktrace . WireBacktrace)
    -- _profile = o .? "p" -- TODO
    atom :: Maybe Datum
    atom     = case results of Just [single] -> Just single; _ -> Nothing
    m .? k   = HM.lookup k m >>= resultToMaybe . fromDatum
    (-->) = flip ($)
    e = fromMaybe "" $ resultToMaybe . fromDatum =<< listToMaybe =<< results
    _ <!< Nothing = ResponseError $ RethinkDBError ErrorUnexpectedResponse q e bt
    f <!< (Just a) = f a
    in case type_ of
    Just SUCCESS_ATOM -> ResponseSingle <!< atom
    Just Response.SERVER_INFO -> ResponseSingle <!< atom
    Just SUCCESS_PARTIAL -> ResponseBatch (Just $ More False h t) <!< results
    Just SUCCESS_SEQUENCE -> ResponseBatch Nothing <!< results
    Just CLIENT_ERROR -> ResponseError $ RethinkDBError ErrorBrokenClient q e bt
    Just COMPILE_ERROR -> ResponseError $ RethinkDBError ErrorBadQuery q e bt
    Just RUNTIME_ERROR -> ResponseError $ RethinkDBError ErrorRuntime q e bt
    Just WAIT_COMPLETE -> ResponseSingle (toDatum True)
    Nothing -> ResponseError $ RethinkDBError ErrorUnexpectedResponse q e bt
  
  convertResponse _ q _ (WireResponse json) =
    ResponseError $
    RethinkDBError ErrorUnexpectedResponse q ("Response is not a JSON object: " ++ show json) []
  
  runQLQuery :: RethinkDBHandle -> WireQuery -> Term -> IO (MVar Response)
  runQLQuery h query term = do
    tok <- newToken h
    let noReply = isNoReplyQuery query
    mbox <- if noReply
            then newEmptyMVar
            else addMBox h tok term
    sendQLQuery h tok query
    when noReply $ putMVar mbox $ ResponseSingle $ Null
    return mbox
  
  isNoReplyQuery :: WireQuery -> Bool
  isNoReplyQuery (WireQuery (Array v)) |
    [_type, _term, (Object optargs)] <- toList v,
    Just (Bool True) <- HM.lookup "noreply" optargs =
      True
  isNoReplyQuery _ = False
  
  addMBox :: RethinkDBHandle -> Token -> Term -> IO (MVar Response)
  addMBox h tok term = do
    chan <- newChan
    mbox <- newEmptyMVar
    weak <- mkWeakMVar mbox $ do
      closeToken h tok -- TODO: don't close if already closed
      atomicModifyIORef' (rdbWait h) $ \mboxes ->
        (M.delete tok mboxes, ())
    atomicModifyIORef' (rdbWait h) $ \mboxes ->
      (M.insert tok (chan, term, finalize weak) mboxes, ())
    _ <- forkIO $ fix $ \loop -> do
      response <- readChan chan
      putMVar mbox response
      when (not $ isLastResponse response) $ do
        nextResponse response
        loop
    return mbox
  
  sendQLQuery :: RethinkDBHandle -> Token -> WireQuery -> IO ()
  sendQLQuery h tok query = do
    let queryS = encode $ queryJSON query
    withSocket h $ \s -> do
      sendAll s $ runPut $ do
        putWord64le tok
        putWord32le (fromIntegral $ B.length queryS)
        putLazyByteString queryS 
  
  data RethinkDBReadError =
    RethinkDBReadError SomeException
    deriving (Show, Typeable)
  instance Exception RethinkDBReadError
  
  readResponses :: (ThreadId -> RethinkDBHandle) -> IO ()
  readResponses h' = do
    tid <- myThreadId
    let h = h' tid
    let handler e@SomeException{} = do
          Socket.close $ rdbSocket h
          modifyMVar (rdbWriteLock h) $ \_ -> return (Just e, ())
          writeIORef (rdbWait h) M.empty
    flip catch handler $ forever $ readSingleResponse h
  
  readSingleResponse :: RethinkDBHandle -> IO ()
  readSingleResponse h = do
    tokenString <- recvAll (rdbSocket h) 8
    when (B.length tokenString /= 8) $
      throwIO $ RethinkDBConnectionError "RethinkDB connection closed unexpectedly"
    let token = runGet getWord64le tokenString
    header <- recvAll (rdbSocket h) 4
    when (B.length header /= 4) $
      throwIO $ RethinkDBConnectionError "RethinkDB connection closed unexpectedly"
    let replyLength = runGet getWord32le header
    rawResponse <- recvAll (rdbSocket h) (fromIntegral replyLength)
    let parsedResponse = eitherDecode rawResponse
    case parsedResponse of
      Left errMsg -> do
        -- TODO: don't give up on the connection, only share the error message with the MVar
        fail errMsg
      Right response -> dispatch token $ WireResponse response
  
    where
    dispatch tok response = do
      mboxes <- readIORef $ rdbWait h
      case M.lookup tok mboxes of
        Nothing -> return ()
        Just (mbox, term, closetok) -> do
          let convertedResponse = convertResponse h term tok response
          writeChan mbox convertedResponse
          when (isLastResponse convertedResponse) $ closetok
  
  isLastResponse :: Response -> Bool
  isLastResponse ResponseError{} = True
  isLastResponse ResponseSingle{} = True
  isLastResponse (ResponseBatch (Just _) _) = False
  isLastResponse (ResponseBatch Nothing _) = True
   
  -- | Set the default database
  --
  -- The new handle is an alias for the old one. Calling close on either one
  -- will close both.
  use :: Database -> RethinkDBHandle -> RethinkDBHandle
  use db' h = h { rdbDatabase = db' }
  
  -- | Close an open connection
  close :: RethinkDBHandle -> IO ()
  close h@RethinkDBHandle{ rdbSocket, rdbThread } = do
    noReplyWait h
    killThread rdbThread
    Socket.close rdbSocket
  
  closeToken :: RethinkDBHandle -> Token -> IO ()
  closeToken h tok = do
    let query = WireQuery $ toDatum [toWire STOP]
    sendQLQuery h tok query
  
  nextResponse :: Response -> IO ()
  nextResponse (ResponseBatch (Just (More _ h tok)) _) = do
    let query = WireQuery $ toDatum [toWire CONTINUE]
    sendQLQuery h tok query
  nextResponse _ = return ()
  
  makeCursor :: MVar Response -> IO (Cursor Datum)
  makeCursor cursorMBox = do
    cursorBuffer <- newMVar (Right ([], False))
    return Cursor{..}
    where cursorMap = return . id
  
  -- | Get the next value from a cursor
  next :: Cursor a -> IO (Maybe a)
  next c@Cursor{ .. } = modifyMVar cursorBuffer $ fix $ \loop mbuffer ->
    case mbuffer of
      Left err -> throwIO err
      Right ([], True) -> return (Right ([], True), Nothing)
      Right (x:xs, end) -> do x' <- cursorMap x; return $ (Right (xs, end), Just x')
      Right ([], False) -> cursorFetchBatch c >>= loop
  
  -- | Get the next batch from a cursor
  nextBatch :: Cursor a -> IO [a]
  nextBatch c@Cursor{ .. } = modifyMVar cursorBuffer $ fix $ \loop mbuffer ->
    case mbuffer of
      Left err -> throwIO err
      Right ([], True) -> return (Right ([], True), [])
      Right (xs@(_:_), end) -> do
        xs' <- mapM cursorMap xs
        return $ (Right ([], end), xs')
      Right ([], False) -> cursorFetchBatch c >>= loop
  
  cursorFetchBatch :: Cursor a -> IO (Either RethinkDBError ([Datum], Bool))
  cursorFetchBatch c = do
    response <- takeMVar (cursorMBox c)
    case response of
      ResponseError e -> return $ Left e
      ResponseBatch more datums -> return $ Right (datums, isNothing more)
      ResponseSingle (Array a) -> return $ Right (toList a, True)
      ResponseSingle _ ->
        return $ Left $ RethinkDBError ErrorUnexpectedResponse (Datum Null)
        "Expected a stream or an array but got a datum" []
  
  -- | A lazy stream of all the elements in the cursor
  collect :: Cursor a -> IO [a]
  collect c = fix $ \loop -> do
      b <- nextBatch c
      case b of
        [] -> return []
        xs -> do
          ys <- unsafeInterleaveIO $ loop
          return $ xs ++ ys
  
  -- | A strict version of collect
  collect' :: Cursor a -> IO [a]
  collect' c = fix $ \loop -> do
      b <- nextBatch c
      case b of
        [] -> return []
        xs -> do
          ys <- loop
          return $ xs ++ ys
  
  -- | Wait for NoReply queries to complete on the server
  --
  -- >>> () <- runOpts h [NoReply] $ table "users" # get "bob" # update (\row -> merge row ["occupation" := "teacher"])
  -- >>> noReplyWait h
  noReplyWait :: RethinkDBHandle -> IO ()
  noReplyWait h = do
    m <- runQLQuery h (WireQuery $ toDatum [toWire NOREPLY_WAIT]) (Datum Null)
    _ <- takeMVar m
    return ()
  
  each :: Cursor a -> (a -> IO b) -> IO ()
  each cursor f = do
    batch <- nextBatch cursor
    if null batch
      then return ()
      else do
        forM_ batch f
        each cursor f
  
  -- | Get information about the server
  serverInfo :: RethinkDBHandle -> IO Datum
  serverInfo h = do
    m <- runQLQuery h (WireQuery $ toDatum [toWire Query.SERVER_INFO]) (Datum Null)
    response <- takeMVar m
    case response of
      ResponseError e -> throwIO e
      ResponseBatch _ _ -> throwIO (RethinkDBError ErrorUnexpectedResponse (Datum Null) "" [])
      ResponseSingle d -> return d
  