{-# LANGUAGE LambdaCase #-}

module Check (checkType, findType) where

import Control.Monad.Except (MonadError (throwError), liftEither)
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
        Star -> pure $ -1 -- HACK: but works, so...
        _ -> throwError $ NotASort a

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

isSort :: Expr -> Bool
isSort Star = True
isSort (Level _) = True
isSort _ = False

compSort :: Expr -> Expr -> Expr -> Expr -> Result Expr
compSort _ _ Star Star = pure Star
compSort _ _ Star r@(Level _) = pure r
compSort _ _ (Level i) Star
    | i == 0 = pure Star
    | otherwise = pure $ Level i
compSort _ _ (Level i) (Level j) = pure $ Level $ max i j
compSort a b left _
    | isSort left = throwError $ NotASort b
    | otherwise = throwError $ NotASort a

findType :: Context -> Expr -> ReaderT Env Result Expr
findType _ Star = pure $ Level 0
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 _ (Axiom n) = envLookupTy n
findType g e@(App m n) = do
    (a, b) <- findType g m >>= matchPi m
    a' <- findType g n
    betaEquiv' e a a'
    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
    aSort <- findType g a
    bSort <- findType (incIndices a : map incIndices g) b
    liftEither $ compSort a b aSort bSort
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)
    _ <- findType g t
    betaEquiv' e t res
    pure t

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