module ZMachine.Base where

import ZMachine.IO.Base

import Data.Word
import Data.Bits (shiftR,(.&.))
import Data.Array.Diff
import Data.Array.Unboxed
import Control.Monad.State

type Addr = Word
type Val  = Word16

data ZFrame = ZFrame { zfReturnAddr :: Addr,
                       zfLocals     :: UArray Val Val,
                       zfReturnVar  :: Maybe Val,
                       zfArgCount   :: Val,
                       zfEvalStack  :: [Val] }

data ZState = ZState { dynMem :: DiffUArray Addr Word8
                     , story :: UArray Addr Word8
                     , stack :: [ZFrame]
                     , localStack :: [Val]
                     , localVars :: UArray Val Val
                     , argCount :: Val
                     , ptr :: Addr
                     , zs_undo :: Maybe ZState
                     , zs_io_ :: IOOperation -> IO ()
                     }

type ZM a = StateT ZState IO a


-- Address unpacking

decodeRoutineAddr, decodeStringAddr :: Val -> ZM Addr
decodeRoutineAddr = decodePaddr
decodeStringAddr  = decodePaddr

decodePaddr :: (Monad m, Num b, Integral a) => a -> m b
decodePaddr = return . (4 *) . fromIntegral

-- Memory access

getByte :: Addr -> ZM Val
getByte addr = do st <- get
                  return $ fromIntegral $
                    if inRange (bounds (dynMem st)) addr
                       then dynMem st ! addr
                       else story  st ! addr

getWord :: Addr -> StateT ZState IO Val
getWord addr = do a <- getByte addr
                  b <- getByte (addr+1)
                  return (256*a + b)

putByte :: Addr -> Val -> ZM ()
putByte addr val =
    do st <- get
       put (st { dynMem = dynMem st // [(addr, fromIntegral val)]})

putWord :: Addr -> Val -> ZM ()
putWord addr val =
    do putByte addr     (val `shiftR` 8)
       putByte (addr+1) (val .&. 0xff)

getNextByte :: ZM Val
getNextByte = do st <- get
                 x <- getByte (ptr st)
                 put (st { ptr = ptr st + 1 })
                 return x

getNextWord :: ZM Val
getNextWord = do a <- getNextByte
                 b <- getNextByte
                 return (256*a + b)

getVar :: Val -> ZM Val
getVar var
    | var == 0  = do st <- get
                     let val = head (localStack st)
                     put (st { localStack = tail (localStack st) })
                     return val
    | var < 16  = liftM ((!(var-1))) $ gets localVars
    | otherwise = do varTable <- getWord 0xC
                     getWord (fromIntegral $ varTable + (var-16)*2)

putVar :: Val -> Val -> ZM ()
putVar var val
    | var == 0  =
        do st <- get
           put (st { localStack = val : localStack st })
    | var < 16  =
        do st <- get
           put (st { localVars = localVars st // [(var-1, val)] })
    | otherwise =
        do varTable <- getWord 0xC
           putWord (fromIntegral $ varTable + (var-16)*2) val