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

import Control.Monad.Reader
import Control.Exception

import Data.Text.Prettyprint.Doc.Render.Terminal (AnsiStyle)
import qualified Data.Map.Strict as Map
import Data.Text.Prettyprint.Doc
import Data.Map.Strict (Map)
import Data.Text (Text)
import Data.Typeable

import qualified Presyntax.Presyntax as P

import Syntax.Pretty (getNameText)
import Syntax

data ElabEnv =
  ElabEnv { getEnv      :: Map Name (NFType, Value)

          , nameMap     :: Map Text Name
          , pingPong    :: Int
          , commHook    :: Value -> IO ()

          , currentSpan :: Maybe (P.Posn, P.Posn)
          , whereBound  :: Map Name (P.Posn, P.Posn)
          }

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

emptyEnv :: ElabEnv
emptyEnv = ElabEnv mempty mempty 0 (const (pure ())) Nothing mempty

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)
           , whereBound = maybe (whereBound x) (flip (Map.insert nm) (whereBound x)) (currentSpan x)
           }

assumes :: [Name] -> Value -> ElabM a -> ElabM a
assumes nm ty = local go where
  go x = x { getEnv = Map.union (Map.fromList (map (\v -> (v, (ty, VVar v))) nm)) (getEnv x)
           , nameMap = Map.union (Map.fromList (map ((,) <$> getNameText <*> id) nm)) (nameMap x)
           , whereBound = maybe (whereBound x) (\l -> Map.union (Map.fromList (zip nm (repeat l))) (whereBound x)) (currentSpan 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 -> throwElab $ NotInScope nm

getNfType :: Name -> ElabM NFType
getNfType nm = do
  vl <- asks (Map.lookup nm . getEnv)
  case vl of
    Just v -> pure (fst v)
    Nothing -> throwElab $ 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) $ throwElab StackOverflow
    local go k
  where go e = e { pingPong = pingPong e + 1 }

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

data AttachedNote = AttachedNote { getNote :: Doc AnsiStyle, getExc :: SomeException }
  deriving          (Show, Typeable)
  deriving anyclass (Exception)

withNote :: ElabM a -> Doc AnsiStyle -> ElabM a
withNote k note = do
  env <- ask
  liftIO $
      runElab k env
    `catch` \e -> throwIO (AttachedNote note e)

data WhileChecking = WhileChecking { startPos :: P.Posn, endPos :: P.Posn, exc :: SomeException }
  deriving (Show, Typeable, Exception)

withSpan :: P.Posn -> P.Posn -> ElabM a -> ElabM a
withSpan a b k = do
  env <- ask
  liftIO $
      runElab k env{ currentSpan = Just (a, b) }
    `catches` [ Handler $ \e@WhileChecking{} -> throwIO e
              , Handler $ \e -> throwIO (WhileChecking a b e)
              ]

data SeeAlso = SeeAlso { saStartPos :: P.Posn, saEndPos :: P.Posn, saExc :: SomeException }
  deriving (Show, Typeable, Exception)

seeAlso :: ElabM a -> Name -> ElabM a
seeAlso k nm = do
  env <- ask
  case Map.lookup nm (whereBound env) of
    Just l ->
      liftIO $ runElab k env
        `catch` \e -> throwIO (SeeAlso (fst l) (snd l) e)
    Nothing -> k

catchElab :: Exception e => ElabM a -> (e -> ElabM a) -> ElabM a
catchElab k h = do
  env <- ask
  liftIO $ runElab k env `catch` \e -> runElab (h e) env

tryElab :: Exception e => ElabM a -> ElabM (Either e a)
tryElab k = do
  env <- ask
  liftIO $ (Right <$> runElab k env) `catch` \e -> pure (Left e)

throwElab :: Exception e => e -> ElabM a
throwElab = liftIO . throwIO