module Text.Search.Sphinx.Get where

import Data.Binary.Get
import Data.Binary.IEEE754

import Data.Int (Int64)
import Prelude hiding (readList)
import Data.ByteString.Lazy hiding (pack, length, map, groupBy)
import Control.Monad
import qualified Text.Search.Sphinx.Types as T
import Data.Maybe (isJust, fromJust)

import qualified Data.Text.ICU.Convert as ICU

-- Utility functions
getNum :: Get Int
getNum :: Get Int
getNum = Get Word32
getWord32be Get Word32 -> (Word32 -> Get Int) -> Get Int
forall a b. Get a -> (a -> Get b) -> Get b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> Get Int
forall a. a -> Get a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Get Int) -> (Word32 -> Int) -> Word32 -> Get Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word32 -> Int
forall a. Enum a => a -> Int
fromEnum

getFloat :: Get Float
getFloat :: Get Float
getFloat = Get Float
getFloat32be

getNum64 :: Get Int64
getNum64 :: Get Int64
getNum64 = Get Word64
getWord64be Get Word64 -> (Word64 -> Get Int64) -> Get Int64
forall a b. Get a -> (a -> Get b) -> Get b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int64 -> Get Int64
forall a. a -> Get a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int64 -> Get Int64) -> (Word64 -> Int64) -> Word64 -> Get Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word64 -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral

readList :: Get a -> Get [a]
readList Get a
f = do Int
num <- Get Int
getNum
                Int
num Int -> Get a -> Get [a]
forall {a}. Int -> Get a -> Get [a]
`times` Get a
f
times :: Int -> Get a -> Get [a]
times = Int -> Get a -> Get [a]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM

getTxt :: Converter -> Get Text
getTxt Converter
conv = (ByteString -> Text) -> Get ByteString -> Get Text
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM (Converter -> ByteString -> Text
ICU.toUnicode Converter
conv) Get ByteString
getStrStr

getStr :: Get ByteString
getStr = do Int
len <- Get Int
getNum
            Int64 -> Get ByteString
getLazyByteString (Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len)

-- Get a strict 'ByteString'.
getStrStr :: Get ByteString
getStrStr = do Int
len <- Get Int
getNum
               Int -> Get ByteString
getByteString (Int -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len)

getResult :: ICU.Converter -> Get (T.SingleResult)
getResult :: Converter -> Get SingleResult
getResult Converter
conv = do
  Int
statusNum <- Get Int
getNum
  case Int -> QueryStatus
T.toQueryStatus Int
statusNum of
    T.QueryERROR Int
n -> do Text
e <- Converter -> Get Text
getTxt Converter
conv
                         SingleResult -> Get SingleResult
forall a. a -> Get a
forall (m :: * -> *) a. Monad m => a -> m a
return (SingleResult -> Get SingleResult)
-> SingleResult -> Get SingleResult
forall a b. (a -> b) -> a -> b
$ Int -> Text -> SingleResult
T.QueryError Int
statusNum Text
e
    QueryStatus
T.QueryOK      -> Get QueryResult
getResultOk Get QueryResult
-> (QueryResult -> Get SingleResult) -> Get SingleResult
forall a b. Get a -> (a -> Get b) -> Get b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= SingleResult -> Get SingleResult
forall a. a -> Get a
forall (m :: * -> *) a. Monad m => a -> m a
return (SingleResult -> Get SingleResult)
-> (QueryResult -> SingleResult) -> QueryResult -> Get SingleResult
forall b c a. (b -> c) -> (a -> b) -> a -> c
. QueryResult -> SingleResult
T.QueryOk
    QueryStatus
T.QueryWARNING -> do Text
w <- Converter -> Get Text
getTxt Converter
conv
                         Get QueryResult
getResultOk Get QueryResult
-> (QueryResult -> Get SingleResult) -> Get SingleResult
forall a b. Get a -> (a -> Get b) -> Get b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= SingleResult -> Get SingleResult
forall a. a -> Get a
forall (m :: * -> *) a. Monad m => a -> m a
return (SingleResult -> Get SingleResult)
-> (QueryResult -> SingleResult) -> QueryResult -> Get SingleResult
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text -> QueryResult -> SingleResult
T.QueryWarning Text
w)
  where
    getResultOk :: Get QueryResult
