module Control.Monad.BrainFuck where

import qualified Control.Monad as M
import qualified Data.Char as Char

newtype BrainFuck a = BrainFuck (DataPointer -> ([Char], DataPointer, a))

type DataPointer = Integer

-- | Retrieve the inner function
func :: BrainFuck f -> DataPointer -> ([Char], DataPointer, f)
func (BrainFuck f) = f

-- | Evaluate the monad and get a brainfuck program
brainfuck :: BrainFuck f -> String
brainfuck (BrainFuck f) = bytes where (bytes, _, _) = f 0

instance Monad BrainFuck where
	return ret = BrainFuck $ \loc -> ([], loc, ret)
	a >>= b = BrainFuck $ \start -> let
		(left, mid, val) = func a start
		(right, end, ret) = func (b val) mid
		in (left ++ right, end, ret)

next, prev, incr, decr, output, input, open, close :: BrainFuck ()
-- | move data pointer right
next = opcode' succ '>'
-- | move data pointer left
prev = opcode' pred '<'
-- | increment data
incr = opcode '+'
-- | decrement data
decr = opcode '-'
-- | output byte at data pointer
output = opcode '.'
-- | input byte, storing at data pointer
input = opcode ','
-- | if byte at data pointer is zero, jump to command after close
open = opcode '['
-- | if byte at data pointer is nonzero, jump to command after matching open
close = opcode ']'

add, sub :: Int -> BrainFuck ()
add = multi incr
sub = multi decr

-- | Adds an arbitrary character to the program.
-- Should not be used directly.
opcode :: Char -> BrainFuck ()
opcode = opcode' id

-- | Adds an arbitrary character to the program,
-- and updates the data pointer.
-- Should not be used directly.
opcode' :: (DataPointer -> DataPointer) -> Char -> BrainFuck ()
opcode' f x =  BrainFuck $ \loc -> ([x], f loc, ())

-- | Run an action multiple times.
multi :: BrainFuck () -> Int -> BrainFuck ()
multi c n = do
	_ <- sequence (replicate n c)
	return ()

-- | Assuming that the data pointer points to a 0, changes
-- it to the value of a Char.
char :: Char.Char -> BrainFuck ()
char = multi incr . Char.ord

-- | Undoes `char`
unchar :: Char.Char -> BrainFuck ()
unchar = multi decr . Char.ord

-- | Runs an action with the data pointer pointing at a Char.
withChar :: Char.Char -> BrainFuck a -> BrainFuck a
withChar c a = do
	char c
	r <- a
	unchar c
	return r

-- | Gets the current address of the data pointer.
addr :: BrainFuck DataPointer
addr = BrainFuck $ \loc -> ([], loc, loc)

-- | Moves the data pointer to a specific address.
setAddr :: Integer -> BrainFuck ()
setAddr n = do
	a <- addr
	if a > n
		then prev >> setAddr n
		else if a < n
			then next >> setAddr n
			else return ()

-- | The loop is only run if the data pointer doesn't point to 0.
--
-- On entry, the loop body is run, and then it loops, until the data
-- pointer points to 0.
loopUnless0 :: BrainFuck () -> BrainFuck ()
loopUnless0 a = do
	open
	a
	close

data Constants = Constants NumConstants Integer

type NumConstants = Int

-- | Let's run programs with data addresses 0..n reserved to always contain
-- the numbers n..0
--
-- These constants will make brainfuck quite a bit easier to use!
brainfuckConstants :: NumConstants -> (Constants -> BrainFuck a) -> String
brainfuckConstants numconstants a = brainfuck $ do
	M.forM_ [0..numconstants] $ \n -> do
		add (numconstants - n)
		next
	a (Constants numconstants 0)

-- | Sets data pointer to point to one of the Constants.
pointAt :: Constants -> Integer -> BrainFuck ()
pointAt (Constants numconstants c) i = setAddr (c + fromIntegral numconstants - i)

-- | Run an action in a loop, until the data pointer points to 0.
loop :: Constants -> BrainFuck () -> BrainFuck ()
loop constants a = do
	here <- addr
	pointAt constants 1
	loopUnless0 $ do
		setAddr here
		a

-- | Runs an action in an infinite loop.
forever :: Constants -> BrainFuck () -> BrainFuck ()
forever constants a = loop constants $ do
	a
	pointAt constants 1

-- | Prints out the ASCII characters, starting with space, repeatedly.
-- 
-- TODO: Only print printable characters.
demo :: String
demo = brainfuckConstants 1 $ \constants -> do
	add $ Char.ord ' '
	forever constants $ do
		add 1
		output

cat :: String
cat = brainfuckConstants 1 $ flip forever $ input >> output

helloworld :: String
helloworld = brainfuck $
	M.forM_ "hello, world!" $
		flip withChar output