{-# LANGUAGE ForeignFunctionInterface #-}

import System.Random.Mersenne.Pure64
import Control.Concurrent
import Control.Monad
import Data.Binary.Put
import qualified Data.ByteString.Lazy as BS
import qualified Data.ByteString      as S
import Data.Word
import Foreign.Marshal.Alloc
import Foreign.Ptr
import Foreign.Storable
import System.Environment
import GHC.Conc
import Foreign.C.Types
import Debug.Trace
import System.IO


--
-- Parallel unfold and reduce.
--
-- -lcrypto
--

main :: IO ()
main = do
  args <- getArgs
  case args of
    [file, size]       -> checkReasonableAndGo (Just (show numCapabilities)) file size
    [cpus, file, size] -> checkReasonableAndGo (Just cpus) file size
    _                  -> dumpUsageAndQuit

dumpUsageAndQuit :: IO ()
dumpUsageAndQuit = mapM_ putStrLn
    ["USAGE: genblock [cpus] filename size"
    ,"where 'size' is given in gigabytes"
    ,"and cpus defaults to the value of +RTS -N"]

checkReasonableAndGo :: Maybe String -> String -> String -> IO ()
checkReasonableAndGo mcpus file size =
  case (maybe [(numCapabilities, "")] reads mcpus, reads size) of
    ([(cpus, "")], [(size', "")]) -> go cpus file size'
    _                             -> dumpUsageAndQuit

go :: Int -> String -> Int -> IO ()
go cpus filename size_in_gb = do
  chan <- newChan
  forM_ [1..cpus] $ \n -> forkIO {-n-} $ do
    baseMT <- pureMT `fmap` getOpenSSLRand
    writeList2Chan chan $
            BS.toChunks $ runPut $ mersennePut baseMT

  ls <- getChanContents chan    -- TODO need a LBS chan
  BS.writeFile filename . BS.take size . BS.fromChunks $ ls

 where
  size = fromIntegral size_in_gb * 1024 * 1024 * 1024


-- unroll the loop
mersennePut :: PureMT -> Put
mersennePut mt = do

  let (k, mt1) = randomWord64 mt
  putWord64host k
  let (k, mt2) = randomWord64 mt1
  putWord64host k
  let (k, mt3) = randomWord64 mt2
  putWord64host k
  let (k, mt4) = randomWord64 mt3
  putWord64host k
  let (k, mt5) = randomWord64 mt4
  putWord64host k
  let (k, mt6) = randomWord64 mt5
  putWord64host k
  let (k, mt7) = randomWord64 mt6
  putWord64host k
  let (k, mt8) = randomWord64 mt7
  putWord64host k
  let (k, mt9) = randomWord64 mt8
  putWord64host k
  let (k, mt10) = randomWord64 mt9
  putWord64host k
  let (k, mt11) = randomWord64 mt10
  putWord64host k
  let (k, mt12) = randomWord64 mt11
  putWord64host k
  let (k, mt13) = randomWord64 mt12
  putWord64host k
  let (k, mt14) = randomWord64 mt13
  putWord64host k
  let (k, mt15) = randomWord64 mt14
  putWord64host k
  let (k, mt16) = randomWord64 mt15
  putWord64host k
  mersennePut mt16

------------------------------------------------------------------------

--
-- Get a decent random seed
--
getOpenSSLRand :: IO Word64
getOpenSSLRand = allocaBytes n $ \ ptr -> do
  fill_RAND_bytes ptr (fromIntegral n)
  peek $ castPtr ptr
 where
  n = sizeOf (undefined :: Word64)

foreign import ccall unsafe "openssl/rand.h RAND_bytes"
  fill_RAND_bytes :: Ptr Word8 -> CInt -> IO ()