module Main where

-- there are bugs yet with missing loop marks at beginning or end.

import Control.Monad.State

data X29AState = X29AState {register::Int,
                            stack :: [X29AFunction]} deriving Show

data X29AFunction = PrimS | PrimK | PrimOut | PrimIn | PrimInc | PrimDec | 
                    Apply X29AFunction X29AFunction deriving Show

type X29AMonad a = StateT X29AState IO a

data X29AInsn = SingleInsn (X29AMonad ()) |
                Loop [X29AInsn]

evaluate :: X29AFunction -> Maybe (X29AMonad ())
evaluate (Apply (Apply (Apply PrimS x) y) z) = Just $ push (Apply (Apply x z) (Apply y z))
evaluate (Apply (Apply PrimK x) y) = Just $ push x
evaluate (Apply (Apply PrimOut x) y) = Just $ do push x
                                                 r <- gets register
                                                 liftIO $ print r
--                                                 liftIO $ putChar (toEnum $ fromIntegral r)
evaluate (Apply (Apply PrimIn x) y) = Just $ do push x
                                                c <- liftIO getChar
                                                modify (\x -> x {register = (fromIntegral $ fromEnum c)})
evaluate (Apply (Apply PrimInc x) y) = Just $ do push x
                                                 r <- gets register
                                                 modify (\x -> x {register = r+1})
evaluate (Apply (Apply PrimDec x) y) = Just $ do push x
                                                 r <- gets register
                                                 modify (\x -> x {register = r-1})
evaluate f = Nothing

push :: X29AFunction -> X29AMonad ()
push i = do s <- gets stack
            modify (\x -> x {stack = i:s})

pop :: X29AMonad X29AFunction
pop = do s <- gets stack
         if null s
           then return $ Apply (Apply PrimS PrimK) PrimS
           else modify (\x -> x {stack = tail s}) >> (return $ head s)

swap :: X29AMonad ()
swap = do a <- pop
          b <- pop
          push a
          push b

apply :: X29AMonad ()
apply = do a <- pop
           b <- pop
           push $ Apply b a

parse0x29A :: String -> Maybe (X29AInsn, String)
parse0x29A ('s':rest) = Just (SingleInsn $ push PrimS, rest)
parse0x29A ('k':rest) = Just (SingleInsn $ push PrimK, rest)
parse0x29A ('.':rest) = Just (SingleInsn $ push PrimOut, rest)
parse0x29A (',':rest) = Just (SingleInsn $ push PrimIn, rest)
parse0x29A ('+':rest) = Just (SingleInsn $ push PrimInc, rest)
parse0x29A ('-':rest) = Just (SingleInsn $ push PrimDec, rest)
parse0x29A ('[':rest) = Just (Loop insns, afterloop) where
     (insns,afterloop) = parse0x29Ablock rest
parse0x29A (']':rest) = Nothing
parse0x29A ('%':rest) = Just (SingleInsn $ swap, rest)
parse0x29A ('~':rest) = Just (SingleInsn $ apply, rest)
paese0x29A (_:rest) = Just (return (),rest) -- FIXME: performance!

parse0x29Ablock :: String -> ([X29AInsn],String)
parse0x29Ablock "" = ([],"")
parse0x29Ablock s = case parse0x29A s of
                       Nothing -> ([],"")
                       Just (i,s') ->
                           case parse0x29Ablock s' of
                             (is,s'') -> (i:is,s'')

parse0x29Aprog :: String -> [X29AInsn]
parse0x29Aprog str =
  case parse0x29Ablock str of
    (prg,"") -> prg
    (prg,rest) -> Loop prg : parse0x29Aprog rest

loop0x29A :: [X29AInsn] -> X29AMonad ()
loop0x29A is = do
                 r <- gets register
                 when (r /= 0) $ exec0x29A is >> loop0x29A is
                             
exec0x29A :: [X29AInsn] -> X29AMonad ()
exec0x29A [] = return ()
exec0x29A (SingleInsn i: is) = i >> execTOS >> exec0x29A is
exec0x29A (Loop l : is) = do
         loop0x29A l
         exec0x29A is

execTOS :: X29AMonad ()
execTOS = do
  s <- gets stack
  when (not $ null s) (case evaluate (head s) of
                             Nothing -> return ()
                             Just f -> pop >> f >> execTOS)

run :: String -> IO ()
run = (\i -> execStateT (exec0x29A i) (X29AState 0 []) >>= print ). parse0x29Aprog