-- factorial code: : fact 1 swap begin dup 1 > if dup rot * swap 1 - else drop exit then again drop ;
-- fibonacci code: : fib 1 dup . 1 dup . begin dup rot + dup . again ;
import Data.HashMap.Strict as H

-- Initial types

type ForthState = (IStack, CStack, Dictionary)

type IStack = [Integer]
initialIStack = []

type CStack = [[String]]
initialCStack = []

-- Type for the symbol dictionary

type Dictionary = H.HashMap String [Entry]

data Entry =
     Prim ([Integer] -> [Integer])
   | Def [String]
   | Num Integer
   | Unknown String

instance Show Entry where
  show (Prim f)    = "Prim"
  show (Def s)     = show s
  show (Num i)     = show i
  show (Unknown s) = "Unknown: " ++ s

-- Dictionary helpers

wrap2 f (x:y:xs) = (f y x):xs
wrap2 f _ = error "Value stack underflow!"

dlookup :: String -> Dictionary -> Entry
dlookup word dict =
  case H.lookup word dict of
    Nothing -> case reads word of
                 [(i,"")] -> Num i
                 _        -> Unknown word
    Just x  -> head x

dinsert :: String -> Entry -> Dictionary -> Dictionary
dinsert key val dict =
   case H.lookup key dict of
      Nothing -> H.insert key [val] dict
      Just x  -> H.insert key (val:x) dict

-- Additional Arithmetic operators

dictionary = dinsert "/" (Prim $ wrap2 ( div )) 
	(dinsert "*" (Prim $ wrap2 (*)) 
	(dinsert "-" (Prim $ wrap2 (-)) 
	(dinsert "+" (Prim $ wrap2 (+))
-- Comparisons
	(dinsert "==" (Prim $ Main.compare (==)) 
	(dinsert "/=" (Prim $ Main.compare (/=)) 
	(dinsert ">=" (Prim $ Main.compare (<=)) 
	(dinsert "<=" (Prim $ Main.compare (>=)) 
	(dinsert "<" (Prim $ Main.compare (>)) 
	(dinsert ">" (Prim $ Main.compare (<)) 
-- Dup drop rot swap
	(dinsert "dup" (Prim $ Main.dup ) 
	(dinsert "drop" (Prim $ Main.drop ) 
	(dinsert "rot" (Prim $ Main.rot ) 
	(dinsert "swap" (Prim $ Main.swap ) 
        H.empty
	)))))))))))))





-- The Evaluator

eval :: [String] -> ForthState -> IO ForthState
eval []    (istack, [],     dict) = return (istack, [], dict)
eval words (istack, cstack, dict) =
  case dlookup (head words) dict of
    Def l        -> eval (l ++ xs) (istack, cstack,dict)
    Num i        -> eval xs (i:istack, cstack, dict)
    Prim f       -> eval xs (f istack, cstack, dict)
    Unknown ":" -> eval (chopdef xs) (istack, cstack, (Main.define xs dict ))
    Unknown "if" -> eval (selectpath xs (head istack)) (tail istack, cstack, dict )

    Unknown "begin" -> eval xs (istack, ((getloop xs [] ):cstack), dict )
    Unknown "again" -> eval ((head cstack) ++ xs)  (istack, cstack, dict )
    Unknown "exit" -> eval (striploop xs) (istack, tail cstack, dict )

    Unknown ".S"  -> do { putStrLn $ stackprint istack;
                             eval xs (istack, cstack, dict) }
    Unknown "."  -> do { putStrLn $ show (head istack);
                             eval xs (tail istack, cstack, dict) }
			     
  where xs = tail words

getloop (w:words) accum | w == "again" = reverse (w:accum)
                        | otherwise    = getloop words (w:accum)
striploop (w:words) | w == "again" = words
                    | otherwise    = striploop words

selectpath (words) val | val == -1 = truepath words
                       | val ==  0 = falsepath words
		       | otherwise = error "incorrect value in conditional"

truepath words = (truepart words []) ++ (stripcond words)
truepart (w:words) accum | w == "then" = reverse $ tail accum
                         | w == "else" = reverse accum
			 | otherwise = truepart words (w:accum)

falsepath words = (falsepart words [] ) ++ (stripcond words)
falsepart (w:words) accum | w == "then" = (head accum):[]
                          | w ==  "else" = falsepart words []
			  | otherwise = falsepart words (w:accum)

stripcond (w:words) | w == "then" = words
                    | otherwise = stripcond words

compare f (x:y:xs) | f x y == True  = -1:xs
		   | otherwise = 0:xs

chopdef [] = []
chopdef (x:xs) | x == ";" = xs
               | otherwise = chopdef xs

define (w:words) dict = dinsert w ( Def $ getdef words [] ) dict

getdef :: [String] -> [String] -> [String]
getdef [] accum = error "malformed definition"
getdef (w:words) accum | w == ";" = reverse (accum)
                       | otherwise = getdef words (w:accum)

stackprint [] = ""
stackprint (x:xs) = (stackprint xs) ++ (show x) ++ " " 

swap (x:y:xs) = y:x:xs
swap _ = error "Not enough on stack to swap"

rot (x:y:z:xs) = z:x:y:xs
rot _ = error "not enough on stack to rotate"

dup (x:xs) =  (x:x:xs)
dup _ = error "empty stack"

drop (x:xs) = xs
drop _ = error "empty stack"

repl :: ForthState -> IO ForthState
repl state =
  do putStr "> " ;
     input <- getLine
     nustate <- eval (words input) state
     repl nustate

main = do
  putStrLn "Welcome to your forth interpreter!"
  repl (initialIStack, initialCStack, dictionary)