{-# 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
findType g (Let _ v b) = do
    a <- findType g v
    validateType_ g a
    res <- findType (incIndices a : map incIndices g) b
    pure $ subst 0 a res

-- this is kinda goofy, it's just like a function, except the resulting type
-- of the body doesn't need to result in a valid function type
-- this means things like `let x := * in ...` would be allowed, even though
-- you couldn't write a function that takes something of type `□` as an argument

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