{-# LANGUAGE LambdaCase #-}

module Check (checkType, findType) where

import Control.Monad.Except (MonadError (throwError))
import Data.List ((!?))
import Errors
import Eval (Env, betaEquiv, envLookupTy, 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

validateType :: Context -> Expr -> ReaderT Env Result Expr
validateType g a = do
    s <- findType g a
    isSort s >>= flip unless (throwError $ NotASort a s)
    pure s

validateType_ :: Context -> Expr -> ReaderT Env Result ()
validateType_ g a = void $ validateType g a

findType :: Context -> Expr -> ReaderT Env Result Expr
findType _ Star = pure Square
findType _ Square = throwError SquareUntyped
findType g (Var x n) = do
    t <- g !? fromInteger n `whenNothing` throwError (UnboundVariable x)
    validateType_ g t
    pure t
findType _ (Free n) = do
    envLookupTy n
findType _ (Axiom n) = envLookupTy n
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
    validateType_ g a
    b <- findType (incIndices a : map incIndices g) m
    validateType_ g (Pi x a b)
    pure $ if occursFree 0 b then Pi x a b else Pi "" a b
findType g (Pi _ a b) = do
    validateType_ g a
    validateType (incIndices a : map incIndices g) b

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