{-# LANGUAGE ScopedTypeVariables, DeriveDataTypeable #-}

module Hoogle.Store.Type(
    Once, once, fromOnce, putOnce, getOnce,
    SPut, runSPut, putByteString, putStorable, putDefer,
    SGet, runSGet, getByteString, getStorable, getDefer, getLazyList
    ) where

import General.Base
import General.System
import Control.Monad.IO.Class
import Control.Monad.Trans.Reader
import qualified Data.IntMap as IntMap
import Data.IORef
import Data.Typeable
import Foreign
import System.IO.Unsafe
import qualified Hoogle.Store.ReadBuffer as R
import qualified Hoogle.Store.WriteBuffer as W


-- | Turn on to see file statistics
stats = False


-- | All once values are equal with respect to keyOnce
--   If you create it with 'once' it will have the same key.
--   If two are loaded from a file they are equal.
data Once a = Once {_onceKey :: Int, valueOnce :: a}
              deriving Typeable

instance NFData a => NFData (Once a) where
    rnf (Once a b) = rnf (a,b)

fromOnce :: Once a -> a
fromOnce = valueOnce

-- | Given how many you would like to allocate, return your base address
onceKeys :: Int -> IO Int
onceKeys = System.IO.Unsafe.unsafePerformIO $ do
    ref <- newIORef 0
    return $ \n -> atomicModifyIORef ref $ \x -> (x+n, x)


---------------------------------------------------------------------
-- PUT

data SPutS = SPutS
    {putBuffer :: W.Buffer
    ,putOnces :: IORef (IntMap.IntMap PutOnce)
    ,putPending :: IORef [SPut ()]
    }

type SPut a = ReaderT SPutS IO a

modifyRef f x = liftIO . (`modifyIORef` x) =<< asks f
readPos = liftIO . W.getPos =<< asks putBuffer


runSPut :: FilePath -> SPut () -> IO ()
runSPut file act = withBinaryFile file WriteMode $ \h -> do
    pending <- newIORef [act]
    once <- newIORef IntMap.empty

    W.withBuffer h $ \buffer -> do
        let flush = do
                xs <- liftIO $ readIORef pending
                liftIO $ writeIORef pending []
                forM_ xs $ \x -> do
                    x
                    flush
        runReaderT flush $ SPutS buffer once pending


putByteString :: BString -> SPut ()
putByteString x = do
    buf <- asks putBuffer
    liftIO $ W.putByteString buf x

putStorable :: Storable a => a -> SPut ()
putStorable x = do
    buf <- asks putBuffer
    liftIO $ W.putStorable buf x


putDefer :: SPut () -> SPut ()
putDefer act = do
    pos <- readPos
    putStorable (0 :: Word32)
    modifyRef putPending $ (:) $ do
        val <- readPos
        buf <- asks putBuffer
        liftIO $ W.patch buf pos val
        act


{-# NOINLINE once #-}
once :: a -> Once a
once x = System.IO.Unsafe.unsafePerformIO $ do
    key <- onceKeys 1
    return $ Once key x


type PutOnce = Either [Word32] Word32

putOnce :: (a -> SPut ()) -> Once a -> SPut ()
putOnce act (Once key x) = do
    ref <- asks putOnces
    mp <- liftIO $ readIORef ref
    case fromMaybe (Left []) $ IntMap.lookup key mp of
        -- written out at this address
        Right val -> putStorable val

        -- [] is has not been added to the defer list
        -- (:) is on defer list but not yet written, these are places that need back patching
        Left poss -> do
            pos <- readPos
            liftIO $ writeIORef ref $ IntMap.insert key (Left $ pos:poss) mp
            putStorable (0 :: Word32)
            when (null poss) $ modifyRef putPending $ (:) $ do
                val <- readPos
                mp <- liftIO $ readIORef ref
                let Left poss = mp IntMap.! key
                buf <- asks putBuffer
                liftIO $ forM_ poss $ \pos -> W.patch buf pos val
                liftIO $ writeIORef ref $ IntMap.insert key (Right val) mp
                act x


---------------------------------------------------------------------
-- GET

-- getPtr is the pointer you have, how much is left valid, 
data SGetS = SGetS {getBuffer :: R.Buffer, onceBase :: Int}

type SGet a = ReaderT SGetS IO a


runSGet :: Typeable a => FilePath -> SGet a -> IO a
runSGet file m = do
    h <- openBinaryFile file ReadMode
    sz <- hFileSize h
    buf <- R.newBuffer h
    one <- onceKeys $ fromIntegral sz
    runReaderT (getDeferFrom 0 m) $ SGetS buf one


getStorable :: Typeable a => Storable a => SGet a
getStorable = do
    buf <- asks getBuffer
    res <- liftIO $ R.getStorable buf
    when stats $ liftIO $ putStrLn $ "Reading storable " ++ show (sizeOf res)
    return res


getByteString :: Word32 -> SGet BString
getByteString len = do
    buf <- asks getBuffer
    when stats $ liftIO $ putStrLn $ "Reading bytestring " ++ show len
    liftIO $ R.getByteString buf $ fromIntegral len


getDefer :: Typeable a => SGet a -> SGet a
getDefer get = do
    pos :: Word32 <- getStorable
    getDeferFrom pos get


getDeferFrom :: forall a . Typeable a => Word32 -> SGet a -> SGet a
getDeferFrom pos get = do
    s <- ask
    liftIO $ unsafeInterleaveIO $ do
        when stats $ putStrLn $ "Read at " ++ show (typeOf (undefined :: a))
        R.setPos (getBuffer s) pos
        runReaderT get s


getOnce :: Typeable a => SGet a -> SGet (Once a)
getOnce get = do
    pos :: Word32 <- getStorable
    x <- getDeferFrom pos get
    one <- asks onceBase
    return $ Once (fromIntegral pos + one) x


getLazyList :: SGet a -> Int -> Int -> SGet [a]
getLazyList get size n = do
    s <- ask
    pos <- liftIO $ R.getPos $ getBuffer s
    liftIO $ forM [0..n-1] $ \i -> unsafeInterleaveIO $ do
        R.setPos (getBuffer s) (pos + fromIntegral (i * size))
        runReaderT get s