module Interpreter where
-- vim: set ts=4 sw=4 et tw=80:
import Data.Maybe (catMaybes)
import Debug.Trace
import Control.Monad.Trans.State
import Data.Array.IArray

data Instruction = TapeLeft | TapeRight | Add | Sub | Out | In | LoopStart |
    LoopEnd deriving (Eq, Show)

parseInstr :: Char -> Maybe Instruction
parseInstr '<' = Just TapeLeft
parseInstr '>' = Just TapeRight
parseInstr '+' = Just Add
parseInstr '-' = Just Sub
parseInstr '.' = Just Out
parseInstr ',' = Just In
parseInstr '[' = Just LoopStart
parseInstr ']' = Just LoopEnd
parseInstr _ = Nothing

newtype Byte = Byte Int deriving (Eq, Ord, Read, Show)
chr :: Byte -> Char
chr (Byte i) = toEnum i :: Char

updateCell :: Instruction -> Byte -> Byte
updateCell Add (Byte 255) = Byte 0
updateCell Sub (Byte 0) = Byte 255
updateCell Add (Byte i) = Byte (i + 1)
updateCell Sub (Byte i) = Byte (i - 1)

type Tape = [Byte]
type Program = Array Int Instruction

data Memory = Memory { tape :: Tape, dp :: Int } deriving (Show)

-- Consider putting IP in Termination output
data Term = Next Int | Crash Error Int | Done deriving (Show)

toProgram :: String -> Maybe Program
toProgram s = case progL of [] -> Nothing
                            _ -> Just $ listArray (1, (length progL)) progL
    where progL = catMaybes $ fmap parseInstr s

tapeHelp :: Tape -> Int -> (Byte -> Byte) -> Tape
tapeHelp (t:tt) 0 alter = alter t : tt
tapeHelp (t:tt) dp alter = t : (tapeHelp tt (dp-1) alter)

data Error = NegTape | LoopNoMatch deriving (Show)

errorDesc :: Error -> String
errorDesc NegTape = "ERROR: '<' not working. Cannot access negative tape positions."
errorDesc LoopNoMatch = "ERROR: '[' or ']' not working. Matching parenthesis not found."

execUpdate :: Instruction -> State Memory (Maybe Error)
execUpdate i = state stateF
    where stateF p = (Nothing, p { tape = newTape p })
          newTape p = tapeHelp (tape p) (dp p) (updateCell i)

execTapeLeft :: State Memory (Maybe Error)
execTapeLeft = state $ \p -> case dp p of 0 -> (Just NegTape, p)
                                          n -> (Nothing, p { dp = n - 1 })

execTapeRight :: State Memory (Maybe Error)
execTapeRight = state $ \p -> (Nothing, newMem p)
    where newMem p = p { tape = tape p ++ newTape p, dp = dp p + 1 }
          newTape p = case dp p == length (tape p) - 1 of True -> [Byte 0]
                                                          _ -> []

execIn :: StateT Memory IO (Maybe Error)
execIn = StateT $ \p -> do chr <- getChar
                           tp <- pure $ fromEnum chr
                           return (Nothing, newTape tp p)
    where newTape tp p = p { tape = tapeHelp (tape p) (dp p) (\_ -> Byte tp) }

execOut :: StateT Memory IO (Maybe Error)
execOut = StateT $ \p -> putChar (chr (tape p !! dp p)) >> (return (Nothing, p))

findMatching :: Program -> Instruction -> Int -> Maybe Int
findMatching prg par = fmHelper $ -1
    where (oppos, limit, next) = case par of
            LoopStart -> (LoopEnd, length prg, (+1))
            LoopEnd -> (LoopStart, 0, (-1+))

          fmHelper :: Int -> Int -> Maybe Int
          fmHelper acc s
            | s == limit              = Nothing
            | acc == 0 && c == oppos  = Just s
            | c == oppos              = fmHelper (acc - 1) $ next s
            | c == par                = fmHelper (acc + 1) $ next s
            | otherwise               = fmHelper acc $ next s
            where c = prg ! s

execLoop :: Program -> Instruction -> Int -> State Memory Term
execLoop p i ip = state $ \w -> case posOrF w of
                                  Just n -> (Next $ n + 1, w)
                                  Nothing -> (Crash LoopNoMatch ip, w)
    where posOrF ps = case ((tape ps !! dp ps) == Byte 0, i) of
                        (True, LoopStart) -> findMatching p i ip
                        (False, LoopEnd) -> findMatching p i ip
                        _ -> Just ip

execute :: Program -> Memory -> IO (Term, Memory)
execute p = runStateT (runInstruction 1)
    where runInstruction :: Int -> StateT Memory IO Term
          runInstruction ip = fmap checkDone (instrFor ip) >>= dbg >>= loopBack
            where loopBack (Next i) = runInstruction i
                  loopBack x = return x
                  dbg a = StateT $ \s -> do _ <- putStrLn . show $ (a, s)
                                            return (a, s)

          instrFor :: Int -> StateT Memory IO Term
          instrFor ip = case (p ! ip) of
                            TapeLeft -> toTerm ip . toStateT $ execTapeLeft
                            TapeRight -> toTerm ip . toStateT $ execTapeRight
                            Add -> toTerm ip . toStateT $ execUpdate Add
                            Sub -> toTerm ip . toStateT $ execUpdate Sub
                            In -> toTerm ip execIn
                            Out -> toTerm ip execOut
                            LoopStart -> toStateT $ execLoop p LoopStart ip
                            LoopEnd -> toStateT $ execLoop p LoopEnd ip

          toTerm ip = fmap termMem
                where termMem Nothing  = Next $ ip + 1
                      termMem (Just e) = Crash e ip

          toStateT :: (Monad m) => State s a -> StateT s m a
          toStateT x = StateT $ return . (runState x)

          checkDone :: Term -> Term
          checkDone (Next x)
            | x == (snd . bounds $ p) + 1 = Done
            | otherwise                   = Next x
          checkDone x = x


initMemory :: Memory
initMemory = Memory { tape = [Byte 0], dp = 0 }

run :: String -> IO ((Term, Memory))
run str = case toProgram str of Nothing -> return (Done, initMemory)
                                Just pp -> execute pp initMemory