getResultOk = do
      [ByteString]
fields     <- Get ByteString -> Get [ByteString]
forall {a}. Get a -> Get [a]
readList Get ByteString
getStr
      [(ByteString, AttrT)]
attrs      <- Get (ByteString, AttrT) -> Get [(ByteString, AttrT)]
forall {a}. Get a -> Get [a]
readList Get (ByteString, AttrT)
readAttrPair
      Int
matchCount <- Get Int
getNum
      Int
id64       <- Get Int
getNum
      [Match]
matches    <- Int
matchCount Int -> Get Match -> Get [Match]
forall {a}. Int -> Get a -> Get [a]
`times` Bool -> [AttrT] -> Converter -> Get Match
readMatch (Int
id64 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (((ByteString, AttrT) -> AttrT) -> [(ByteString, AttrT)] -> [AttrT]
forall a b. (a -> b) -> [a] -> [b]
map (ByteString, AttrT) -> AttrT
forall a b. (a, b) -> b
snd [(ByteString, AttrT)]
attrs) Converter
conv
      [Int
total, Int
totalFound, Int
time, Int
numWords] <- Int
4 Int -> Get Int -> Get [Int]
forall {a}. Int -> Get a -> Get [a]
`times` Get Int
getNum
      [(Text, Int, Int)]
wrds       <- Int
numWords Int -> Get (Text, Int, Int) -> Get [(Text, Int, Int)]
forall {a}. Int -> Get a -> Get [a]
`times` Converter -> Get (Text, Int, Int)
readWord Converter
conv
      QueryResult -> Get QueryResult
forall a. a -> Get a
forall (m :: * -> *) a. Monad m => a -> m a
return (QueryResult -> Get QueryResult) -> QueryResult -> Get QueryResult
forall a b. (a -> b) -> a -> b
$ [Match]
-> Int -> Int -> [(Text, Int, Int)] -> [ByteString] -> QueryResult
T.QueryResult [Match]
matches Int
total Int
totalFound [(Text, Int, Int)]
wrds (((ByteString, AttrT) -> ByteString)
-> [(ByteString, AttrT)] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map (ByteString, AttrT) -> ByteString
forall a b. (a, b) -> a
fst [(ByteString, AttrT)]
attrs)


readWord :: Converter -> Get (Text, Int, Int)
readWord Converter
conv = do
    ByteString
s <- Get ByteString
getStrStr
    [Int
doc, Int
hits] <- Int
2 Int -> Get Int -> Get [Int]
forall {a}. Int -> Get a -> Get [a]
`times` Get Int
getNum
    (Text, Int, Int) -> Get (Text, Int, Int)
forall a. a -> Get a
forall (m :: * -> *) a. Monad m => a -> m a
return (Converter -> ByteString -> Text
ICU.toUnicode Converter
conv ByteString
s, Int
doc, Int
hits)

readMatch :: Bool -> [AttrT] -> Converter -> Get Match
readMatch Bool
isId64 [AttrT]
attrs Converter
conv = do
    Int64
doc <- if Bool
isId64 then Get Int64
getNum64 else (Get Int
getNum Get Int -> (Int -> Get Int64) -> Get Int64
forall a b. Get a -> (a -> Get b) -> Get b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int64 -> Get Int64
forall a. a -> Get a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int64 -> Get Int64) -> (Int -> Int64) -> Int -> Get Int64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral)
    Int
weight <- Get Int
getNum
    [Attr]
matchAttrs <- (AttrT -> Get Attr) -> [AttrT] -> Get [Attr]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM AttrT -> Get Attr
readAttr [AttrT]
attrs
    Match -> Get Match
forall a. a -> Get a
forall (m :: * -> *) a. Monad m => a -> m a
return (Match -> Get Match) -> Match -> Get Match
forall a b. (a -> b) -> a -> b
$ Int64 -> Int -> [Attr] -> Match
T.Match Int64
doc Int
weight [Attr]
matchAttrs
  where
    readAttr :: AttrT -> Get Attr
