{-# LANGUAGE GeneralizedNewtypeDeriving, DeriveDataTypeable, OverloadedStrings, FlexibleInstances, FlexibleContexts #-}

-- | This module is a very thin wrapper around "Network.HTTP.Enumerator" using the aeson package to parse
--   and encode JSON.  The Couch DB HTTP API is the best place to learn about how to use this library.
--   <http://wiki.apache.org/couchdb/Complete_HTTP_API_Reference>
--
-- > {-# LANGUAGE OverloadedStrings #-}
-- > import Control.Monad.IO.Class (liftIO)
-- > import Data.Aeson
-- > import qualified Data.ByteString.Lazy as BL
-- > import Data.ByteString.UTF8 (fromString)
-- > import Data.Enumerator (($$), run_)
-- > import qualified Data.Enumerator.List as EL
-- > import Database.CouchDB.Enumerator
-- >
-- > testCouch :: IO ()
-- > testCouch = runCouch "localhost" 5984 "test" $ do
-- >    
-- >    -- Insert some documents.   Note that the dbname passed to withCouchConnection
-- >    -- is prepended to the given path, so this is a put to
-- >    -- http://localhost:5984/test/doc1
-- >    rev1 <- couchPut "doc1" [] $ object [ "foo" .= (3 :: Int), "bar" .= ("abc" :: String) ]
-- >    rev2 <- couchPut "doc2" [] $ object [ "foo" .= (7 :: Int), "baz" .= (145 :: Int) ]
-- >
-- >    -- Load the document and print it out
-- >    couchGet "doc1" [] >>= liftIO . BL.putStrLn . encode . Object
-- >
-- >    -- Overwite the document.  We supply the revision, otherwise Couch DB would give an error.
-- >    -- (The revision could also have been passed in the query arguments.)
-- >    rev3 <- couchPut "doc1" [] $ object [ "foo" .= (10 :: Int)
-- >                                        , "bar" .= ("def" :: String)
-- >                                        , "_rev" .= rev1 
-- >                                        ]
-- >
-- >    -- Create a view
-- >    couchPut_ "_design/testdesign" [] $ 
-- >        object [ "language" .= ("javascript" :: String)
-- >               , "views"    .= object [ "myview" .= object [ "map" .=
-- >                    ("function(doc) { emit(doc.foo, doc); }" :: String)
-- >                    ]]
-- >               ]
-- >
-- >    -- Read from the view using couchGet and print it out.
-- >    couchGet "_design/testdesign/_view/myview" [] >>= liftIO . BL.putStrLn . encode . Object
-- >    couchGet "_design/testdesign/_view/myview" [(fromString "key", Just $ fromString "10")]
-- >            >>= liftIO . BL.putStrLn . encode . Object
-- >
-- >    -- Read the view using couchView and print it out.
-- >    run_ $ couchView "testdesign/_view/myview" [] $$
-- >        EL.foldM (\_ o -> liftIO $ BL.putStrLn $ encode $ Object o) ()
-- >    run_ $ couchView "testdesign/_view/myview" [(fromString "key", Just $ fromString "10")] $$
-- >        EL.foldM (\_ o -> liftIO $ BL.putStrLn $ encode $ Object o) ()
-- >
-- >    -- Delete the objects
-- >    couchDelete "doc1" rev3
-- >    couchDelete "doc2" rev2
module Database.CouchDB.Enumerator(
    -- * Couch DB Connection
      CouchConnection(..)
    , runCouch
    , withCouchConnection
    , CouchError(..)
    , MonadCouch(..)

    -- * Accessing Couch DB
    , Path
    , Revision
    , couchGet
    , couchPut
    , couchPut_
    , couchDelete
    , couchView
    , extractViewValue
    , couch

    -- * Connection Pooling
    -- $pool

    -- * Yesod Integration
    -- $yesod
) where

import           Control.Applicative
import           Control.Exception (Exception, throw, bracket)
import           Control.Monad.IO.Class (MonadIO, liftIO)
import           Control.Monad.Trans.Class (MonadTrans, lift)
import           Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp)
import           Control.Monad.Trans.Reader (ReaderT, ask, runReaderT)
import qualified Data.Aeson as A
import qualified Data.Aeson.Encode as AE
import           Data.Attoparsec
import           Data.Attoparsec.Enumerator (iterParser)
import qualified Data.ByteString as B
import qualified Data.ByteString.UTF8 as BU8
import           Data.Enumerator hiding (map)
import qualified Data.HashMap.Lazy as M
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import           Data.Typeable (Typeable)
import qualified Network.HTTP.Enumerator as H
import qualified Network.HTTP.Types as HT

-- | Represents a connection to a single Couch DB Database.  
--
--   A connection contains a 'H.Manager' and reuses it for multiple requests, which means a 
--   single open HTTP connection to CouchDB will be kept around until the manager is closed
--   (http-enumerator will create more connections if needed, it just keeps only one and
--   closes the rest.)  See the Pool section for more information.
data CouchConnection = CouchConnection {
      host      :: B.ByteString
    , port      :: Int
    , manager   :: H.Manager
    , dbname    :: String
}

-- | Connect to a CouchDB database, call the supplied function, and then close the connection.
-- 
--   If you create your own instance of 'MonadCouch' instead of using 'runCouch', this function
--   will help you create the 'CouchConnection'.  On the other hand, if you want to implement
--   connection pooling, you will not be able to use withCouchConnection and must create the
--   connection yourself.
withCouchConnection :: (MonadBaseControl IO m) => String    -- ^ host
                                               -> Int       -- ^ port
                                               -> String    -- ^ database name
                                               -> (CouchConnection -> m a) -- ^ function to run
                                               -> m a
withCouchConnection h p db f = liftBaseOp (bracket H.newManager H.closeManager) go
    where go m = f $ CouchConnection (BU8.fromString h) p m db

-- | A Couch DB Error.  If the error comes from http, the http status code is also given.  Non-http errors
--   include things like errors parsing the response.
data CouchError = CouchError (Maybe Int) String
    deriving (Show,Typeable)
instance Exception CouchError

-- | A monad which allows access to the connection.
class MonadIO m => MonadCouch m where
    couchConnection :: m CouchConnection

instance (MonadIO m) => MonadCouch (ReaderT CouchConnection m) where
    couchConnection = ask

-- | Run a sequence of CouchDB actions.
--
--   The functions below to access CouchDB require a 'MonadCouch' instance to access the connection
--   information.  'ReaderT' is an instance of 'MonadCouch', and /runCouch/ runs a sequence of database
--   actions using 'ReaderT'.  See the top of this page for an example using /runCouch/.
--
--   The main reason to not use /runCouch/ is to obtain more control over connection pooling.
--   Also, if your db code is part of a larger monad, it makes sense to just make the larger monad
--   an instance of 'MonadCouch' and skip the intermediate ReaderT, since then performance is
--   improved by eliminating one monad from the final transformer stack.
--
--   This function is a combination of 'withCouchConnection' and 'runReaderT'
runCouch :: (MonadIO m, MonadBaseControl IO m) => String    -- ^ host
                                               -> Int       -- ^ port
                                               -> String    -- ^ database name
                                               -> ReaderT CouchConnection m a -- ^ CouchDB actions
                                               -> m a
runCouch h p d = withCouchConnection h p d . runReaderT

-- | A path to a Couch DB Object.
type Path = String

-- | Represents a revision of a Couch DB Document.
type Revision = T.Text

-- | The most general method of accessing CouchDB.  This is a very thin wrapper around 'H.http'.  Most of the
--   time you should use one of the other access functions, but this function is needed for example to write
--   and read attachments that are not in JSON format.
couch :: MonadCouch m
      => HT.Method                    -- ^ Method
      -> Path                         -- ^ The dbname from the connection is prepended to
                                      --   this path.
      -> HT.Query                     -- ^ Query arguments
      -> Iteratee B.ByteString m a    -- ^ Iteratee to process the response if no error occurs.
      -> H.RequestBody m              -- ^ Body
      -> Iteratee B.ByteString m a
couch m p q i b = Iteratee $ do
    conn <- couchConnection
    let req = H.def { H.method          = m
                    , H.host            = host conn
                    , H.port            = port conn
                    , H.path            = BU8.fromString ("/" ++ dbname conn ++ "/" ++ p)
                    , H.queryString     = q
                    , H.requestBody     = b
                    }
    runIteratee $ H.http req (\s _ -> checkStatus s i) (manager conn)

-- | Load a single object from couch DB.
couchGet :: MonadCouch m
         => Path       -- ^ the dbname is prepended to this string to form the full path.
         -> HT.Query   -- ^ Query arguments.
         -> m A.Object
couchGet p q = do v <- run_ $ couch HT.methodGet p q (iterParser A.json) (H.RequestBodyBS B.empty)
                  either throw return $ valToObj v

-- | Put an object in Couch DB, returning the new Revision.
couchPut :: (MonadCouch m, A.ToJSON a) 
         => Path        -- ^ the dbname is prepended to this string to form the full path.
         -> HT.Query    -- ^ Query arguments.
         -> a           -- ^ The object to store.
         -> m Revision
couchPut p q val = do v <- run_ $ couch HT.methodPut p q (iterParser A.json) body
                      either (liftIO . throw) return (valToObj v >>= objToRev)
    where body = jsonToReqBody val

-- | A version of 'couchPut' which ignores the return value.  This is slightly faster than / _ <- couchPut .../
--   since the JSON parser is not run.
couchPut_ :: (MonadCouch m, A.ToJSON a) 
          => Path       -- ^ the dbname is prepended to this string to form the full path.
          -> HT.Query   -- ^ Query arguments.
          -> a          -- ^ The object to store.
          -> m ()
couchPut_ p q val = run_ $ couch HT.methodPut p q (yield () EOF) body
    where body = jsonToReqBody val

-- | Delete the given revision of the object.
couchDelete :: MonadCouch m
            => Path     -- ^ the dbname is prepended to this string to form the full path. 
            -> Revision
            -> m ()
couchDelete p r = run_ $ couch HT.methodDelete p [("rev", Just $ TE.encodeUtf8 r)]
                             (yield () EOF) (H.RequestBodyBS B.empty)

-- | Load from a Couch DB View.
--
-- While you can use 'couchGet' on a view object, this function combines the
-- incredible power of http-enumerator and attoparsec to allow you to process objects in constant space.
-- As data is read from the network, it is fed into attoparsec.  When attoparsec completes parsing an object
-- it is sent out the enumerator.
--
-- The objects enumerated are the entries in the \"rows\" property of the view result, which
-- means they are not directly the objects you put into the database.  See <http://wiki.apache.org/couchdb/HTTP_view_API>
-- for more information.  The objects inserted into the database are available in the \"value\" entry, and can be extracted
-- with the 'extractViewValue' enumeratee, for example:
--
-- >  couchView "mydesigndoc/_view/myview" [(fromString "key", Just $ fromString "3")] $= extractViewValue
couchView :: MonadCouch m
          => Path      -- ^ \/dbname\/_design\/  is prepended to the given path
          -> HT.Query  -- ^ Query arguments.
          -> Enumerator A.Object m a
couchView p q step = do s <- lift $ run $ couch HT.methodGet ("_design/" ++ p) q (parseView step) (H.RequestBodyBS B.empty)
                        either throwError returnI s


-- | An enumeratee to extract the \"value\" member of JSON objects.
--
-- This is useful to extract the object from the data returned from a view.  For example, Couch DB will return
-- objects that look like the following:
--
-- >    { "id":"64ACF01B05F53ACFEC48C062A5D01D89", "key":null, "value": { some object } }
--
-- and this enumeratee will extract /{some object}/
extractViewValue :: Monad m => Enumeratee A.Object A.Object m a
extractViewValue = mapEither f
    where f v = case M.lookup "value" v of
                   (Just (A.Object o)) -> Right o
                   _                   -> Left $ CouchError Nothing "view does not contain value"

-----------------------------------------------------------------------------------------
--- Helper Code
-----------------------------------------------------------------------------------------

-- | Converts a json object into a 'H.RequestBodyEnumChunked'
jsonToReqBody :: (A.ToJSON a, Monad m) => a -> H.RequestBody m
jsonToReqBody val = H.RequestBodyEnumChunked enum
    where jbuilder = AE.fromValue $ A.toJSON val
          enum (Continue k) = k (Chunks [jbuilder])
          enum step         = returnI step

-- | Check status codes from couch db.
checkStatus :: Monad m => HT.Status -> Iteratee B.ByteString m b 
                       -> Iteratee B.ByteString m b
checkStatus (HT.Status 200 _) i = i
checkStatus (HT.Status 201 _) i = i
checkStatus (HT.Status 202 _) i = i
checkStatus (HT.Status 304 _) i = i
checkStatus (HT.Status c   m) _ = iterParser A.json >>= \v -> throwError $ CouchError (Just c) $ msg v
    where msg v = BU8.toString m ++ reason v
          reason (A.Object v) = case M.lookup "reason" v of
                                    Just (A.String t) -> ":  " ++ T.unpack t
                                    _                 -> ""
          reason _ = []

-- | Convers a value to an object
valToObj :: A.Value -> Either CouchError A.Object
valToObj (A.Object o) = Right o
valToObj _            = Left $ CouchError Nothing "Couch DB did not return an object"

-- | Converts an object to a revision
objToRev :: A.Object -> Either CouchError Revision
objToRev o =  case M.lookup "rev" o of
                (Just (A.String r)) -> Right r
                _  -> Left $ CouchError Nothing "unable to find revision"        

data CommaOrCloseBracket = Comma | CloseBracket

commaOrClose :: Parser CommaOrCloseBracket
commaOrClose = do
    skipWhile (\c -> c /= 44 && c /= 93) <?> "Checking for next comma"
    w <- anyWord8
    if w == 44 then return Comma else return CloseBracket


-- | The main loop of processing the view rows.
viewLoop :: MonadIO m => Enumeratee B.ByteString A.Object m a
viewLoop (Yield a _) = return $ Yield a EOF
viewLoop (Error err) = return $ Error err
viewLoop (Continue k) = do 
    v <- iterParser (A.json <?> "json object")

    vobj <- case v of
             (A.Object o) -> return o
             _            -> throwError $ CouchError Nothing "view entry is not an object"

    step' <- lift $ runIteratee $ k $ Chunks [vobj]

    res <- iterParser (commaOrClose <?> "comma or close")
    case res of
         Comma        -> viewLoop step'
         CloseBracket -> case step' of
                           (Continue k') -> lift $ runIteratee $ k' EOF
                           _             -> return step'

viewStart :: Parser Bool
viewStart = do _ <- string "{\"total_rows\":"
               skipWhile (\x -> x >= 48 && x <= 57)
               _ <- string ",\"offset\":"
               skipWhile (\x -> x >= 48 && x <= 57)
               _ <- string ",\"rows\":["
               (string "]}" >> return False) <|> return True

-- | Enumeratee to parse the data returned by a view.
parseView :: MonadIO m => Enumeratee B.ByteString A.Object m a
parseView (Yield a _)  = return $ Yield a EOF
parseView (Error err)  = return $ Error err
parseView (Continue k) = do b <- iterParser (viewStart <?> "start of view") 
                            if b
                              then viewLoop $ Continue k
                              else lift $ runIteratee $ k EOF


mapEither :: (Exception e, Monad m) => (a -> Either e b) -> Enumeratee a b m c
mapEither f = checkDone (continue . step) where
    step k EOF         = yield (Continue k) EOF
    step k (Chunks xs) = case Prelude.mapM f xs of
                             Left err  -> throwError err
                             Right xs' -> k (Chunks xs') >>== mapEither f

-- $pool
-- The 'H.Manager' stored in the CouchConnection maintains a pool of open connections in an IORef,
-- but keeps a maximum of one open connection per (host,port) pair.  Also, each time 'runCouch' or
-- 'withCouchConnection' is called, a new manager (and thus new connections) is created and destroyed.
--
-- For more precise control over pooling, use the 
-- <http://hackage.haskell.org/package/resource-pool> or
-- <http://hackage.haskell.org/package/pool> packages combined with the
-- 'H.newManager' and 'H.closeManager' functions.
--
-- For example, the following code using the resource-pool package runs a /ReaderT CouchConnection m/ action using a
-- HTTP connection from a pool.
-- 
-- > runPooledCouch :: MonadCatchIO m
-- >                => Pool Manager -> String -> Int -> String -> ReaderT CouchConnection m a -> m a
-- > runPooledCouch p host port dbname c = withResource p $ \m -> do
-- >    runReaderT c $ CouchConnection (BU8.fromString host) port m dbname

-- $yesod
-- Integrating couchdb-enumerator with yesod looks something the way the scaffold sets up the
-- YesodPersist instance.
--
-- > data MyFoundation = MyFoundation 
-- >     { ... (normal yesod stuff in the foundation type)
-- >     , connPool     :: Data.Pool.Pool H.Manager
-- >     , dbLocation   :: B.ByteString
-- >     , databaseName :: String
-- >     }
-- >
-- > newtype CouchDBPersist m a = CouchDBPersist { unCouchDBPersist :: ReaderT CouchConnection m a }
-- >     deriving (Monad, MonadIO, MonadTrans, Functor, Applicative, MonadTransControl,
-- >               MonadBaseControl, MonadPlus, MonadCouch)
-- >
-- > instance YesodPersist MyFoundation where
-- >     type YesodPersistBackend = CouchDBPersist
-- >     runDB r = do pool <- connPool <$> getYesod
-- >                  loc <- dbLocation <$> getYesod
-- >                  db <- databaseName <$> getYesod
-- >                  Data.Pool.withPool' pool $ \m ->
-- >                      runReaderT (unCouchDBPersist r) $ CouchConnection loc 5984 m db
--
-- Then you can write handler code as follows:
--
-- > getFooR :: PersonID -> Handler RepPlain
-- > getFooR p = do
-- >    person <- runDB $ couchGet p []
-- >    return $ RepPlain $ toContent $ Aeson.encode $ maybe Aeson.null person
--
-- Alternatively, you don't need to make your Foundation an instance of YesodPersist, you could
-- supply your own runCouchDB function which is just a version of runDB specialized to your foundation
-- and just use that from the handlers.

-- vim: set expandtab:set tabstop=4: