{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE DeriveAnyClass #-}
module Elab.Monad where

import Control.Monad.Reader
import Control.Exception

import qualified Data.Map.Strict as Map
import Data.Map.Strict (Map)
import Data.Text (Text)
import Data.Typeable

import Syntax
import qualified Data.Text as T

data ElabEnv = ElabEnv { getEnv :: Map Name (NFType, Value), nameMap :: Map Text Name, pingPong :: Int }

newtype ElabM a = ElabM { runElab :: ElabEnv -> IO a }
  deriving (Functor, Applicative, Monad, MonadReader ElabEnv, MonadIO)
    via ReaderT ElabEnv IO

newtype NotInScope = NotInScope { nameNotInScope :: Name }
  deriving          (Show, Typeable)
  deriving anyclass (Exception)

emptyEnv :: ElabEnv
emptyEnv = ElabEnv mempty (Map.singleton (T.pack "Type") (Builtin (T.pack "Type") WiType)) 0

assume :: Name -> Value -> ElabM a -> ElabM a
assume nm ty = local go where
  go x = x { getEnv = Map.insert nm (ty, VVar nm) (getEnv x), nameMap = Map.insert (getNameText nm) nm (nameMap x) }

getNameText :: Name -> Text
getNameText (Bound x) = x
getNameText (Defined x) = x
getNameText (Builtin x _) = x

define :: Name -> Value -> Value -> ElabM a -> ElabM a
define nm ty vl = local go where
  go x = x { getEnv = Map.insert nm (ty, vl) (getEnv x), nameMap = Map.insert (getNameText nm) nm (nameMap x) }

getValue :: Name -> ElabM Value
getValue nm = do
  vl <- asks (Map.lookup nm . getEnv)
  case vl of
    Just v -> pure (snd v)
    Nothing -> liftIO . throwIO $ NotInScope nm

getNfType :: Name -> ElabM NFType
getNfType nm = do
  vl <- asks (Map.lookup nm . getEnv)
  case vl of
    Just v -> pure (fst v)
    Nothing -> liftIO . throwIO $ NotInScope nm

getNameFor :: Text -> ElabM Name
getNameFor x = do
  vl <- asks (Map.lookup x . nameMap)
  case vl of
    Just v -> pure v
    Nothing -> liftIO . throwIO $ NotInScope (Bound x)

switch :: ElabM a -> ElabM a
switch k =
  do
    depth <- asks pingPong
    when (depth >= 128) $ liftIO $ throwIO StackOverflow
    local go k
  where go e = e { pingPong = pingPong e + 1 }