module Check (TypeCheckError (..), CheckResult, checkType) where

import Control.Monad (unless)
import Control.Monad.Except (MonadError (throwError))
import Control.Monad.Reader
import Data.List ((!?))
import qualified Data.Map as M
import Data.Text (Text)
import qualified Data.Text as T

import Eval
import Expr

type Context = [Expr]

data TypeCheckError = SquareUntyped | UnboundVariable Text | NotASort Expr Expr | ExpectedPiType Expr Expr | NotEquivalent Expr Expr Expr deriving (Eq, Ord)

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

type CheckResult = Either TypeCheckError

matchPi :: Expr -> Expr -> ReaderT Env CheckResult (Expr, Expr)
matchPi _ (Pi _ a b) = pure (a, b)
matchPi m e = throwError $ ExpectedPiType m e

findType :: Context -> Expr -> ReaderT Env CheckResult Expr
findType _ Star = pure Square
findType _ Square = throwError SquareUntyped
findType g (Var n x) = do
    t <- maybe (throwError $ UnboundVariable x) pure $ g !? fromInteger n
    s <- findType g t
    unless (isSort s) $ throwError $ NotASort t s
    pure t
findType g (Free n) = asks (M.lookup n) >>= maybe (throwError $ UnboundVariable n) (findType g)
findType g e@(App m n) = do
    (a, b) <- findType g m >>= matchPi m
    a' <- findType g n
    equiv <- asks $ runReader (betaEquiv a a')
    unless equiv $ 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

checkType :: Env -> Context -> Expr -> CheckResult Expr
checkType env g t = runReaderT (findType g t) env