{-# 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.Set (Set)
import Data.Typeable

import qualified Presyntax.Presyntax as P

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)
          , definedNames :: Set Name
          }

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 mempty

assume :: Name -> Value -> (Name -> ElabM a) -> ElabM a
assume nm ty k = defineInternal nm ty VVar k

define :: Name -> Value -> Value -> (Name -> ElabM a) -> ElabM a
define nm vty val = defineInternal nm vty (const val)

assumes :: [Name] -> Value -> ([Name] -> ElabM a) -> ElabM a
assumes nms ty k = do
  let
    go acc [] k = k acc
    go acc (x:xs) k = assume x ty $ \n -> go (n:acc) xs k
   in go [] nms  k

defineInternal :: Name -> Value -> (Name -> Value) -> (Name -> ElabM a) -> ElabM a
defineInternal nm vty val k =
  do
    env <- ask
    let (env', name') = go env
    local (const env') (k name')
  where
    go x =
      let
        nm' = case Map.lookup (getNameText nm) (nameMap x) of
          Just name -> incName nm name
          Nothing -> nm
      in ( x { getEnv     = Map.insert nm' (vty, val nm') (getEnv x)
             , nameMap    = Map.insert (getNameText nm) nm' (nameMap x)
             , whereBound = maybe (whereBound x) (flip (Map.insert nm') (whereBound x)) (currentSpan x)
             }
         , nm')

redefine :: Name -> Value -> Value -> ElabM a -> ElabM a
redefine nm vty val = local go where
  go x = x { getEnv     = Map.insert nm (vty, val) (getEnv x)
           , nameMap    = Map.insert (getNameText nm) nm (nameMap x)
           , whereBound = maybe (whereBound x) (flip (Map.insert nm) (whereBound x)) (currentSpan 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 0)

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

incName :: Name -> Name -> Name
incName (Bound x _) n = Bound x (getNameNum n + 1)
incName (Defined x k) n = Defined x (getNameNum n + 1)