{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE OverloadedStrings      #-}
{-# LANGUAGE RankNTypes             #-}
module Boots.Random(
    RD(..)
  , HasRandom(..)
  , MonadRandom(..)
  , newRD
  , RDType(..)
  , makeRD
  , makeRD0
  , forkRD
  , hex32
  , hex64
  , nextWord64
  , splitSMGen
  ) where

import           Boots.App.Internal
import           Boots.Prelude
import           Control.Concurrent.MVar
import           Control.Monad.Factory
import           Data.IORef
import           Data.Text               (toLower, unpack)
import           Data.Tuple
import           Foreign
import           Numeric                 (showHex)
import           Salak
import           System.Random.SplitMix

-- | Random value generator.
newtype RD = RD { unRD :: forall a. (SMGen -> (a, SMGen)) -> IO a }

-- | Seed container type.
data RDType = RDIORef | RDMVar

instance FromProp m RDType where
  fromProp = readEnum (go . toLower)
    where
      {-# INLINE go #-}
      go "ioref" = Right RDIORef
      go "mvar"  = Right RDMVar
      go v       = Left $ "unknown <" <> unpack v <> ">"

-- | Environment values with `RD`.
class HasRandom env where
  askRandom :: Lens' env RD

instance HasRandom RD where
  askRandom = id
  {-# INLINE askRandom #-}

-- | Create a new random value generator.
{-# INLINE newRD #-}
newRD :: RDType -> IO RD
newRD RDIORef = initSMGen >>= newIORef >>= \ref -> return (RD $ \f -> atomicModifyIORef' ref (swap.f))
newRD _       = initSMGen >>= makeRD

-- | Create random value generator with a seed.
{-# INLINE makeRD #-}
makeRD :: SMGen -> IO RD
makeRD seed = newMVar seed >>= \ref -> return (RD $ \f -> modifyMVar ref (return . swap . f))

-- | Create a thread unsafe `RD`, which should be used only in single thread.
-- It is faster than `RD` generated by `newRD` and `makeRD`.
{-# INLINE makeRD0 #-}
makeRD0 :: SMGen -> (RD -> IO a) -> IO a
makeRD0 smg f = do
  let (seed, gamma) = unseedSMGen smg
  allocaArray 2 $ \ps -> do
    pokeArray ps [seed,gamma]
    f $ RD $ \func -> do
      [s0,g0] <- peekArray 2 ps
      let (a, smg2) = func $ seedSMGen s0 g0
          (s1,g1)   = unseedSMGen smg2
      pokeArray ps [s1,g1]
      return a

-- | Fork a new `RD` from old `RD`.
{-# INLINE forkRD #-}
forkRD :: RD -> IO RD
forkRD (RD f) = f splitSMGen >>= makeRD

class Monad m => MonadRandom env m | m -> env where
  nextW64   :: m Word64

-- | Convert `Word64` into 64 bit hex.
hex64 :: IsString a => Word64 -> a
hex64 i = fromString $ let x = showHex i "" in replicate (16 - length x) '0' ++ x
{-# INLINE hex64 #-}

-- | Convert `Word64` into 32 bit hex.
hex32 :: IsString a => Word64 -> a
hex32 i = fromString $ let x = showHex i "" in drop 8 $ replicate (16 - length x) '0' ++ x
{-# INLINE hex32 #-}

instance (HasRandom env, MonadMask n, MonadIO n) => MonadRandom env (Factory n env) where
  nextW64 = do
    vr <- asksEnv (view askRandom)
    liftIO $ unRD vr nextWord64
  {-# INLINE nextW64 #-}

instance (HasRandom env, MonadIO n) => MonadRandom env (AppT env n) where
  nextW64 = do
    vr <- asks (view askRandom)
    liftIO $ unRD vr nextWord64
  {-# INLINE nextW64 #-}