{-# LANGUAGE LambdaCase #-}

module Check (checkType, findType) where

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

findLevel :: Context -> Expr -> ReaderT Env Result Integer
findLevel g a = do
    s <- findType g a
    whnf s >>= \case
        Level i -> pure i
        _ -> throwError $ NotASort a s

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

findType :: Context -> Expr -> ReaderT Env Result Expr
findType _ (Level i) = pure $ Level (i + 1)
findType g (Var x n) = do
    t <- g !? fromInteger n `whenNothing` throwError (UnboundVariable x)
    validateType g t
    pure t
findType _ (Free 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
    i <- findLevel g a
    j <- findLevel (incIndices a : map incIndices g) b
    pure $ Level $ max (i - 1) j -- This feels very sketchy, but certainly adds impredicativity
findType g (Let _ Nothing v b) = findType g (subst 0 v b)
findType g e@(Let _ (Just t) v b) = do
    res <- findType g (subst 0 v b)
    equiv <- betaEquiv t res
    unless equiv $ throwError $ NotEquivalent t res e
    pure t

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