module Lib
    ( interpret
    ) where

import Data.Char
import Data.Maybe
import System.IO

data BFInstruction = MemoryRight | MemoryLeft | Increment | Decrement | Output | Input | LoopBegin | LoopEnd deriving (Enum)
data BFProgram = BFProgram [BFInstruction] BFInstruction [BFInstruction]

newtype BFMemoryCell = BFMemoryCell Int
data BFMemory = BFMemory [BFMemoryCell] BFMemoryCell [BFMemoryCell]

startProgram :: [BFInstruction] -> BFProgram
startProgram instructions = BFProgram [] (head instructions) (tail instructions)

advance :: BFProgram -> BFProgram
advance (BFProgram past current next) = BFProgram (past ++ [current]) (head next) (tail next)

advanceBy :: Int -> BFProgram -> BFProgram
advanceBy advancement = case advancement of
    x | x < 0   -> decreaseBy' (-advancement)
    0           -> id
    x | x > 0   -> advanceBy' advancement

advanceBy' :: Int -> BFProgram -> BFProgram
advanceBy' advancement (BFProgram past current next) = do
    let next' = current:next
    let next'' = drop advancement next'
    BFProgram (past ++ take advancement next') (head next'') (tail next'')

decreaseBy' :: Int -> BFProgram -> BFProgram
decreaseBy' decrement (BFProgram past current next) = do
    let past' = past ++ [current]
    let length' = length past' - decrement
    let past'' = take length' past'
    BFProgram (init past'') (last past'') (drop length' past' ++ next)

makeCell :: Int -> BFMemoryCell
makeCell = BFMemoryCell . wrap

incrementCell :: BFMemoryCell -> BFMemoryCell
incrementCell = makeCell . (+1) . getCell

decrementCell :: BFMemoryCell -> BFMemoryCell
decrementCell = makeCell . subtract 1 . getCell

getCell :: BFMemoryCell -> Int
getCell (BFMemoryCell value) = value

wrap :: Int -> Int
wrap input = mod input 256

moveMemoryRight :: BFMemory -> BFMemory
moveMemoryRight (BFMemory previous current []) = BFMemory (previous ++ [current]) (makeCell 0) []
moveMemoryRight (BFMemory previous current next) = BFMemory (previous ++ [current]) (head next) (tail next)

moveMemoryLeft :: BFMemory -> BFMemory
moveMemoryLeft (BFMemory [] current next) = BFMemory [] (makeCell 0) (current:next)
moveMemoryLeft (BFMemory previous current next) = BFMemory (init previous) (last previous) (current:next)

onCurrentCell :: (BFMemoryCell -> BFMemoryCell) -> BFMemory -> BFMemory
onCurrentCell func (BFMemory previous current next) = BFMemory previous (func current) next

toInstructions :: String -> [BFInstruction]
toInstructions = mapMaybe toInstruction

toInstruction :: Char -> Maybe BFInstruction
toInstruction instruction = case instruction of
    '>' -> Just MemoryRight
    '<' -> Just MemoryLeft
    '+' -> Just Increment
    '-' -> Just Decrement
    '.' -> Just Output
    ',' -> Just Input
    '[' -> Just LoopBegin
    ']' -> Just LoopEnd
    _   -> Nothing

interpret :: String -> IO (Either String BFMemory)
interpret program = step (startProgram $ toInstructions program) (BFMemory [] (makeCell 0) [])

step :: BFProgram -> BFMemory -> IO (Either String BFMemory)
step (BFProgram _ _ []) memory = return . Right $ memory
step program@(BFProgram previousProgram instruction nextProgram) memory@(BFMemory previousMemory currentMemory nextMemory) = case instruction of
    MemoryRight -> step (advance program) (moveMemoryRight memory)
    MemoryLeft  -> step (advance program) (moveMemoryLeft memory)
    Increment   -> step (advance program) (onCurrentCell incrementCell memory)
    Decrement   -> step (advance program) (onCurrentCell decrementCell memory)
    Output      -> do
        putChar . chr . getCell $ currentMemory
        hFlush stdout
        step (advance program) memory
    Input       -> do
        newCurrentChar <- getChar
        let newCurrent = makeCell . ord $ newCurrentChar
        step (advance program) (BFMemory previousMemory newCurrent nextMemory)
    LoopBegin   -> case getCell currentMemory of
        0   -> case findMatchingLoopClose nextProgram 0 0 of
            Left err            -> return . Left $ err
            Right advancement   -> step (advanceBy (advancement + 1) program) memory
        _   -> step (advance program) memory
    LoopEnd     -> case getCell currentMemory of
        0   -> step (advance program) memory
        _   -> case findMatchingLoopOpen previousProgram 0 0 of
            Left err        -> return . Left $ err
            Right decrease  -> step (advanceBy (-(decrease + 1)) program) memory

findMatchingLoopClose :: [BFInstruction] -> Int -> Int -> Either String Int
findMatchingLoopClose [] _ _ = Left "findMatchingLoopClose: No matching ] found"
findMatchingLoopClose (instruction:remainingProgram) index nestingCounter = case instruction of
    LoopBegin   -> findMatchingLoopClose remainingProgram (index + 1) (nestingCounter + 1)
    LoopEnd     -> if nestingCounter == 0
        then Right index
        else findMatchingLoopClose remainingProgram (index + 1) (nestingCounter - 1)
    _           -> findMatchingLoopClose remainingProgram (index + 1) nestingCounter

findMatchingLoopOpen :: [BFInstruction] -> Int -> Int -> Either String Int
findMatchingLoopOpen = findMatchingLoopOpen' . reverse

findMatchingLoopOpen' :: [BFInstruction] -> Int -> Int -> Either String Int
findMatchingLoopOpen' [] _ _ = Left "findMatchingLoopOpen: No matching [ found"
findMatchingLoopOpen' (instruction:remainingProgram) index nestingCounter = case instruction of
    LoopEnd     -> findMatchingLoopOpen' remainingProgram (index + 1) (nestingCounter + 1)
    LoopBegin   -> if nestingCounter == 0
        then Right index
        else findMatchingLoopOpen' remainingProgram (index + 1) (nestingCounter - 1)
    _           -> findMatchingLoopOpen' remainingProgram (index + 1) nestingCounter