readAttr (T.AttrTMulti AttrT
attr)  = (Get Attr -> Get [Attr]
forall {a}. Get a -> Get [a]
readList (AttrT -> Get Attr
readAttr AttrT
attr)) Get [Attr] -> ([Attr] -> Get Attr) -> Get Attr
forall a b. Get a -> (a -> Get b) -> Get b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Attr -> Get Attr
forall a. a -> Get a
forall (m :: * -> *) a. Monad m => a -> m a
return (Attr -> Get Attr) -> ([Attr] -> Attr) -> [Attr] -> Get Attr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Attr] -> Attr
T.AttrMulti
    readAttr AttrT
T.AttrTBigInt    = Get Int64
getNum64 Get Int64 -> (Int64 -> Get Attr) -> Get Attr
forall a b. Get a -> (a -> Get b) -> Get b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Attr -> Get Attr
forall a. a -> Get a
forall (m :: * -> *) a. Monad m => a -> m a
return (Attr -> Get Attr) -> (Int64 -> Attr) -> Int64 -> Get Attr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> Attr
T.AttrBigInt
    readAttr AttrT
T.AttrTString    = Get ByteString
getStrStr  Get ByteString -> (ByteString -> Get Attr) -> Get Attr
forall a b. Get a -> (a -> Get b) -> Get b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Attr -> Get Attr
forall a. a -> Get a
forall (m :: * -> *) a. Monad m => a -> m a
return (Attr -> Get Attr)
-> (ByteString -> Attr) -> ByteString -> Get Attr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Attr
T.AttrString (Text -> Attr) -> (ByteString -> Text) -> ByteString -> Attr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Converter -> ByteString -> Text
ICU.toUnicode Converter
conv
    readAttr AttrT
T.AttrTUInt      = Get Int
getNum Get Int -> (Int -> Get Attr) -> Get Attr
forall a b. Get a -> (a -> Get b) -> Get b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Attr -> Get Attr
forall a. a -> Get a
forall (m :: * -> *) a. Monad m => a -> m a
return (Attr -> Get Attr) -> (Int -> Attr) -> Int -> Get Attr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Attr
T.AttrUInt
    readAttr AttrT
T.AttrTFloat     = Get Float
getFloat Get Float -> (Float -> Get Attr) -> Get Attr
forall a b. Get a -> (a -> Get b) -> Get b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Attr -> Get Attr
forall a. a -> Get a
forall (m :: * -> *) a. Monad m => a -> m a
return (Attr -> Get Attr) -> (Float -> Attr) -> Float -> Get Attr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Float -> Attr
T.AttrFloat
    readAttr AttrT
_                = Get Int
getNum  Get Int -> (Int -> Get Attr) -> Get Attr
forall a b. Get a -> (a -> Get b) -> Get b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Attr -> Get Attr
forall a. a -> Get a
forall (m :: * -> *) a. Monad m => a -> m a
return (Attr -> Get Attr) -> (Int -> Attr) -> Int -> Get Attr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Attr
T.AttrUInt


readAttrPair :: Get (ByteString, AttrT)
readAttrPair = do
    ByteString
s <- Get ByteString
getStr
    Int
t <- Get Int
getNum
    (ByteString, AttrT) -> Get (ByteString, AttrT)
forall a. a -> Get a
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString
s, Int -> AttrT
forall a. Enum a => Int -> a
toEnum Int
t)

readHeader :: ByteString -> (Status, Word16, Word32)
readHeader = Get (Status, Word16, Word32)
-> ByteString -> (Status, Word16, Word32)
forall a. Get a -> ByteString -> a
runGet (Get (Status, Word16, Word32)
 -> ByteString -> (Status, Word16, Word32))
-> Get (Status, Word16, Word32)
-> ByteString
-> (Status, Word16, Word32)
forall a b. (a -> b) -> a -> b
$ do Word16
status  <- Get Word16
getWord16be
                         Word16
version <- Get Word16
getWord16be
                         Word32
length  <- Get Word32
getWord32be
                         (Status, Word16, Word32) -> Get (Status, Word16, Word32)
forall a. a -> Get a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Status
T.toStatus (Int -> Status) -> Int -> Status
forall a b. (a -> b) -> a -> b
$ Word16 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word16
status, Word16
version, Word32
length)