{-# LANGUAGE FlexibleContexts, LambdaCase #-}

module Task12 where

import Control.Applicative
import Control.Monad.Writer
import Control.Monad.State
import Data.Map ( Map )
import qualified Data.Map as M
import System.IO ( hPutStrLn, stderr )

type Varname = String
type Assignment = Map Varname Integer

data AExpr = Con Integer
           | Var Varname
           | Add AExpr AExpr
           | Mul AExpr AExpr
           deriving ( Eq, Show, Read )

data BExpr = Equal AExpr AExpr
           | LEQ AExpr AExpr
           | And BExpr BExpr
           | Not BExpr
           deriving ( Eq, Show, Read )

data Command = Def Varname AExpr
             | Assign Varname AExpr
             | Print Varname
             | If BExpr Command Command
             | While BExpr Command
             | Seq [Command]
             deriving ( Eq, Show, Read )

val :: (MonadState Assignment m, MonadIO m) => String -> m Integer
val v = gets (M.lookup v) >>= \case
          Nothing -> liftIO (hPutStrLn stderr $ "ERROR: use of undefined variable " ++ v) >> pure 0
          Just x  -> pure x

aeval :: (MonadState Assignment m, MonadIO m) => AExpr -> m Integer
aeval (Con x) = pure x
aeval (Var x) = val x
aeval (Add x y) = liftA2 (+) (aeval x) (aeval y)
aeval (Mul x y) = liftA2 (*) (aeval x) (aeval y)

beval :: (MonadState Assignment m, MonadIO m) => BExpr -> m Bool
beval (Equal x y) = liftA2 (==) (aeval x) (aeval y)
beval (LEQ x y)   = liftA2 (<=) (aeval x) (aeval y)
beval (And x y)   = liftA2 (&&) (beval x) (beval y)
beval (Not x)     = fmap not (beval x)

ceval :: (MonadState Assignment m, MonadWriter [Integer] m, MonadIO m)
      => Command -> m ()
ceval (Def n x) = aeval x >>= \ex -> modify (M.insert n ex)
ceval (Assign n x) = gets (M.member n) >>= \case
                        False -> liftIO (putStrLn $ "ERROR: assignment to an "
                                         ++ "undefined variable " ++ n)
                        True -> ceval (Def n x)
ceval (Print n) = val n >>= tell . (:[])
ceval (If b x y) = beval b >>= \case
                    True -> ceval x
                    False -> ceval y
ceval w@(While b x) = ceval (If b (Seq [x, w]) (Seq []))
ceval (Seq xs) = mapM_ ceval xs

eval :: Command -> IO ()
eval cmd = do
    (_, prints) <- runWriterT $ evalStateT (ceval cmd) M.empty
    putStrLn . unlines $ map show prints