{-# OPTIONS_GHC -fglasgow-exts -package mtl -cpp #-}
module MachineState 
   ( Platter, Offset, RegisterId, PlatterArrayId, PlatterArray, RegisterValue
   , loadProgram, putProgram, getProgram, getCurrentPlatter
   , allocateArray, deallocateArray, ammendArray, indexArray
   , advanceFinger, getFinger, retreatFinger
   , getRegister, putRegister
   , halt, isHalted 
   , readCodex
   , saveState, loadState, checkIsState
   ) where

import Control.Monad
import qualified Data.Map as M
import Data.Array.IO
import Data.Array(assocs)
import Data.IORef
import Data.Word
import Data.Bits
import Data.Char (chr,ord)
import Foreign.Marshal.Alloc
import Foreign.Ptr
import Foreign.Storable
import System.IO
import System.IO.Unsafe

import Prelude hiding (readList)

programRef   :: IORef PlatterArray
arraysRef    :: IORef PlatterArrayCollection
freshIdRef   :: IORef PlatterArrayId
registersRef :: Registers
fingerRef    :: IORef Offset
isRunningRef :: IORef Bool

programRef   = unsafePerformIO $ newIORef $ error "programRef not initialized"
arraysRef    = unsafePerformIO $ newIORef M.empty
freshIdRef   = unsafePerformIO (newIORef 1)
registersRef = unsafePerformIO $ newArray (0, 7) 0
fingerRef    = unsafePerformIO (newIORef 0)
isRunningRef = unsafePerformIO (newIORef True)

type Platter      = Word32
type Offset       = Word32
data PlatterArray = PlatterArray Int (Ptr Platter)
type PlatterArrayId         = Platter
type PlatterArrayCollection = M.Map PlatterArrayId PlatterArray

type RegisterId    = Platter
type RegisterValue = Platter
type Registers     = IOArray RegisterId RegisterValue

getRegister     :: RegisterId -> IO RegisterValue
getRegister rid = readArray registersRef rid

putRegister        :: RegisterId -> RegisterValue -> IO ()
putRegister rid rv = do
   writeArray registersRef rid rv

getArray     :: PlatterArrayId -> IO PlatterArray
getArray 0 = getProgram
getArray aid =  do
  m <- readIORef arraysRef
  M.lookup aid m
  
putArray        :: PlatterArrayId -> PlatterArray -> IO ()
putArray 0 pa   = putProgram pa
putArray aid pa = 
   modifyIORef arraysRef (M.insert aid pa)

getProgram :: IO PlatterArray
getProgram = readIORef programRef

putProgram    :: PlatterArray -> IO ()
putProgram pa = writeIORef programRef pa 

getCurrentPlatter :: IO Platter
getCurrentPlatter = do
  prog   <- getProgram
  f      <- getFinger 
  (prog !!! f) 
  
indexArray       :: PlatterArrayId -> Offset -> IO Platter
indexArray aid j =  do
  pa <- getArray aid
  p  <- pa !!! j
  return p 

ammendArray         :: PlatterArrayId -> Offset -> Platter -> IO ()
ammendArray aid j p =  do
  pa <- getArray aid  
  pa /// (j, p)
  return () 

allocateArray   :: Platter -> IO PlatterArrayId
allocateArray p =  let n = toInt p
                   in  do
  aid <- fresh
  pa  <- platterArray n -- (replicate n 0)
  putArray aid pa
  return aid

deallocateArray :: Platter -> IO ()
deallocateArray p = do
   (PlatterArray _ ptr) <- getArray p
   modifyIORef arraysRef (M.delete p)
   free ptr

loadProgram       :: PlatterArrayId -> Offset -> IO ()
loadProgram 0   j =  putFinger j
loadProgram aid j =  do
  pa <- getArray aid
  paCopy <- dupPlatterArray pa
  (PlatterArray _ oldPtr) <- getProgram
  putProgram paCopy
  free oldPtr
  putFinger j

getFreshId :: IO PlatterArrayId
getFreshId =  readIORef freshIdRef

fresh :: IO PlatterArrayId
fresh =  do
   aid <- readIORef freshIdRef
   writeIORef freshIdRef (aid+1)
   return aid
   
getFinger :: IO Offset
getFinger = readIORef fingerRef

putFinger   :: Offset -> IO ()
putFinger j = writeIORef fingerRef j 

advanceFinger :: IO ()
advanceFinger = modifyIORef fingerRef (+1)

retreatFinger :: IO ()
retreatFinger = modifyIORef fingerRef (\x -> x-1)

halt :: IO ()
halt = writeIORef isRunningRef False

isHalted :: IO Bool
isHalted =  do
  ms <- readIORef isRunningRef
  return (not ms)
  
(!!!)                    :: PlatterArray -> Offset -> IO Platter
PlatterArray _ ptr !!! j =  peekElemOff ptr (toInt j)

(///)                         :: PlatterArray -> (Offset, Platter) -> IO ()
PlatterArray _ ptr /// (j, p) =  pokeElemOff ptr (toInt j) p

platterArray    :: Int -> IO PlatterArray
platterArray n =  do
  ptr <- mallocBytes (4 * n)
  memset ptr 0 n
  return (PlatterArray n ptr)

toInt :: Platter -> Int
toInt =  fromInteger . toInteger

fromInt :: Int -> Platter
fromInt =  fromInteger . toInteger

readCodex :: String -> IO PlatterArray 
readCodex file =  do
  h   <- openBinaryFile file ReadMode
  sz  <- hFileSize h >>= return . fromInteger -- hack: large files?
  pa <- readPlatterArray h (sz `div` 4)
  hClose h

  return pa


-- Swaps the bytes of every word in PlatterArray
-- TODO: implement in C, because right now this is slow
-- Used to make the program portable across different endian architectures
swapPlatterArray :: PlatterArray -> IO ()

#if defined(i386_TARGET_ARCH)
foreign import ccall unsafe "swap.h swapArray" swapArray
    :: Ptr Word32 -> Int -> IO ()

swapPlatterArray (PlatterArray size ptr) = swapArray ptr size
#else
swapPlatterArray _ = return ()
#endif


-- Save and restore functions
saveStateIO :: Handle -> IO ()
saveStateIO h = do
  saveWord32 h 0xabcd
  readIORef programRef >>= savePlatterArray h
  readIORef arraysRef >>= saveMap h
  readIORef freshIdRef >>= saveWord32 h
  saveRegs h
  readIORef fingerRef >>= saveWord32 h
  readIORef isRunningRef >>= saveBool h
  
readState :: Handle -> IO ()
readState h = do
  readWord32 h
  readPlatterArrayIO h >>= writeIORef programRef
  putStrLn "program loaded..."
  readMap h >>= writeIORef arraysRef
  putStrLn "arrays loaded..."
  readWord32 h >>= writeIORef freshIdRef
  readRegs h
  readWord32 h >>= writeIORef fingerRef
  readBool h >>= writeIORef isRunningRef

saveState file = do
  h <- openFile file WriteMode
  saveStateIO h
  hClose h

loadState name = do
  putStrLn "Loading state..."
  h <- openFile name ReadMode
  readState h
  hClose h
  putStrLn "State loaded."

checkIsState name = do
  h <- openFile name ReadMode
  w <- readWord32 h
  hClose h
  return (w== 0xabcd)

saveMap h m = do
  let (keys,elems) = unzip $ M.toList m
  saveWord32 h (fromInt $ length elems)
  saveList h saveWord32 keys
  saveList h savePlatterArray elems
  return ()

readMap h = do
  wl <- readWord32 h
  let nelems = toInt wl
  putStrLn ("Loading "++show nelems++" arrays...")
  keys <- readList h readWord32 nelems []
  elems <- readList h readPlatterArrayIO nelems []
  return (M.fromList (zip keys elems))

saveBool h True = saveWord32 h 1
saveBool h False = saveWord32 h 0

readBool h = readWord32 h >>= \x -> return (x==1)

saveRegs h = do
--  let elems = map snd (assocs registersRef)
  mapM getRegister [0..7] >>= saveList h saveWord32

readRegs h = do
  elems <- readList h readWord32 8 []
  let idxelems = zip [0..7] (reverse elems)
  mapM_ (\(idx,e) -> writeArray registersRef idx e) idxelems
  



saveList h f ls = mapM_ (f h) ls

-- Warning, reads in reverse!
readList h f 0 res = return res
readList h f n res = f h >>= \ x -> readList h f (n-1) (x:res)
--readList h f n = mapM (\_ -> f h) [1..n]

savePlatterArray :: Handle -> PlatterArray -> IO ()
savePlatterArray h pa@(PlatterArray sz ptr) = do
  -- first correct endianness for storage
  swapPlatterArray pa

  -- and store it
  saveWord32 h (fromInt sz)
  hPutBuf h ptr (sz * 4)

  -- Fix it back
  swapPlatterArray pa

readPlatterArrayIO :: Handle -> IO PlatterArray
readPlatterArrayIO h = do
  w <- readWord32 h
  readPlatterArray h (toInt w)

readPlatterArray :: Handle -> Int -> IO PlatterArray
readPlatterArray h sz = do
  ptr <- mallocBytes (sz * 4)
  n <- hGetBuf h ptr (sz * 4)
  let pa = (PlatterArray sz ptr)  

  -- Corrects endianness, if needed
  swapPlatterArray pa

  return pa

readWord32 :: Handle -> IO Word32
readWord32 h
  = do chars <- mapM (const (hGetChar h)) [1..4]
       let nums = map (toInteger . ord) (reverse chars)
           exps = iterate ((*0x100) .) id
       return (fromInteger $ sum $ zipWith ($) exps nums)

saveWord32 :: Handle -> Word32 -> IO ()
saveWord32 h num
  = mapM_ (hPutChar h) (chars num)
  where
  chars num = map (chr . fromInteger . toInteger) $ map (\ e -> (num `shiftR` (8*e)) .&. 0xff) (reverse [0..3])


-- copies a PlatterArray
dupPlatterArray :: PlatterArray -> IO PlatterArray
dupPlatterArray (PlatterArray sz ptr) = do
  ptrDst <- mallocBytes (4 * sz)
  memcpy ptrDst ptr sz
  return (PlatterArray sz ptrDst)

memcpy p1 p2 s = memcpy_c p1 p2 (s*4)
memset p c s = memset_c p c (s*4)

#if defined(__GLASGOW_HASKELL__)
foreign import ccall unsafe "string.h memcpy" memcpy_c
    :: Ptr Word32 -> Ptr Word32 -> Int -> IO ()
foreign import ccall unsafe "string.h memset" memset_c
    :: Ptr Word32 -> Int -> Int -> IO ()
#endif