{-# LANGUAGE Rank2Types, RankNTypes #-}
module Halfs.Blocks (getDiskAddrOfBlockRead, getDiskAddrOfBlockRaw,
                     getDiskAddrOfBlockWrite, getBlockFromCacheOrDeviceWrite,
                     numBlocksForSize, markBlockDirtyWrite)
    where

import Halfs.Inode(Inode(..), InodeMetadata(..), getPointerAt, updatePointers,
                   goodInodeMagic)
import Data.Integral(INInt, intToINInt, inIntToInt)
import Halfs.Utils(BlockNumber, DiskAddress,
                   unimplemented, bytesPerBlock, blockPointersPerInode,
                   blockPointersPerIndirectBlock)
import Halfs.BufferBlock (BufferBlock(..), getDiskAddr,
                        bufferBlockCursorIntoPointers,
                        BufferBlockCursor,
                        writeToBufferBlock,
                        readFromBufferBlock,
                        zeroBufferBlock,
--UNUSED:               deadbeefBufferBlock,
                        Alloc(..))
import Halfs.FSRW (unsafeLiftIORW)
import Halfs.BinaryMonad(FSRW)
import {-# SOURCE #-} Halfs.BasicIO (getInodeWrite, putPointersBin)
import Halfs.TheBlockMap(findAddress, allocateBlock)
import Halfs.FSRoot (FSRoot(..))
import Halfs.FSState (FSWrite, FSRead, unsafeLiftIOWrite,
                      unsafeReadGet, updateInodeCacheWrite,
                      modifyFSWrite, unsafeWriteGet)
import System.RawDevice(RawDevice)
import Halfs.BufferBlockCache (markBlockDirty,
                               getBlockFromCacheOrDevice,
                               getNewBlockFromCacheOrDevice,
                               getBlockFromDevice)

-- base
import Control.Exception(assert)
import Control.Monad(when)
import Data.Array(listArray)
import Control.Monad.Error(MonadError)

-- |Traverses the tree of DiskAddresses implied by this block number
-- and inode.  Gets the block at that disk address, and positions the
-- pointer in that block according to the last address in the tree,
-- then reads the disk address there.  This function used only for
-- reading.
getDiskAddrFun :: (MonadError e m,FSRW m)
               => Inode
               -> BlockNumber
               -> (forall m1 . forall e1 . (MonadError e1 m1, FSRW m1) => DiskAddress -> (forall s . BufferBlock s -> m1 DiskAddress) -> m1 DiskAddress)
               -- ^function for getting the block since might be raw
               -- device or fsroot.
               -> m DiskAddress -- ^a block of indirect pointers
getDiskAddrFun inode blockNum findBlock = do
  -- special case for the first block; it's in the inode itself.
  let inLevel = level $ metaData (assert (goodInodeMagic inode) inode)
  let addrs = findAddress inode blockNum
  let firstSeek = head (assert ((intToINInt $ length addrs) == inLevel) addrs)
  let firstAddr = getPointerAt inode firstSeek
  findBlock firstAddr (\ nextBlock ->
    looper (tail addrs) nextBlock (\ b c ->
       unsafeLiftIORW (readFromBufferBlock b c)))

  where -- guts of the block recursion; handles all but first case

        looper :: (MonadError e m,FSRW m)
               => [INInt] -- Dereference locations
               -> BufferBlock t -- Block to start at
               -> (forall s . forall m1 . forall e1 . (MonadError e1 m1,FSRW m1) => BufferBlock s -> BufferBlockCursor s -> m1 DiskAddress)
               -> m DiskAddress
        looper [] _b _cont  = error "never happens" -- cont b -- this shouldn't happen.
        looper (h:t) b cont = do
          if length t > 0
           then do addr <- unsafeLiftIORW (readFromBufferBlock b (bufferBlockCursorIntoPointers h))
                   findBlock addr (\ nextBlock ->
                     looper t nextBlock cont)
           else cont b (bufferBlockCursorIntoPointers h)

-- |Get this inode from the inode cache.  It SHOULD be in there, but
-- just to be safe, we will look on disk if it's not.
internalGetInodeFromCache :: INInt -> FSWrite Inode
internalGetInodeFromCache inodeNum = getInodeWrite inodeNum Nothing

-- |Returns Nothing iff address is zero but inode number and block
-- number aren't zero.  FIX: maybe also check for out of range block
-- number?
getDiskAddrOfBlockRead :: Inode
                       -> BlockNumber
                       -> FSRead DiskAddress
getDiskAddrOfBlockRead inode blockNumber = do
  FSRoot{device=dev, bbCache=cache} <- unsafeReadGet
  addr <- case (level $ metaData inode) of
            0 -> unimplemented "Zero optimization (getDiskAddrOfBlockRead)"
            1 -> return $ getPointerAt inode blockNumber
            _ -> do getDiskAddrFun inode blockNumber
                             (getBlockFromCacheOrDevice cache dev)
  case checkedAddr addr inode blockNumber of
    Nothing -> error $ "UNINITIALIZED READ. Address is zero but inode number and block number aren't zero: "
                          ++ (show blockNumber) ++ (show inode)
    Just a -> return a

-- |calls to getBlockFromDevice or getPointerAt, depending on level.
getDiskAddrOfBlockRaw :: (MonadError e m,FSRW m)
                      => Inode
                      -> BlockNumber
                      -> RawDevice
                      -> m (Maybe DiskAddress)
getDiskAddrOfBlockRaw inode blockNumber rawDevice = do
  addr <- case (level $ metaData inode) of
            0 -> unimplemented "Zero optimization (getDiskAddrOfBlockRaw)"
            1 -> return $ getPointerAt inode blockNumber
                 -- FIX: make sure that the block allocated below gets freed.
            _ -> do getDiskAddrFun inode blockNumber
                            (getBlockFromDevice rawDevice)
  return $ checkedAddr addr inode blockNumber

-- |Mark this block dirty.  Doesn't modify the fsroot (except the bbcache).
markBlockDirtyWrite :: BufferBlock s -> FSWrite ()
markBlockDirtyWrite block = do
  addr <- unsafeLiftIOWrite $ getDiskAddr block
  FSRoot{device=dev, bbCache=cache} <- unsafeWriteGet -- safe
  markBlockDirty cache dev addr

{-
-- |Mark this block dirty.  Doesn't modify the fsroot (except the bbcache).
markBlockDirtyRead :: DiskAddress -> FSRead ()
markBlockDirtyRead addr = do
  addr <- unsafeLiftIOWrite $ getDiskAddr block
  FSRoot{device=dev, bbCache=cache} <- unsafeReadGet -- safe
  markBlockDirty cache dev addr
-}

-- |Might modify inode because it might grow it.  Returns Alters inode
-- cache.  FIX: maybe also check for out of range block number?
getDiskAddrOfBlockWrite :: Inode
                        -> BlockNumber
                        -> FSWrite DiskAddress
getDiskAddrOfBlockWrite inodeIn blockNumber = do
  inode <- ensureCapacity inodeIn blockNumber
  updateInodeCacheWrite inode -- just to make sure its in the cache
  addr <- case level $ metaData inode of
            0 -> unimplemented "Zero optimization (getDiskAddrWrite)"
            1 -> do -- unsafeLiftIOWrite $ print "getDiskAddrOfBlockWrite<0>"
                    getAndPositionBlockWrite inode blockNumber (\ block _ -> unsafeLiftIOWrite $ getDiskAddr block)
            _ ->
             do -- unsafeLiftIOWrite $ print "getDiskAddrOfBlockWrite<*>"
                getAndPositionBlockWrite inode blockNumber (\ bb (Just pos) -> do
                    addr <- unsafeLiftIORW $ readFromBufferBlock bb pos
                    return addr)

  -- Have to copy that block if its negative, but it won't be
  -- indirect, so no negations:

  case checkedAddr (assert (addr >= 0) addr) inode blockNumber of
    Nothing -> error $ "internal: ensureCapacity failed to ensure: "
                          ++ (show blockNumber) ++ (show inode)
    Just a -> return a

checkedAddr :: DiskAddress
            -> Inode
            -> BlockNumber
            -> Maybe DiskAddress
checkedAddr addr inode blockNumber=
      if addr == 0
             && inode_num (metaData inode) == 0
             && blockNumber == 0
         then Just addr
         else if addr == 0
              then Nothing
              else Just addr

-- |This is for reading a block.
getBlockFromCacheOrDeviceWrite :: DiskAddress
                               -> (forall s . BufferBlock s -> FSWrite a)
                               -> FSWrite a
getBlockFromCacheOrDeviceWrite da cont = do
  FSRoot{device=dev, bbCache=cache} <- unsafeWriteGet
  getBlockFromCacheOrDevice cache dev da cont

-- |Just like allocateBlock, but happens in the FSWrite monad.
allocateBlockWrite :: (forall s. BufferBlock s -> FSWrite a) -> FSWrite a
allocateBlockWrite cont = do
   mDa <- modifyFSWrite (\fsr@FSRoot{blockMap=bm} -> do
       case allocateBlock bm of
        Nothing          -> return (fsr, Nothing)
        Just (da, bmNew) -> return (fsr{blockMap=bmNew}, Just da))
   case mDa of
     Just da -> do FSRoot{device=dev, bbCache=cache} <- unsafeWriteGet
                   getNewBlockFromCacheOrDevice cache dev da cont
     Nothing -> error "Out of disk space."

-- |Like getDiskAddrFun, but for writing.  If the
-- inode level is 1, then this will return the actual block.
-- Otherwise, it'll return an indirect block, from which you must read
-- or write the address you're interested in.
getAndPositionBlockWrite :: Inode
                         -> BlockNumber
                         -> (forall s . BufferBlock s -> (Maybe (BufferBlockCursor s)) -> FSWrite a) -- ^a block of indirect pointers
                         -> FSWrite a
getAndPositionBlockWrite inode blockNum cont = do
  -- special case for the first block; it's in the inode itself.
  let inLevel = level $ metaData inode
  let inodeNum = inode_num $ metaData inode
  let addrs = findAddress inode blockNum
  let firstSeek = head (assert ((intToINInt $ length addrs) == inLevel) addrs)
  let firstAddr = getPointerAt inode firstSeek
  updateInodeCacheWrite inode -- just to make sure its in the cache, see below
  getBlockFromCacheOrDeviceWrite firstAddr (\ mightBeNextBlock -> do
      -- update the pointer, as it may now point to a new block.
    newAddr <- unsafeLiftIOWrite $ getDiskAddr mightBeNextBlock
    -- inode may have changed while copying; get it out of the
    -- cache. Must be there because we made sure of that above.
    inode1 <- internalGetInodeFromCache inodeNum
    let i = updatePointers inode1 [(inIntToInt firstSeek, newAddr)]
    updateInodeCacheWrite i
    looper (tail addrs) inodeNum mightBeNextBlock cont)
  where -- guts of the block recursion; handles all but first case
        looper :: [INInt] -- Dereference locations
               -> INInt -- InodeNum.
               -> BufferBlock s -- Block to start at
               -> (forall s1 . BufferBlock s1 -> (Maybe (BufferBlockCursor s1)) -> FSWrite a)
               -> FSWrite a
        looper [] _ b ncont = ncont b Nothing
        looper (h:t) inodeNum b ncont = do
          -- seek to the position, but don't actually read:
          let cursor = bufferBlockCursorIntoPointers h
          if length t > 0 -- b is an indirect block
           then do addr <- unsafeLiftIORW $ readFromBufferBlock b cursor
                   getBlockFromCacheOrDeviceWrite addr
                    (\  mightBeNextBlock -> do
                        newAddr <- unsafeLiftIOWrite $ getDiskAddr mightBeNextBlock
                        -- the block might be at a diff. address now. if so, write it back.
                        when (newAddr /= addr) $ do
                                          unsafeLiftIORW $ writeToBufferBlock b cursor  [Alloc newAddr]
                                          markBlockDirtyWrite b
                        looper t inodeNum mightBeNextBlock ncont)
           else ncont b (Just cursor)

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

-- |Conditionally grows the inode.  Ensures that the inode now
-- containes THIS BLOCK NUMBER, not this size (watch for off-by-one
-- errors).  NOTE: Be sure to bump the size when you're done!
ensureCapacity :: Inode
               -> BlockNumber  -- ^new capacity
               -> FSWrite Inode
ensureCapacity inode newBlockNum
    = ensureCapacity' inode (numBlocksForSize inode) newBlockNum

ensureCapacity' :: Inode
               -> BlockNumber  -- ^current largest block number
               -> BlockNumber  -- ^new largest block number
               -> FSWrite Inode
ensureCapacity' inodeIn currBlocksInFile newBlockNum
    | (abs newBlockNum) < currBlocksInFile =
           updateInodeCacheWrite inodeIn
--                  >> putStrLnWriteRead ("no need to grow: " ++(show (inode_num (metaData inodeIn))) ++ ":" ++ (show newBlockNum) ++ ":" ++ (show currBlocksInFile) ++ ":" ++ (show (num_bytes $ metaData inodeIn)))
                  >> return inodeIn
    | otherwise = do
        -- may change inode's level:
--        putStrLnWriteRead $ "growing capacity: " ++ (show newBlockNum)
        inode@Inode{metaData=InodeMetadata{inode_num=inodeNum}}
            <- ensureLevel inodeIn newBlockNum
        allocateBlockWrite (\ bb -> do
                da <- unsafeLiftIORW $ getDiskAddr bb
--                 putStrLnWriteRead ("allocated new address: " ++ (show da)) >>
               -- safe to coerce since it's a small file:
                case level $ metaData inode of
                   0 -> unimplemented "Zero optimization (ensureCapacity')"
                   1 -> ensureCapacity' (updatePointers inode
                                                   [(inIntToInt currBlocksInFile, da)])
                                        (currBlocksInFile + 1)
                                        newBlockNum
                   _ -> do when (abs newBlockNum `mod` intToINInt blockPointersPerIndirectBlock == 0) $ do
                                allocateBlockWrite (\ new_bb -> do
--                                FIX: This does not work for 32M+ files (stops with array out of bounds).
--                                 unsafeLiftIORW $ print "allocating extra space ..."
                                   zeroBlock new_bb
                                   new_da <- unsafeLiftIORW $ getDiskAddr new_bb
                                   let ix = inIntToInt $ abs newBlockNum `div` intToINInt blockPointersPerIndirectBlock
                                   let i = updatePointers inode [(ix, new_da)]
--                                 unsafeLiftIORW $ print ("allocated extra space ...",ix,new_da,i)
                                   updateInodeCacheWrite i)
                           inode1 <- getInodeWrite inodeNum Nothing
                           writeAddressAt inode1 newBlockNum (assert (da /= 0) da)
                           -- inode may have changed:
                           inode2 <- getInodeWrite inodeNum Nothing
                           ensureCapacity' inode2 (currBlocksInFile + 1) newBlockNum)

maxBlockNum :: INInt -- ^Level
            -> BlockNumber
maxBlockNum 1 = intToINInt blockPointersPerInode - 1
maxBlockNum n = (intToINInt blockPointersPerIndirectBlock ^ (n - 1)
                 * intToINInt blockPointersPerInode) - 1
                -- -1, since its counting from zero


-- |Put a bunch of zeros into this block.
zeroBlock :: BufferBlock s -> FSWrite ()
zeroBlock buffer = unsafeLiftIORW $ zeroBufferBlock buffer


-- |When debugging, pPut a bunch of 0xDEADBEEF's into this block;
-- uses just after allocating a block, so we know if we have an uninitalized value.

{- UNUSED:
deadBeefBlock :: DiskAddress -> FSWrite ()
deadBeefBlock addr =
  getBlockFromCacheOrDeviceWrite addr (\ bb -> unsafeLiftIORW $ deadbeefBufferBlock bb)
-}

-- |Writes this disk address into the indirect block pointed to by
-- this list of addresses.  Doesn't update the inode, because
-- ensureLevel should have already done that.
writeAddressAt :: Inode -> BlockNumber -> DiskAddress -> FSWrite ()
writeAddressAt inode blockNum da = do
--    unsafeLiftIOWrite $ print $ "writeAddressAt: " ++ show da
    getAndPositionBlockWrite inode blockNum (\ bb (Just cursor) -> do
--          unsafeLiftIOWrite $ print $ "(1)writeAddressAt" ++ show (bbDiskAddr bb)
            unsafeLiftIOWrite $ writeToBufferBlock bb cursor [Alloc da]
            markBlockDirtyWrite bb)

-- |Optionally increase the level of this inode to include the given
-- block number.  This allocates new blocks, so is in the FSWrite
-- monad.  We grab a new block (which is an indirect block of block
-- pointers), dump the inode's block pointers into it, then update the
-- inode's block pointers to be empty except for an entry pointing to
-- ths new block.  Finally, we update the level.

ensureLevel :: Inode
            -> BlockNumber -- ^Block number we want to include
            -> FSWrite Inode
ensureLevel inode newBlockNum
    | maxBlockNum (level $ metaData inode) >= (abs newBlockNum)
        = updateInodeCacheWrite inode
             >> return inode
    | otherwise = do
        allocateBlockWrite (\ newBB ->  do
            let currLevel = level $ metaData inode
            let inodeNum = inode_num $ metaData inode

             -- dump the block pointers from inode into new block
            putPointersBin newBB inode

            da <- unsafeLiftIORW $ getDiskAddr newBB
            -- inode may have changed after "getBlockFromCacheOrDeviceWrite"
            -- that probably doesn't matter, because we're blowing away the pointers
            -- here anyway, but just in case:
            inode1 <-
             (do inode2 <- getInodeWrite inodeNum Nothing
                 -- update the inode to point to this new block instead;
                 -- erasing existing list.
                 return inode2{blockPtrs=
                                  listArray (0, blockPointersPerInode - 1)
                              (da:(replicate (blockPointersPerInode -1) 0))
                              -- bump the actual level and repeat.
                              ,metaData=(metaData inode){level = currLevel+1}
                             })
            ensureLevel inode1 newBlockNum)

-- |Performs a VERY complex computation.  Don't peak inside.
numBlocksForSize :: Inode -> BlockNumber
numBlocksForSize Inode{metaData=InodeMetadata{num_bytes=numBytes}}
       -- Rational is arbitrary precision, so this should be safe:
    = ceiling $ (toRational numBytes) / (toRational bytesPerBlock)