module Check where

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

import Control.Monad (unless)
import Expr

type Context = [Expr]

data TypeCheckError = SquareUntyped | UnboundVariable String | NotASort Expr Expr | ExpectedPiType Expr Expr | NotEquivalent Expr Expr Expr

instance Show TypeCheckError where
    show SquareUntyped = "□ does not have a type"
    show (UnboundVariable x) = "Unbound variable: " ++ x
    show (NotASort x t) = "Expected " ++ pretty x ++ " to have type * or □, instead found " ++ pretty t
    show (ExpectedPiType m a) = pretty m ++ " : " ++ pretty a ++ " is not a function"
    show (NotEquivalent a a' e) = "Cannot unify " ++ pretty a ++ " with " ++ pretty a' ++ " when evaluating " ++ pretty e

type CheckResult = Either TypeCheckError

matchPi :: Expr -> Expr -> CheckResult (Expr, Expr)
matchPi _ (Pi _ a b) = Right (a, b)
matchPi m e = Left $ ExpectedPiType m 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 x) = do
    t <- maybe (Left $ UnboundVariable x) Right $ g !? fromInteger n
    s <- findType g t
    unless (isSort s) $ throwError $ NotASort t s
    pure t
findType g e@(App m n) = do
    (a, b) <- findType g m >>= matchPi m
    a' <- findType g n
    unless (betaEquiv a a') $ throwError $ NotEquivalent a a' e
    pure $ subst 0 n b
findType g (Abs x a m) = do
    s1 <- findType g a
    unless (isSort s1) $ throwError $ NotASort a s1
    b <- findType (incIndices a : map incIndices g) m
    s2 <- findType g (Pi x a b)
    unless (isSort s2) $ throwError $ NotASort (Pi x a b) s2
    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 a s1
    s2 <- findType (incIndices a : map incIndices g) b
    unless (isSort s2) $ throwError $ NotASort b s2
    pure s2