module Check where

import Control.Monad.Except
import Data.List (intercalate, (!?))

import Control.Monad (unless)
import Expr

type Context = [Expr]

data TypeCheckError = Err | SquareUntyped | UnboundVariable | NotASort Expr Int | ExpectedFunctionType Expr | NotEquivalent Expr Expr
    deriving (Show)

type CheckResult = Either TypeCheckError

matchPi :: Expr -> CheckResult (Expr, Expr)
matchPi (Pi _ a b) = Right (a, b)
matchPi e = Left $ ExpectedFunctionType e

showContext :: Context -> String
showContext g = "[" ++ intercalate ", " (map show g) ++ "]"

findType :: Context -> Expr -> CheckResult Expr
findType _ Star = Right Square
findType _ Square = Left SquareUntyped
findType g (Var n _) = do
    t <- maybe (Left UnboundVariable) Right $ g !? fromInteger n
    s <- findType g t
    unless (isSort s) $ throwError $ NotASort s 32
    pure t
findType g (App m n) = do
    (a, b) <- findType g m >>= matchPi
    a' <- findType g n
    unless (betaEquiv a a') $ throwError $ NotEquivalent a a'
    pure $ subst 0 n b
findType g (Abs x a m) = do
    s1 <- findType g a
    unless (isSort s1) $ throwError $ NotASort s1 43
    b <- findType (incIndices a : map incIndices g) m
    s2 <- findType g (Pi x a b)
    unless (isSort s2) $ throwError $ NotASort s2 46
    pure $ if occursFree 0 b then Pi x a b else Pi "" a b
findType g (Pi _ a b) = do
    s1 <- findType g a
    unless (isSort s1) $ throwError $ NotASort s1 51
    s2 <- findType (incIndices a : map incIndices g) b
    unless (isSort s2) $ throwError $ NotASort s2 53
    pure s2