module Language.Brainfuck (
          evaluate
        , makeState
        , BFState(..)
        ) where

import Data.ByteString (ByteString)
import Data.List (unfoldr, foldl')
import Data.Maybe
import Data.Word8 (Word8)

import qualified Data.ByteString       as BS
import qualified Data.ByteString.Char8 as C8
import qualified Data.Word8            as W8

data BFState = BFState { program :: ByteString,
                         input   :: ByteString,
                         memory  :: ByteString,
                         eip     :: Int       ,
                         eax     :: Int       ,
                         jstack  :: [Int]     ,
                         output  :: String    }

makeState :: String -> String -> Int -> BFState
makeState prg inp memSize = BFState { program = C8.pack prg,
                                      input   = C8.pack inp,
                                      memory  = BS.replicate memSize (0 :: Word8),
                                      eip     = 0,
                                      eax     = 0,
                                      jstack  = [],
                                      output  = "" }

data FNWSState = FNWSState { stackv :: Int ,
                             lasti  :: Int }

evaluate :: BFState -> String
evaluate s =
    catMaybes $ unfoldr (\ state ->
            let eip' = eip state in
            let eax' = eax state in
            let prg  = program state in
            let mem  = memory  state in
            let ns   = state { eip = eip' + 1 } in
            if eip' >= BS.length prg then Nothing
            else
                case prg `C8.index` eip' of
                    '+' -> Just (Nothing, ns { memory = putW8At mem (mem `BS.index` eax' + 1) eax' })
                    '-' -> Just (Nothing, ns { memory = putW8At mem (mem `BS.index` eax' - 1) eax' })
                    '>' -> Just (Nothing, ns { eax = eax' + 1 })
                    '<' -> Just (Nothing, ns { eax = eax' - 1 })
                    '.' -> Just (Just $ mem `C8.index` eax', ns)
                    ',' ->
                        let inp  = input state in
                        let mem' = (if BS.null inp then mem else putChAt mem (C8.head inp) eax') in
                        let inp' = if BS.null inp then inp else BS.tail inp
                        in Just (Nothing, ns { memory = mem', input = inp' })
                    '[' ->
                        if mem `BS.index` eax' == 0
                        then case findNextWithStack prg '[' ']' (eip' + 1) of
                                 Just i  -> Just (Nothing, state { eip = i + 1 })
                                 Nothing -> Nothing -- RIP
                        else Just (Nothing, ns { jstack = eip' : jstack state })
                    ']' ->
                        let st = jstack state in
                        if null st || head st == 0 || mem `BS.index` eax' == 0
                        then Just (Nothing, ns)
                        else Just (Nothing, state { eip = head st + 1, jstack = tail st })

                    _ -> Just (Nothing, ns)
        ) s
    where
        findNextWithStack bs inc dec start =
            let bs' = BS.drop (start - 1) bs in
            let i = lasti $ foldl' (\ s (i, e) ->
                    case e of
                        c | c == inc -> s { stackv = stackv s + 1 }
                        c | c == dec -> if stackv s == 0
                               then s { stackv = -1, lasti = i }
                               else s { stackv = stackv s - 1 }
                        _   -> s
                    ) (FNWSState 0 (-1)) $ (zip [ 0 .. (BS.length bs' - 1) ] $ C8.unpack bs')
            in if i < 0 then Nothing else Just i
        map2     f l   = map (\ (a, b) -> f a b) l
        mapIndex f l   = map2 f $ zip [0 .. (length l - 1) ] l
        putChAt bs c i = C8.pack $ mapIndex (\ i' e -> if i == i' then c else e) $ C8.unpack bs
        putW8At bs b i = BS.pack $ mapIndex (\ i' e -> if i == i' then b else e) $ BS.unpack bs