--
-- Module : HSClient
-- Author : Wu Xingbo
-- Copyright (c) 2011 Wu Xingbo (wuxb45@gmail.com)
-- New BSD License
--
-- |
-- Haskell implementation of a HandlerSocket client (API).
-- HandlerSocket is a MySQL plugin with better performance then SQL.
-- HandlerSocket site: https://github.com/ahiguti/HandlerSocket-Plugin-for-MySQL
--


module Database.HandlerSocketClient (
  -- * Basic Type
  HandSock (..),
  -- * Connection
  connectTo,
  closeHS,
  -- * Do Generic Request
  runReqNoRespData,
  runReqGeneric,
  -- * Do Special Request
  runReqOpen,
  runReqFind,
  runReqInsert,
  runReqUpdate,
  runReqDelete
  ) where
import Data.Word
import Data.Char
import Data.List
import Network.Socket hiding (send,sendTo,recv,recvFrom)
import Network.Socket.ByteString
import System.IO
import qualified Data.ByteString as B

-- | basic data structure for one Socket and one Handle (or Index).
data HandSock = HandSock {getSocket :: Socket, getHost :: String, getPort :: String}

recvOne :: Socket -> Int -> IO B.ByteString
recvOne sock maxlen = do
  s <- recvOneR sock maxlen
  if B.elem 10 s
    then return s
    else return B.empty
  
recvOneR :: Socket -> Int -> IO B.ByteString
recvOneR sock maxlen = do
  s1 <- recv sock maxlen
  if B.null s1
    then return s1
    else
      if B.elem 10 s1
        then return s1
        else fmap (B.append s1) $ recvOneR sock maxlen

escapeW :: [Word8] -> [Word8]
escapeW [] = []
escapeW (x:xs) = if x < 16 then 1:(0x40+x):(escapeW xs) else x:(escapeW xs)

escape :: String -> [Word8]
escape str = escapeW $ map (fromIntegral . ord) str

unescapeW :: [Word8] -> [Word8]
unescapeW [] = []
unescapeW (x:xs) = 
  if x == 1 
  then case xs of
    (y:ys) -> (y-0x40):(unescapeW ys)
    _ -> [] -- error
  else x:(unescapeW xs)

unescape :: [Word8] -> String
unescape xs = map (chr . fromIntegral) $ unescapeW xs

decodeResp :: B.ByteString -> [String]
decodeResp = map unescape . mysplit . B.unpack
  where
    isspace' = (\x -> (x == 9) || (x == 10)) -- \t \n
    mysplit s = case dropWhile isspace' s of -- drop leading space'
      [] -> []
      (10:_) -> []
      s1 -> w : (mysplit s2) -- first : rest.
            where
              (w,s2) = break isspace' s1

makeReq :: [String] -> B.ByteString
makeReq = B.pack . (++ [10]) . intercalate [9] . map escape

makeReqOpen :: String -> String -> String -> String -> B.ByteString
makeReqOpen db tb ind cs = makeReq ["P","0",db,tb,ind,cs]

makeReqFind :: (Integral a) => String -> String -> a -> B.ByteString
makeReqFind op v lim = makeReq ["0",op,"1",v,show lim,"0"]

makeReqInsert :: [String] -> B.ByteString
makeReqInsert vs = makeReq $ ["0","+",show (length vs)] ++ vs

makeReqUpdate :: String -> [String] -> B.ByteString
makeReqUpdate key newvs = makeReq $ ["0","=","1",key,"1","0","U"] ++ newvs

makeReqDelete :: String -> B.ByteString
makeReqDelete key = makeReq ["0","=","1",key,"1","0","D"]

-- | Connect to specific address & port.
-- Returns the 'HandSock' for later request or close.
connectTo :: String -> String -> IO HandSock
connectTo host port = do
  addrinfos <- getAddrInfo Nothing (Just host) (Just port)
  let addr = head addrinfos
  let addr' = addr {addrFamily = AF_INET, addrSocketType = Stream, addrProtocol = 6}
  sock <- socket AF_INET Stream 6
  connect sock (addrAddress addr')
  return $ HandSock sock host port

runReqBool :: HandSock -> B.ByteString -> IO Bool
runReqBool hs req = do
  sendAll (getSocket hs) req
  resp <- recvOne (getSocket hs) 1024
  if B.null resp
    then return False
    else return $ 48 == (B.head resp)

-- | Do Generic Request with given Strings. Return True or False, discarding all other messages.
runReqNoRespData :: HandSock -> [String] -> IO Bool
runReqNoRespData hs vs = runReqBool hs $ makeReq vs

-- | Do Generic Request, Returns all values received, \"0\" in the head means SUCCESS.
runReqGeneric :: HandSock -> [String] -> IO [String]
runReqGeneric hs vs = do
  let req = makeReq vs
  sendAll (getSocket hs) req
  fmap decodeResp $ recvOne (getSocket hs) 4096

-- | Do \'open_index\' request.
runReqOpen :: HandSock -> String -> String -> String -> String -> IO Bool
runReqOpen hs db tb ind cs = runReqBool hs $ makeReqOpen db tb ind cs

-- | Do \'find\' request. Returns at most lim results.
runReqFind :: (Integral a) => HandSock -> String -> String -> a -> IO [[String]]
runReqFind hs op v lim = do
  let req = makeReqFind op v lim
  sendAll (getSocket hs) req
  resp <- fmap decodeResp $ recvOne (getSocket hs) 4096
  case resp of
    (o:i:xs) -> return $ mygroup (read i :: Int) xs
    _ -> return []
  where
    mygroup _ [] = []
    mygroup i xs = l:(mygroup i r) 
      where (l,r) = splitAt i xs

-- | Do \'insert\' request, insert one record at a time.
runReqInsert :: HandSock -> [String] -> IO Bool
runReqInsert hs vs = runReqBool hs $ makeReqInsert vs

-- | Do \'update\' request, give one value to be compared with the KEY, then provide the new data.
-- Only ONE record will be updated.
runReqUpdate :: HandSock -> String -> [String] -> IO Bool
runReqUpdate hs v0 vs = runReqBool hs $ makeReqUpdate v0 vs

-- | Do \'delete\' request, give one value to be compared with the KEY. 
-- Only ONE record will be deleted.
runReqDelete :: HandSock -> String -> IO Bool
runReqDelete hs key = runReqBool hs $ makeReqDelete key

-- | Close the connection.
closeHS :: HandSock -> IO ()
closeHS hs = sClose (getSocket hs)

reConnect :: HandSock -> IO HandSock
reConnect hs = do
  closeHS hs
  connectTo (getHost hs) (getPort hs)