{-# language BangPatterns #-}
{-# language LambdaCase #-}
module Main where

import Data.ByteString.Lazy (toStrict)
import qualified Data.ByteString.Lazy as BL
import qualified Data.ByteString as B
import Data.ByteString.Builder (toLazyByteString)
import Data.Vector as V(fromList, empty)
import Data.IORef
import Data.Int
import Data.Foldable
import Data.Maybe
import Control.Concurrent
import Control.Applicative
import Control.Monad
import Data.Monoid
import System.Clock
import Options.Applicative

import qualified Database.PostgreSQL.LibPQ as LibPQ

import Database.PostgreSQL.Protocol.Types
import Database.PostgreSQL.Protocol.DataRows
import Database.PostgreSQL.Driver
--
-- CREATE TABLE _bytes_100_of_1k(b bytea);
-- CREATE TABLE _bytes_400_of_200(b bytea);
-- CREATE TABLE _bytes_10_of_20k(b bytea);
-- CREATE TABLE _bytes_1_of_200(b bytea);
-- CREATE TABLE _bytes_300_of_100(b bytea);

-- INSERT INTO _bytes_100_of_1k(b)
--   (SELECT repeat('a', 1000)::bytea FROM generate_series(1, 100));
--
-- INSERT INTO _bytes_400_of_200(b)
--   (SELECT repeat('a', 200)::bytea FROM generate_series(1, 400));
--
-- INSERT INTO _bytes_10_of_20k(b)
--   (SELECT repeat('a', 20000)::bytea FROM generate_series(1, 10));
--
-- INSERT INTO _bytes_1_of_200(b) VALUES(repeat('a', 200)::bytea);
--
-- INSERT INTO _bytes_300_of_100(b)
--  (SELECT repeat('a', 100)::bytea FROM generate_series(1, 300));

data Action 
    = BenchPW RowsType
    | BenchLibPQ RowsType
    | BenchLoop
    deriving (Show, Eq)

data RowsType
    = Bytes100_1k
    | Bytes400_200
    | Bytes10_20k
    | Bytes1_200
    | Bytes300_100
    deriving (Show, Eq)

cli :: Parser Action
cli = hsubparser $
       cmd "pw" "benchmark postgres-wire" (BenchPW <$> rowTypeParser)
    <> cmd "libpq" "benchmark libpq" (BenchLibPQ <$> rowTypeParser) 
    <> cmd "loop" "benchmark datarows decoding loop" (pure BenchLoop)
  where
    cmd c h p = command c (info (helper <*> p) $ header h)
    rowTypeParser = hsubparser $
           cmd "b100_1k"  "100 rows of 1k bytes"  (pure Bytes100_1k)
        <> cmd "b400_200" "400 rows of 200 bytes" (pure Bytes400_200)
        <> cmd "b10_20k"  "10 rows of 20k bytes"  (pure Bytes10_20k)
        <> cmd "b1_200"   "1 row of 200 bytes"    (pure Bytes1_200)
        <> cmd "b300_100" "300 rows of 100 bytes" (pure Bytes300_100)  

main :: IO ()
main = execParser (info (helper <*> cli) $ header "Postgres-wire benchmark") 
        >>= execAction

execAction :: Action -> IO ()
execAction (BenchPW rows)    = benchPw $ queryStatement rows
execAction (BenchLibPQ rows) = benchLibpq $ queryStatement rows
execAction BenchLoop         = benchLoop

queryStatement :: RowsType -> B.ByteString
queryStatement = \case
    Bytes100_1k  -> "SELECT * from _bytes_100_of_1k"
    Bytes400_200 -> "SELECT * from _bytes_400_of_200"
    Bytes10_20k  -> "SELECT * from _bytes_10_of_20k"
    Bytes1_200   -> "SELECT * fromm _bytes_1_of_200"
    Bytes300_100 -> "SELECT * from _bytes_300_of_100"

benchPw :: B.ByteString -> IO ()
benchPw statement = benchRequests createConnection $ \c -> do
    sendBatchAndSync c [q]
    d <- readNextData c
    waitReadyForQuery c
  where
    q = Query statement [] Binary Binary AlwaysCache
    createConnection = connect defaultSettings >>= 
        either (error . ("Connection error " <>) . show) pure

    defaultSettings = defaultConnectionSettings
        { settingsHost     = "localhost"
        , settingsDatabase = "travis_test"
        , settingsUser     = "postgres"
        , settingsPassword = ""
        }

benchLibpq :: B.ByteString -> IO ()
benchLibpq statement = benchRequests libpqConnection $ \c -> do
    r <- fromJust <$> LibPQ.execPrepared c "" [] LibPQ.Binary
    rows <- LibPQ.ntuples r
    parseRows r (rows - 1)
  where
    libpqConnection = do
        conn <- LibPQ.connectdb "host=localhost user=postgres dbname=travis_test"
        LibPQ.prepare conn "" "SELECT * from _bytes_300_of_100" Nothing
        pure conn
    parseRows r (-1) = pure ()
    parseRows r n = LibPQ.getvalue r n 0 >> parseRows r (n - 1)

benchRequests :: IO c -> (c -> IO a) -> IO ()
benchRequests connectAction queryAction = do
    results <- replicateM 8 newThread
    threadDelay $ durationSeconds * 1000 * 1000
    for_ results $ \(_, _, tid) -> killThread tid
    s <- sum <$> traverse (\(ref, _, _) -> readIORef ref) results
    latency_total <- sum <$> traverse (\(_, ref, _) -> readIORef ref) results

    print $ "Requests per second: " ++ show (s `div` durationSeconds)
    print $ "Average latency, ms: " ++ displayLatency latency_total s
  where
    durationSeconds = 10
    newThread  = do
        ref_count   <- newIORef 0 :: IO (IORef Int)
        ref_latency <- newIORef 0 :: IO (IORef Int64)
        c <- connectAction
        tid <- forkIO $ forever $ do
            t1 <- getTime Monotonic
            r <- queryAction c
            r `seq` pure ()
            t2 <- getTime Monotonic
            modifyIORef' ref_latency (+ (getDifference t2 t1))
            modifyIORef' ref_count (+1)
        pure (ref_count, ref_latency, tid)

    getDifference (TimeSpec end_s end_ns) (TimeSpec start_s start_ns) = 
        (end_s - start_s) * 1000000000 + end_ns - start_ns

    displayLatency latency reqs = 
        let a = latency `div` fromIntegral reqs
            (ms, ns) = a `divMod` 1000000
        in show ms <> "." <> show ns

benchLoop :: IO ()
benchLoop = do
    counter <- newIORef 0  :: IO (IORef Word)
    content <- newIORef "" :: IO (IORef BL.ByteString)
    -- File contains a PostgreSQL binary response on the query:
    --   "SELECT typname, typnamespace, typowner, typlen, typbyval,
    --          typcategory, typispreferred, typisdefined, typdelim,
    --          typrelid, typelem, typarray from pg_type"
    !bs <- B.readFile "bench/pg_type_rows.out"
    writeIORef content . BL.cycle $ BL.fromStrict bs

    let handler dm = case dm of
            DataMessage _ -> modifyIORef' counter (+1)
            _             -> pure ()
        newChunk preBs = do
            b <- readIORef content
            let (nb, rest) = BL.splitAt 4096 b
            writeIORef content rest
            let res = preBs <> ( BL.toStrict nb)
            res `seq` pure res
    tid <- forkIO . forever $ loopExtractDataRows newChunk handler
    threadDelay $ durationSeconds * 1000 * 1000
    killThread tid
    s <- readIORef counter
    print $ "Data messages parsed per second: " 
            ++ show (s `div` fromIntegral durationSeconds)
  where
    durationSeconds = 10
