{-# LANGUAGE ExistentialQuantification, OverloadedStrings #-}

-- | Currently, only Exact and Range Queries via Secondary Indexes are
-- implemented.

module Network.Riak.HTTP.MapReduce
    (Input (..),
     -- * Specifying Functions
     MRFunction (..), Arg (Arg),
     -- * Phases
     Phase (..), LinkTag,
     -- * Running MapReduce
     mapReduce, Timeout (..)
     )
where

import Network.Riak.HTTP.Internal
import Network.Riak.HTTP.Types
import Data.Aeson
import Data.Aeson.Types (Pair)
import Data.Attoparsec.Number (Number (..))
import Data.ByteString.Lazy (ByteString)
import Data.String.Class (toString, toText)
import Network.HTTP
import Network.URI
import qualified Network.Riak.HTTP.Types as Types
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BSL
import qualified Data.ByteString.Char8 as BSC
import qualified Data.ByteString.Lazy.Char8 as BLC

data Input =
    BucketInput Bucket [KeyFilter] |
    BucketKeyInput [(Bucket, Key)] |
    SearchInput Bucket SearchQuery |
    IndexInput Bucket IndexKey IndexValue |
    RangeIndexInput Bucket IndexKey IndexValue IndexValue
    deriving (Show)

data KeyFilter = KeyFilter deriving (Show)

type SearchQuery = String

data MRFunction =
    JSNamed {
        functionModule :: String,
        functionName :: String
    } |
    -- | JavaScript function declared in the content of a Riak object
    JSObject Bucket Key |
    JSCode String |
    ErlangNamed {
        functionModule :: String,
        functionName :: String
    }
    deriving (Show)

type LinkTag = String

data Phase =
    MapPhase {
        mapper :: MRFunction,
        keep :: Bool,
        arg :: Maybe Arg
    } |
    ReducePhase {
        reducer :: MRFunction,
        keep :: Bool,
        arg :: Maybe Arg
    } |
    LinkPhase {
        bucket :: Bucket,
        tag :: LinkTag,
        keep :: Bool
    }
    deriving (Show)

instance ToJSON Phase where
    toJSON (MapPhase func keep arg) = object
        ["map" .= object (["keep" .= keep] ++ argJSONAttr arg ++
            functionJSONAttrs func)]
    toJSON (ReducePhase func keep arg) = object
        ["reduce" .= object (["keep" .= keep] ++ argJSONAttr arg ++
            functionJSONAttrs func)]
    toJSON (LinkPhase bucket tag keep) = object ["link" .=
        ["bucket" .= bucket, "tag" .= tag, "keep" .= keep]]

data Arg = forall a. ToJSON a => Arg a

instance ToJSON Arg where toJSON (Arg x) = toJSON x

instance Show Arg where
    show (Arg x) = toString $ encode x

data Timeout = DefaultTimeout | Timeout Integer deriving (Show)

mapReduce :: Client -> Input -> [Phase] -> Timeout
          -> IO (Either String ByteString)
mapReduce client input phases timeout = do
    let uri = buildMRURI client
    let requestBody = buildReqBody input phases timeout'
    let headers = [mkHeader HdrContentType "application/json",
                   mkHeader HdrContentLength (show $ BLC.length requestBody)]
    let request = Request { rqURI = uri, rqMethod = POST, rqHeaders = headers,
                            rqBody = requestBody }
    result <- tryHTTPRequest request
    return $ case result of
        Left err -> Left err
        Right response -> case rspCode response of
            (2, 0, 0) -> Right $ rspBody response
            (4, 0, 0) -> Left "An invalid job was submitted"
            (5, 0, 0) -> Left ("There was an error in processing a map or\
                \ reduce function : " ++ toString (rspBody response))
            (5, 0, 3) -> Left "The job timed out before it could complete"
    where timeout' = case timeout of
            DefaultTimeout -> 60000 -- http://wiki.basho.com/MapReduce.html
            Timeout t -> t

-- Clearly, this needs some cleanup.
buildReqBody :: (ToJSON a, ToJSON b) => Input -> a -> b -> ByteString
buildReqBody (RangeIndexInput bucket ik start end) phases t = encode $ object
        ["inputs" .= object
            ["bucket" .= bucket,
             "index" .= (ik ++ indexType),
             "start" .= startValue,
             "end" .= endValue],
         "query" .= phases,
         "timeout" .= t
        ]
    where
        (indexType, startValue, endValue) = case (start, end) of
            (BinValue s, BinValue e) -> ("_bin",
                String $ toText s, String $ toText e)
            (IntValue s, IntValue e) -> ("_int",
                Number $ fromIntegral s, Number $ fromIntegral e)
            _ -> error "Start and end values must be the same type"
buildReqBody (BucketKeyInput bucketKeys) phases t = encode $ object [
        "inputs" .= map (\(b, k) -> [b, k]) bucketKeys, "query" .= phases,
        "timeout" .= t
    ]
buildReqBody (IndexInput bucket ik val) phases t = encode $ object [
        "inputs" .= object [
            "bucket" .= bucket,
            "index" .= (ik ++ indexType),
            "key" .= val'
        ],
        "query" .= phases,
        "timeout" .= t
    ]
    where
        (indexType, val') = case val of
            BinValue v -> ("_bin", String $ toText v)
            IntValue v -> ("_int", Number $ fromIntegral v)
buildReqBody _ _ _ = error "Not yet implemented :("

buildMRURI client = nullURI {
    uriScheme = "http:",
    uriAuthority = Just URIAuth {
        uriUserInfo = "",
        uriRegName = Types.host client,
        uriPort = ":" ++ (Types.port client)
    },
    uriPath = "/mapred"
}

functionJSONAttrs :: MRFunction -> [Pair]
functionJSONAttrs func = case func of
    ErlangNamed m n -> ["language" .= e, "module" .= m, "function" .= n]
    JSObject b k -> ["language" .= j, "bucket" .= b, "key" .= k]
    JSNamed m n -> ["language" .= j, "module" .= m, "function" .= n]
    JSCode code -> ["language" .= j, "source" .= code]
    where
        e = "erlang" :: String
        j = "javascript" :: String

-- | Because it's possible to pattern-match in Erlang on @none@ on the
-- argument parameter with the expectation that no @arg@ parameter has
-- been passed in, even though that's probably the wrong thing to do. For
-- safety, only put @arg@ in the JSON if we passed in @Just x@.
argJSONAttr :: Maybe Arg -> [Pair]
argJSONAttr (Just arg) = ["arg" .= arg]
argJSONAttr Nothing = []
