{-# LANGUAGE BangPatterns #-}
module Check where

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

import Control.Monad (unless)
import Expr
import Debug.Trace

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) ++ "]"

-- TODO: Debug these problem cases
-- λ (S : *) (P : S -> *) (H : forall (x : S), P x) (y : S) => P y
findType :: Context -> Expr -> CheckResult Expr
findType _ Star = trace "star" $ Right Square
findType _ Square = trace "square" $ Left SquareUntyped
findType g (Var n _) = do
    !_ <- trace ("var:\t" ++ showContext g ++ "\n  n:\t" ++ show n) (Right Star)
    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
    !_ <- trace ("app:\t" ++ showContext g ++ "\n  m:\t" ++ show m ++ "\n  n: \t" ++ show n) (Right Star)
    (a, b) <- findType g m >>= matchPi
    a' <- findType g n
    unless (betaEquiv a a') $ throwError $ NotEquivalent a a'
    pure $ subst n b
findType g (Abs x a m) = do
    !_ <- trace ("abs:\t" ++ showContext g ++ "\n  a:\t" ++ show a ++ "\n  m:\t" ++ show m) (Right Star)
    s1 <- findType g a
    !_ <- trace ("back in abs:\t" ++ showContext g ++ "\n\t\t" ++ show a ++ " : " ++ show s1) (Right Star)
    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
    !_ <- trace ("pi:\t" ++ showContext g ++ "\n  a:\t" ++ show a ++ "\n  b:\t" ++ show b) (Right Star)
    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