{-# LANGUAGE DeriveFunctor #-}
module Main where

import Control.Monad ((>=>), ap)
import qualified Data.Map.Strict as Map
import Data.Map.Strict (Map)


data Value m = VInt Int
             | VFun (Value m -> Heap m -> m (Value m, Heap m))

instance Show (Value m) where
  showsPrec d (VInt i) = showParen (d > 10) $ showString "VInt " . showsPrec 11 i
  showsPrec d (VFun _) = showParen (d > 10) $ showString "VFun <function>"

type Name = String

newtype Program = Program [Stmt]
  deriving (Show)

data Stmt
  = Eval Expr
  | Print Expr
  | While Expr [Stmt]
  deriving (Show)

data Expr
  = Var Name
  | LitInt Int
  | Prim Primitive
  | Set Name Expr
  | Lambda Name Expr
  | Apply Expr Expr
  deriving (Show)

data Primitive = ADD | SUB | MUL  -- on integers
               | LESS | EQUAL     -- true = 1, false = 0
  deriving (Show)

example :: Program
example = Program
  [Eval (Set "a" (LitInt 0))
  ,While (Prim LESS `Apply` Var "a" `Apply` LitInt 10)
     [Print (Var "a")
     ,Eval (Set "a" (Prim ADD `Apply` Var "a" `Apply` LitInt 1))]]

type Stack = Map Name Int
type Heap m = Map Int (Value m)

class MonadPrint m where
  mprint :: String -> m ()

class MonadIdGen m where
  genId :: m Int

runP :: (Monad m, MonadPrint m, MonadIdGen m) => Program -> (Stack, Heap m) -> m (Stack, Heap m)
runP (Program ss) = runSs ss

runSs :: (Monad m, MonadPrint m, MonadIdGen m) => [Stmt] -> (Stack, Heap m) -> m (Stack, Heap m)
runSs ss = foldr (>=>) return (map runS ss)

runS :: (Monad m, MonadPrint m, MonadIdGen m) => Stmt -> (Stack, Heap m) -> m (Stack, Heap m)
runS stmt env = case stmt of
  Eval ex -> snd <$> runE ex env
  Print ex -> do (x, env') <- runE ex env
                 case x of
                   VInt i -> mprint (show i)
                   VFun _ -> mprint "<function>"
                 return env'
  While ex body -> do
    (cond, env1) <- runE ex env
    case cond of
      VInt i
        | i >= 1 -> do
            env2 <- runSs body env1
            runS (While ex body) env2
        | otherwise -> do
            return env1
      _ -> error "non-integer in condition"

runE :: (Monad m, MonadPrint m, MonadIdGen m) => Expr -> (Stack, Heap m) -> m (Value m, (Stack, Heap m))
runE expr env = case expr of
  Var n -> case Map.lookup n (fst env) of
             Just addr -> return (snd env Map.! addr, env)
             Nothing -> error "Variable out of scope"
  LitInt i -> return (VInt i, env)
  Prim prim -> do
    let binaryintfun f =
          VFun (\(VInt a) env' ->
                   return (VFun (\(VInt b) env'' ->
                                    return (VInt (f a b), env''))
                          ,env'))
        fun = case prim of  -- TODO type checking
                ADD -> binaryintfun (+)
                SUB -> binaryintfun (-)
                MUL -> binaryintfun (*)
                LESS -> binaryintfun ((fromEnum .) . (<))
                EQUAL -> binaryintfun ((fromEnum .) . (==))
    return (fun, env)
  Set n ex -> do
    (x, env1) <- runE ex env
    addr <- genId
    let env2 = (Map.insert n addr (fst env1)
               ,Map.insert addr x (snd env1))
    return (x, env2)
  Lambda n ex -> return (VFun (\arg heap -> do
                                  addr <- genId
                                  let stack = fst env
                                      stack' = Map.insert n addr stack
                                      heap' = Map.insert addr arg heap
                                  (res, (_, heap'')) <- runE ex (stack', heap')
                                  return (res, heap''))
                        ,env)
  Apply ex1 ex2 -> do
    (f', env1) <- runE ex1 env
    (x, env2) <- runE ex2 env1
    case f' of
      VFun f -> do (res, heap') <- f x (snd env2)
                   return (res, (fst env2, heap'))
      _ -> error "Cannot call non-function"

newtype M a = M { runM :: Int -> IO (a, Int) }
  deriving (Functor)
instance Applicative M where
  pure x = M (\i -> return (x, i))
  (<*>) = ap
instance Monad M where
  M g >>= f = M (\i -> do (x, i') <- g i
                          runM (f x) i')
instance MonadPrint M where
  mprint s = M (\i -> putStrLn s >> return ((), i))
instance MonadIdGen M where
  genId = M (\i -> return (i, i + 1))

main :: IO ()
main = do
  ((stk, heap), outid) <- runM (runP example (mempty, mempty)) 0
  putStrLn $ "out stack: " ++ show stk
  putStrLn $ "out heap: " ++ show heap
  putStrLn $ "out address: " ++ show outid