{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE GADTs #-}

module Crypto.RandomMonad (RndT, RndST, RndIO, Rnd, RndState, getRandomM, getRandom2M, runRndT, newRandomElementST, getRandomElement, randomElementsLength, replaceSeedM, addSeedM, getRandomByteStringM, RandomElementsListST, RndStateList(..), RndStateParallelization(..), BitStringToRandomExceptions(..), seedFromBytestrings, seedFromBytestringsThreaded, seedFromBytestringsM) where

import Control.DeepSeq (deepseq, force)
import Control.Exception (Exception, throw)
import Control.Monad.Identity (Identity)
import Control.Monad.Primitive (PrimMonad, PrimState, primitive)
import Control.Monad.ST (ST)
import Control.Monad.Trans.Class (MonadTrans, lift)
import Control.Monad.Trans.State.Lazy (StateT, put, state, runStateT)
import Control.Parallel (par, pseq)
import Control.Parallel.Strategies (parMap, rpar, using, rseq)
import Data.Bits (shiftL, (.|.), xor)
import Data.STRef (STRef, newSTRef, readSTRef, writeSTRef)
import Data.Typeable (Typeable)

import qualified Data.BitString as BS
import qualified Data.ByteString as ByS
import qualified Data.ByteString.Lazy as ByLS
import qualified Data.Vector.Unboxed as V
import qualified Data.Vector.Unboxed.Mutable as VM

data BitStringToRandomExceptions = OutOfElementsException deriving (Show, Typeable)
data RandomElementsListST a s = VM.Unbox a => RandomElementsListST (STRef s (V.Vector a))
instance Exception BitStringToRandomExceptions

bitsNeeded :: Integer -> Integer
bitsNeeded x = (+) 1 $ floor $ logBase 2 (fromIntegral x)

convertBitStringToInteger = BS.foldl' convert' 0
  where
  convert' :: Integer -> Bool -> Integer
  convert' prev cur = (shiftL prev 1) .|. (case cur of True -> 1 ; False -> 0)

multipleBitstringsSplitAt i (RndStateList RndStateParallel [x]) = let (takers, droppers) = BS.splitAt i x in ([takers], RndStateList RndStateParallel [droppers])
multipleBitstringsSplitAt i (RndStateList RndStateParallel x) = join' (split' x) [] []
 where
 split' = parMap rpar (\bs -> let (take,drop) = BS.splitAt i bs in (take `using` rseq, drop))
 join' [] takers droppers = (takers, RndStateList RndStateParallel droppers)
 join' (x:xs) takers droppers = let (newTake, newDrop) = x in join' xs (newTake:takers) (newDrop:droppers)

multipleBitstringsSplitAt i (RndStateList p [x]) = let (takers, droppers) = BS.splitAt i x in ([takers], RndStateList p [droppers])
multipleBitstringsSplitAt i (RndStateList p x) = join' (split' x) [] []
 where
 split' = map $ BS.splitAt i
 join' [] takers droppers = (takers, RndStateList p droppers)
 join' (x:xs) takers droppers = let (newTake, newDrop) = x in join' xs (newTake:takers) (newDrop:droppers)

multipleBitstringsAssertLength _ [] = False
multipleBitstringsAssertLength len x = len' x
  where
  len' [] = True
  len' (x:xs) = if (BS.length x) == len
                   then len' xs
                   else False


