module Main where

import Control.Arrow
import Control.Monad
import Data.List
import Data.Maybe
import System
import Text.Printf

import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as B 

main = do
	args <- getArgs
	case length args of
		0 -> run =<< B.getContents
		1 -> run =<< B.readFile (head args)
		_ -> putStr =<< readFile "usage.txt"

run contents = do
	let (prog,input) = parse contents
	putStrLn (printf "Program parsed: %d rules, %d bytes of input" (length prog) (B.length input))
	print prog
	(finalState,n,_) <- untilM (not . (\(_,_,x) -> x)) (step prog) (input,0::Int,True)
	putStrLn ""
	putStrLn (printf "Program completed after %d rewrites" n)
	putStrLn ("Residual state: \""++B.unpack finalState++"\"")

type Rule = (ByteString, ByteString)

thueSep = B.pack "::="

--parse :: (IsString a) => a -> (Rule a, a)
parse prog = (parsedRules, B.unlines input)
	where
		(rules,_:input) = break (== thueSep) (filter (\xs -> not (B.null xs) && B.head xs /= '#') (B.lines prog))
		parsedRules = mapMaybe (splitWith thueSep) rules

fromMaybeM a (Just x) = return x
fromMaybeM a Nothing = return a

fmapFmapSecond f m = do
	(x,y) <- m
	y' <- f y
	return (x,y')

splitWith sep x = do
	i <- B.findSubstring sep x
	return $ second (B.drop $ B.length sep) $ B.splitAt i x

step rules (state,n,_) = foldM stepRule (state,n,False) rules
	where
		stepRule :: (ByteString,Int,Bool) -> Rule -> IO (ByteString,Int,Bool)
		stepRule (state,n,True)  _  = return (state,n,True)
		stepRule (state,n,x) (str,replace) =
			maybe
				(return (state,n,x))
				(doReplace replace n)
				(splitWith str state)
		doReplace :: ByteString -> Int -> (ByteString,ByteString) -> IO (ByteString,Int,Bool)
		doReplace replace n (pre,suf)
			| not (B.null replace) && B.head replace == '~' =
				B.putStr (B.tail replace) >> doReplace B.empty n (pre,suf)
			| otherwise = do
				let newState = B.concat [pre,replace,suf]
				(n `mod` 10000 == 0) `when` putStr (printf "After %d rewrites: %s\n" n (B.unpack newState))
				return (newState, n+1, True)

untilM :: (Monad m) => (a -> Bool) -> (a -> m a) -> a -> m a
untilM cond f x =
	if (cond x)
		then return x
		else (f x >>= untilM cond f)