module Main where

import Data.IORef
import Data.Bits
import Data.Word
import Data.Array.Unboxed
import Data.Array.IO
import Data.Array.MArray
import qualified Data.Map as Map
import Control.Monad.State
import Control.Monad
import Data.List
import System.IO
import System
import Data.Maybe

data MachineState = MachineState { registers :: IOUArray Int Word32
                                 , array0 :: IOUArray Word32 Word32
                                 , arrays :: Map.Map Word32 (IOUArray Word32 Word32)
                                 , finger :: IORef Word32}

mkMachineState r a f = MachineState r (fromJust (Map.lookup 0 a)) a f

updateArrays m a = mkMachineState (registers m) a (finger m)

updateBackup m b = MachineState (registers m) (array0 m) (arrays m) (finger m)

copyState :: MachineState -> IO MachineState
copyState m = do fx <- readIORef (finger m)
                 f' <- newIORef fx
                 r' <- copyArray' (registers m)
                 arrays' <- sequence [do {a' <- copyArray a; return (i,a')}
                                     |(i,a) <- Map.assocs (arrays m)]
                 return $ mkMachineState r' (Map.fromAscList arrays') f'

type Machine a = StateT MachineState IO a

initalState :: Handle -> IO MachineState
initalState program = do f<- newIORef 0
                         l<- hFileSize program
                         a<- newArray (0,8) 0
                         let maxi = (fromInteger $ (l `div` 4)-1)::Word32
                         p<- newArray_ (0,maxi)
                         sequence_ [bigE p i | i<-[0..maxi]]
                         return $ mkMachineState a (Map.singleton 0 p) f
 where 
  bigE p i = do a <- hGetChar program
                b <- hGetChar program
                c <- hGetChar program
                d <- hGetChar program
                let  r = ((fromIntegral $ fromEnum a) `shiftL` 24) .|.
                         ((fromIntegral $ fromEnum b) `shiftL` 16) .|.
                         ((fromIntegral $ fromEnum c) `shiftL` 8) .|.
                         (fromIntegral $ fromEnum d) :: Word32
                writeArray p i r

getRegister :: Int -> Machine Word32
getRegister i = do m <- get
                   liftIO $ readArray (registers m) i

setRegister :: Int -> Word32 -> Machine ()
setRegister i v = do m <- get
                     liftIO $ writeArray (registers m) i v

copyArray :: (IOUArray Word32 Word32) -> IO (IOUArray Word32 Word32)
copyArray a = do a' <- freeze a
                 unsafeThaw (a'::UArray Word32 Word32)

-- Help me not duplicate code! 
copyArray' :: (IOUArray Int Word32) -> IO (IOUArray Int Word32)
copyArray' a = do a' <- freeze a
                  unsafeThaw (a'::UArray Int Word32)

getArray :: Word32 -> Word32 -> Machine Word32
getArray 0 i = do m <- get
                  liftIO $ readArray (array0 m) i
getArray n i = do m <- get
                  a <- Map.lookup n (arrays m)
                  liftIO $ readArray a i

setArray :: Word32 -> Word32 -> Word32 -> Machine ()
setArray n i v = do m <- get
                    a <- Map.lookup n (arrays m)
                    liftIO $ writeArray a i v

addArray :: Word32 -> Machine Word32
addArray s = do m <- get
                a <- liftIO $ newArray (0,s-1) 0
                let (mi,_) = Map.findMax (arrays m)
                let mn = mi+1
                put $ updateArrays m (Map.insert mn a (arrays m))
                return mn

removeArray :: Word32 -> Machine ()
removeArray n = do m <- get
                   put $ updateArrays m (Map.delete n (arrays m))

dupArray :: Word32 -> Word32 -> Machine ()
dupArray x y = do m <- get
                  a <- Map.lookup x (arrays m)
                  a' <- liftIO $ copyArray a
                  put $ updateArrays m (Map.insert y a' (arrays m))

setIP :: Word32 -> Machine ()
setIP ip = do m <- get
              liftIO $ writeIORef (finger m) ip

getIP :: Machine Word32
getIP = do m <- get
           liftIO $ readIORef (finger m)

nextIP :: Machine ()
nextIP = do m <- get
            ip <- liftIO $ readIORef (finger m)
            liftIO $ writeIORef (finger m) (succ ip)

op0 a b c = do cx <- getRegister c
               unless (cx==0) (getRegister b >>= setRegister a)
               nextIP

op1 a b c = do bx <- getRegister b
               cx <- getRegister c
               v <- getArray bx cx 
               setRegister a v
               nextIP

op2 a b c = do ax <- getRegister a
               bx <- getRegister b
               cx <- getRegister c
               setArray ax bx cx
               nextIP

op3 a b c = do bx <- getRegister b
               cx <- getRegister c
               setRegister a (bx+cx)
               nextIP

op4 a b c = do bx <- getRegister b
               cx <- getRegister c
               setRegister a (bx*cx)
               nextIP

op5 a b c = do bx <- getRegister b
               cx <- getRegister c
               setRegister a (bx `div` cx)
               nextIP

op6 a b c = do bx <- getRegister b
               cx <- getRegister c
               setRegister a (complement $ bx .&. cx)
               nextIP

op7 :: Machine ()
op7 = fail "Halting"

op8 b c = do cx <- getRegister c
             r <- addArray cx
             setRegister b r
             nextIP

op9 c = do cx <- getRegister c
           removeArray cx
           nextIP

op10 c = do cx <- getRegister c
            liftIO . putChar . toEnum $ fromIntegral cx
            nextIP

op11 c = do v <- liftIO getChar
            setRegister c (fromIntegral (fromEnum v))
            nextIP

op12 b c = do bx <- getRegister b
              cx <- getRegister c
              when (bx/=0) $ dupArray bx 0
              setIP cx

op13 a v = do setRegister a v
              nextIP

decode w = case op of 0 -> op0 a b c
                      1 -> op1 a b c
                      2 -> op2 a b c
                      3 -> op3 a b c
                      4 -> op4 a b c
                      5 -> op5 a b c
                      6 -> op6 a b c
                      7 -> op7
                      8 -> op8 b c
                      9 -> op9 c
                      10 -> op10 c
                      11 -> op11 c
                      12 -> op12 b c
                      13 -> op13 a' v
 where
  c = fromIntegral $ 0x7 .&. w
  b = fromIntegral $ 0x7 .&. (w `shiftR` 3)
  a = fromIntegral $ 0x7 .&. (w `shiftR` 6)
  a' = fromIntegral $ 0x7 .&. (w `shiftR` 25)
  v = fromIntegral $ 0x1ffffff .&. w
  op = (w `shiftR` 28)

tick :: Machine ()
tick = do m <- get
          ip <- getIP
          w <- getArray 0 ip
          decode w

main = do [fn] <- getArgs
          h <- openBinaryFile fn ReadMode
          s <- initalState h
          hSetBuffering stdout NoBuffering
          runStateT (sequence_ (repeat tick)) s

fork :: Machine MachineState
fork = do m <- get
          liftIO $ copyState m