module CS4400.Result where

import Control.Monad (ap, liftM)
import Data.Map (Map)
import Data.Map qualified as Map

data Expr
  = EVar String
  | EBool Bool
  | ENum Int
  | ELet String Expr Expr
  | EIf Expr Expr Expr
  | EAdd Expr Expr

data Value = VBool Bool | VNum Int
  deriving (Show)

data Result e a = Err e | Ok a
data EvalError = ExpectedNum | ExpectedBool | UnknownVar

instance Monad (Result e) where
  (>>=) = andThen

instance Applicative (Result e) where
  pure = Ok
  (<*>) = ap

instance Functor (Result e) where
  fmap = liftM

andThen :: Result e a -> (a -> Result e b) -> Result e b
andThen r f = case r of
  Err e -> Err e
  Ok x -> f x

toBool :: Value -> Result EvalError Bool
toBool v = case v of
  VBool b -> Ok b
  VNum _ -> Err ExpectedBool

toNum :: Value -> Result EvalError Int
toNum v = case v of
  VBool _ -> Err ExpectedNum
  VNum n -> Ok n

interp :: Map String Value -> Expr -> Result EvalError Value
interp env expr = case expr of
  EVar x -> case Map.lookup x env of
    Nothing -> Err UnknownVar
    Just v -> Ok v
  EBool b -> Ok (VBool b)
  ENum n -> Ok (VNum n)
  EAdd e1 e2 -> do
    v1 <- interp env e1
    n1 <- toNum v1
    v2 <- interp env e2
    n2 <- toNum v2
    Ok (VNum (n1 + n2))
  EIf e1 e2 e3 -> do
    v1 <- interp env e1
    b <- toBool v1
    if b then interp env e2 else interp env e3
  ELet x e1 e2 -> do
    v1 <- interp env e1
    let env' = Map.insert x v1 env
    interp env' e2