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

import qualified Presyntax.Presyntax as P

import Syntax

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

          , nameMap     :: Map Text Name
          , pingPong    :: {-# UNPACK #-} !Int
          , commHook    :: Doc AnsiStyle -> IO ()

          , currentSpan :: Maybe (P.Posn, P.Posn)
          , currentFile :: Maybe Text

          , whereBound   :: Map Name (P.Posn, P.Posn)
          , definedNames :: Set Name

          , boundaries   :: Map Name Boundary

          , unsolvedMetas :: {-# UNPACK #-} !(IORef (Map MV [(Seq Projection, Value)]))

          , loadedFiles :: [String]
          }

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

emptyEnv :: IO ElabEnv
emptyEnv = do
  u <- newIORef mempty
  pure $ ElabEnv { getEnv      = mempty
                 , nameMap     = mempty
                 , pingPong    = 0
                 , commHook    = const (pure ())

                 , currentSpan = Nothing
                 , currentFile = Nothing

                 , whereBound   = mempty
                 , definedNames = mempty

                 , boundaries   = mempty

                 , unsolvedMetas = u
                 , loadedFiles = []
                 }

addBoundary :: Name -> Boundary -> ElabM a -> ElabM a
addBoundary nm boundary = local (\e -> e { boundaries = Map.insert nm boundary (boundaries e)} )

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)

makeLetDef :: Name -> Value -> Value -> (Name -> ElabM a) -> ElabM a
makeLetDef nm vty val = defineInternal nm vty (\nm -> GluedVl (HVar nm) mempty val)

replaceLetDef :: Name -> Value -> Value -> ElabM a -> ElabM a
replaceLetDef nm vty val = redefine nm vty (GluedVl (HVar nm) mempty 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 _) n = Defined x (getNameNum n + 1)
incName (ConName x _ s a) n = ConName x (getNameNum n + 1) s a