getRandom :: Integer -> RndStateList -> (Integer, RndStateList)
getRandom 0 x = (0, x)
getRandom max string = if has_error
                          then error "There was an error acquiring random data"
                          else if random <= max
                                  then (random, unused)
                                  else getRandom max unused
  where
    bitsNeeded' = bitsNeeded max
    has_error = not $ multipleBitstringsAssertLength (fromInteger bitsNeeded') used
    random = foldl (\i cur -> xor i $ convertBitStringToInteger cur) 0 used
    (used, unused) = multipleBitstringsSplitAt (fromIntegral bitsNeeded') string

getRandom2 :: Integer -> Integer -> RndStateList -> (Integer, RndStateList)
getRandom2 a b string = getRandom2' (getRandom (max' - min') string)
  where
  min' = min a b
  max' = max a b
  getRandom2' (random, unused) = (random + min', unused)

getRandomByteString :: Integer -> RndStateList -> (ByS.ByteString, RndStateList)
getRandomByteString 0 x = (ByS.pack [], x)
getRandomByteString len x = let (byte, newState) = getRandom 255 x ; (allBytes, lastState) = getRandomByteString (len - 1) newState in (ByS.cons (fromIntegral byte) allBytes, lastState)

newRandomElementST :: VM.Unbox a => [a] -> ST s (RandomElementsListST a s)
newRandomElementST acc = (newSTRef $ V.fromList acc) >>= \ref -> return $ RandomElementsListST ref

getRandomElement :: (RandomElementsListST a s) -> RndST s a
getRandomElement (RandomElementsListST ref) = do
  mut <- lift $ readSTRef ref >>= V.unsafeThaw
  let n = toInteger $ VM.length mut
  j <- if n > 0 then getRandomM $ n - 1
                else throw OutOfElementsException
  let j' = fromInteger j
  lift $ do
    aa <- VM.read mut 0
    ab <- VM.read mut j'
    VM.write mut j' aa
    vec <- V.unsafeFreeze mut
    writeSTRef ref $ V.unsafeTail vec
    return ab

randomElementsLength :: RandomElementsListST a s -> RndST s Int
randomElementsLength (RandomElementsListST ref) = do
  vec <- lift $ readSTRef ref
  return $ V.length vec

type RndStatePrimitive = [BS.BitString]
data RndStateParallelization = RndStateSeq | RndStatePrecalculate | RndStateParallel
data RndStateList = RndStateList RndStateParallelization RndStatePrimitive

type RndState = RndStateList
newtype RndT m a = RndT
  { unRndT :: StateT RndState m a }
  deriving (Functor, Applicative, Monad, MonadTrans)

instance PrimMonad m => PrimMonad (RndT m) where
  type PrimState (RndT m) = PrimState m
  primitive = lift . primitive
  {-# INLINE primitive #-}

type RndST s a = RndT (ST s) a
type RndIO a = RndT IO a
type Rnd a = RndT Identity a

replaceSeedM :: Monad m => RndStatePrimitive -> RndT m ()
replaceSeedM s = RndT $ state $ replaceSeedM s
  where
  replaceSeedM newState (RndStateList p _) = ((),RndStateList p newState)

addSeedM :: Monad m => RndStatePrimitive -> RndT m ()
addSeedM s = RndT $ state $ addSeedM s
  where
  addSeedM x (RndStateList p y) = ((),RndStateList p (x ++ y))


getRandomM :: Monad m => Integer -> RndT m Integer
getRandomM x = RndT $ state $ getRandom x

getRandom2M :: Monad m => Integer -> Integer -> RndT m Integer
getRandom2M x y = RndT $ state $ getRandom2 x y

getRandomByteStringM :: Monad m => Integer -> RndT m ByS.ByteString
getRandomByteStringM x = RndT $ state $ getRandomByteString x

runRndT :: RndState -> RndT m a -> m (a, RndState)
runRndT rnd m = runStateT (unRndT m) rnd

withOneElementPrecalculated (x:xs@(y:_)) = par (force y) x :
                                           withOneElementPrecalculated xs

seedFromBytestrings :: [ByS.ByteString] -> RndStatePrimitive
seedFromBytestrings list = [BS.bitStringLazy $ ByLS.fromChunks list]

seedFromBytestringsThreaded :: [ByS.ByteString] -> RndStatePrimitive
seedFromBytestringsThreaded list = [BS.bitStringLazy $ ByLS.fromChunks $ withOneElementPrecalculated list]

seedFromBytestringsM :: Monad m => [ByS.ByteString] -> RndT m RndStatePrimitive
seedFromBytestringsM list = RndT $ state $ getSeed list
  where
  getSeed list st@(RndStateList RndStatePrecalculate _) = (seedFromBytestringsThreaded list, st)
  getSeed list st@(RndStateList _ _) = (seedFromBytestrings list, st)