{-# LANGUAGE LambdaCase #-}

module Check (checkType) where

import Control.Monad (unless)
import Control.Monad.Except (MonadError (throwError))
import Control.Monad.Reader
import Data.List ((!?))
import Errors
import Eval (Env, betaEquiv, envLookup, isSort, subst, whnf)
import Expr (Expr (..), incIndices, occursFree)

type Context = [Expr]

matchPi :: Expr -> Expr -> ReaderT Env Result (Expr, Expr)
matchPi x mt =
    whnf mt >>= \case
        (Pi _ a b) -> pure (a, b)
        t -> throwError $ ExpectedPiType x t

findType :: Context -> Expr -> ReaderT Env Result Expr
findType _ Star = pure Square
findType _ Square = throwError SquareUntyped
findType g (Var n x) = do
    t <- maybe (throwError $ UnboundVariable x) pure $ g !? fromInteger n
    (sSort, s) <- findType g t >>= isSort
    unless sSort $ throwError $ NotASort t s
    pure t
findType g (Free n) = envLookup n >>= findType g
findType g e@(App m n) = do
    (a, b) <- findType g m >>= matchPi m
    a' <- findType g n
    equiv <- betaEquiv a a'
    unless equiv $ throwError $ NotEquivalent a a' e
    pure $ subst 0 n b
findType g (Abs x a m) = do
    (s1Sort, s1) <- findType g a >>= isSort
    unless s1Sort $ throwError $ NotASort a s1
    b <- findType (incIndices a : map incIndices g) m
    (s2Sort, s2) <- findType g (Pi x a b) >>= isSort
    unless s2Sort $ throwError $ NotASort (Pi x a b) s2
    pure $ if occursFree 0 b then Pi x a b else Pi "" a b
findType g (Pi _ a b) = do
    (s1Sort, s1) <- findType g a >>= isSort
    unless s1Sort $ throwError $ NotASort a s1
    (s2Sort, s2) <- findType (incIndices a : map incIndices g) b >>= isSort
    unless s2Sort $ throwError $ NotASort b s2
    pure s2

checkType :: Env -> Expr -> Result Expr
checkType env t = runReaderT (findType [] t) env