{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE DerivingVia #-}
module Elaboration.Monad where

import Control.Monad.Except
import Control.Monad.Reader
import Control.Applicative

import qualified Data.HashMap.Strict as HashMap
import qualified Data.Sequence as Seq
import Data.HashMap.Strict (HashMap)
import Data.Text (Text)

import Syntax

import Value


data ElabState =
  ElabState
    { elabEnv       :: {-# UNPACK #-} !Env
    , elabLevel     :: {-# UNPACK #-} !Level
    , elabSwitches  :: {-# UNPACK #-} !Int
    , elabNames     :: HashMap Text (Level, VTy)
    , elabConstrs   :: HashMap Text VTy
    , elabBound     :: [BoundDef]  
    , elabSourcePos :: ((Int, Int), (Int, Int))
    }
  deriving (Eq)

emptyElabState :: ElabState
emptyElabState = ElabState emptyEnv (Lvl 0) 0 mempty mempty [] ((0, 0), (0, 0))

getNames :: MonadReader ElabState m => m [Text]
getNames = asks (map go . elabBound) where
  go (BDBound n) = n
  go (BDDefined n) = n

data ElabError
  = NotInScope     Text
  | NotFunction    [Text] Term
  | NotEqual       [Text] Term Term
  | CantSolveMeta  [Text] Term Term
  deriving (Show)

data ProgError
  = ProgError { peErr :: ElabError
              , peSL :: !Int
              , peSC :: !Int
              , peEL :: !Int
              , peEC :: !Int
              }
  deriving (Show)

newtype ElabM a
  = ElabM { runElab :: ElabState -> IO (Either [ProgError] a) }
  deriving
    ( Functor
    , Applicative
    , Monad

    , Alternative
    , MonadPlus

    , MonadReader ElabState
    , MonadError [ProgError]
    , MonadIO
    )
  via ReaderT ElabState (ExceptT [ProgError] IO)

typeError :: ElabError -> ElabM a
typeError err = do
  (s, e) <- asks elabSourcePos
  throwError [uncurry (uncurry (ProgError err) s) e]

assumeLocal :: Text -> VTy -> ElabM a -> ElabM a
assumeLocal name tipe = local go where
  go r =
    r { elabLevel = succ (elabLevel r)
      , elabNames = HashMap.insert name (elabLevel r, tipe) (elabNames r)
      , elabEnv = (elabEnv r) {
          locals = VGlued (HVar (Bound (unLvl (elabLevel r)))) mempty Nothing
            Seq.<| locals (elabEnv r)
        } 
      , elabBound = BDBound name:elabBound r  
      }

defineLocal :: Text -> VTy -> Value -> ElabM a -> ElabM a
defineLocal name tipe val = local go where
  go r =
    r { elabLevel = succ (elabLevel r)
      , elabNames = HashMap.insert name (elabLevel r, tipe) (elabNames r)
      , elabEnv = (elabEnv r) {
          locals = val Seq.<| locals (elabEnv r)
        } 
      , elabBound = BDDefined name:elabBound r  
      }