{-# 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 Syntax
import qualified Presyntax.Presyntax as P

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)
           }

getNameText :: Name -> Text
getNameText (Bound x) = x
getNameText (Defined 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 }

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