{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase #-}

module Brainfuck where

import Control.Monad
import Control.Monad.State.Class
import Control.Monad.Trans.State (StateT, evalStateT)
import Control.Monad.IO.Class
import Data.Word
import Data.Maybe
import Data.Char

-- | A simple Zipper Structure used to navigate bothe the code and the cells
newtype Zipper a = Zipper ([a], [a])

-- | Create a Zipper from a list
fromList :: [a] -> Zipper a
fromList l = Zipper ([], l)

-- | Create an infinite zipper from a default value
infiniteZipper :: a -> Zipper a
infiniteZipper x = Zipper (repeat x, repeat x)

-- | Move the Zipper forward
forward :: Zipper a -> Zipper a
forward (Zipper ( _,   [])) = error "Can't go forward on the last element"
forward (Zipper (xs, y:ys)) = Zipper (y:xs, ys)

-- | Move the Zipper back
back :: Zipper a -> Zipper a
back (Zipper (  [],  _)) = error "Can't go back on the first element"
back (Zipper (x:xs, ys)) = Zipper (xs, x:ys)

-- | Peek at the Zipper's current value
peek :: Zipper a -> Maybe a
peek (Zipper (_, x:_)) = Just x
peek                 _ = Nothing

-- | Write at the Zipper's current position
write :: a -> Zipper a -> Zipper a
write _ (Zipper ( _,   [])) = error "Can't write after last element"
write x (Zipper (ys, _:xs)) = Zipper (ys, x:xs)

-- | Apply a function to value at the Zipper's current position
applyTo :: (a -> a) -> Zipper a -> Zipper a
applyTo f z = case peek z of
        Nothing -> error "Can't apply to the end of the zipper"
        Just x -> flip write z . f $ x

-- | The Monad for executing brainfuck code
newtype Brainfuck a = Brainfuck { exitBrainfuck :: StateT BrainfuckState IO a }
    deriving (Functor, Applicative, Monad, MonadState BrainfuckState, MonadIO)

data BrainfuckState = BrainfuckState {
          cells :: Zipper Word8
        , code  :: Zipper Char
        }

runBrainfuck :: String -> IO ()
runBrainfuck code = flip evalStateT initial . exitBrainfuck $ interpret
                    >> (liftIO . putStr $ "
")
    where initial = BrainfuckState (infiniteZipper 0) (fromList code)

-- | peek raised into the Brainfuck Monad
peekB :: (BrainfuckState -> Zipper a) -> Brainfuck (Maybe a)
peekB = fmap peek . gets

-- | Helper to modify the code
modifyCode :: (Zipper Char -> Zipper Char) -> Brainfuck ()
modifyCode f = modify $
        \(BrainfuckState cells code) -> BrainfuckState cells . f $ code

-- | Helper to modify the cells
modifyCells :: (Zipper Word8 -> Zipper Word8) -> Brainfuck ()
modifyCells f = modify $
        \(BrainfuckState cells code) -> flip BrainfuckState code . f $ cells

-- | Move the Zipper to the matching bracket in the code
seekBracket :: Char -> Char -> (Zipper Char -> Zipper Char) -> Brainfuck ()
seekBracket target opposite move = go 0
    where go n = modifyCode move >> peekB code >>= \case
                    Nothing                 -> return ()
                    Just  c | c == target   -> unless (n == 0) $ go (n-1)
                    Just  c | c == opposite -> go (n+1)
                    _                       -> go n

-- | The interpretation loop
interpret :: Brainfuck ()
interpret =
    peekB code >>= \case
        Nothing  -> return ()
        Just c -> do
            handleChar c
            modifyCode forward
            interpret

-- | This is where the magic happens.
handleChar :: Char -> Brainfuck ()
-- | Print the current cell
handleChar '.' = peekB cells >>=
                 liftIO . putStr . (:[]) . chr . fromIntegral . fromJust
-- | Read into the current cell
handleChar ',' = liftIO getChar >>=
                 modifyCells . write . fromIntegral . ord
-- | Move the cell-Zipper forward
handleChar '>' = modifyCells forward
-- | Move the cell-Zipper back
handleChar '<' = modifyCells back
-- | Increment the current cell
handleChar '+' = modifyCells . applyTo $ (+1)
-- | Decrement the current cell
handleChar '-' = modifyCells . applyTo $ (\x -> x-1)
-- | Move the code-Zipper onto the matching ']'
handleChar '[' = (fmap (fromIntegral . fromJust) . peekB $ cells)
                 >>= \case
                    0 -> seekBracket ']' '[' forward
                    _ -> return ()
-- | Move the code-Zipper onto the matching '['
handleChar ']' = (fmap (fromIntegral . fromJust) . peekB $ cells)
                 >>= \case
                    0 -> return ()
                    _ -> seekBracket '[' ']' back
handleChar  _  = return ()