module Core.Check

import Control.Monad.RWS
import Control.Monad.Either
import Control.Monad.Identity
import Data.Vect
import Data.Fin
import Data.IOArray
import Data.IORef

import Core.Term
import Core.Value
import Core.Normalize
import Core.Misc
import Core.Convert

%default total
 
-- extend environment, used to ensure environment is always in normal form
extV : {auto deftrs : RefA DTR Value}
    -> {auto frst   : Ref NST Nat}
    -> Ctx n -> Value -> PI (Ctx (S n))
extV ctx val = whnf val >>= pure . (`Data.Vect.(::)` ctx)

-- to extend, closure env, term
extT : {auto deftrs : RefA DTR Value}
    -> {auto frst   : Ref NST Nat}
    -> Ctx m -> Ctx n -> Term n -> PI (Ctx (S m))
extT ctx env = extV ctx . VClos env

mutual
  public export
  -- terms types expected term
  check : {auto deftrs    : RefA DTR Value}
       -> {auto deftys : RefA DTY Value}
       -> {auto frst   : Ref NST Nat}
       -> Ctx n -> Ctx n -> Value -> Term n -> PI Bool
  check trs tys xpt' tr = do
    xpt <- whnf xpt'
    case tr of
      TLam sc    => case xpt of
                         VClos env (TPi a b) => do
                           v <- VGen <$> fresh
                           check (v :: trs) !(extT tys env a) (VClos (v :: env) b) sc
                         _ => oops "expected pi"

      TPair  x y => case xpt of
                         (VClos env (TSigma a b)) => do
                           guardS "Pair a" =<< check trs tys (VClos env a) x
                           check trs tys (VClos env b `VApp` VClos trs x) y
                         _ => oops "expected sigma"

      TLet ty tr tri => do
        guardS "let" =<< check trs tys (VClos trs ty) tr
        check !(extT trs trs tr) !(extT tys trs ty) xpt tri

      _       => convert xpt =<< infer trs tys tr

  -- terms types term
  public export
  infer : {auto deftrs : RefA DTR Value}
       -> {auto deftys : RefA DTY Value}
       -> {auto frst   : Ref NST Nat}
       -> Ctx n -> Ctx n -> Term n -> PI Value
  infer trs tys TType      = pure VType
  infer trs tys TTop       = pure VType
  infer trs tys TBot       = pure VType
  infer trs tys TNat       = pure VType
  infer trs tys TStar      = pure VTop
  infer trs tys TZero      = pure VNat
  infer trs tys (TVar i)   = pure (index i tys)
  infer trs tys (TDef i)   = do
    res <- getArr DTY i
    case res of
         Just x  => pure x
         Nothing => oops "TDef type lookup"
  infer trs tys (TApp f x) = infer trs tys f >>= whnf >>=
    \case
       VClos env (TPi a b) => do
         guardS ("app x:\n" ++ show !(whnf (VClos env a))) =<< check trs tys (VClos env a) x
         tr <- whnf (VClos trs x)
         pure (VClos (tr :: env) b)

       _ => oops "expected infer pi"

  infer trs tys (TPi a b) = do
    v <- VGen <$> fresh
    guardS "Pi a" =<< check trs tys VType a 
    guardS "Pi b" =<< check (v :: trs) !(extT tys trs a) VType b
    pure VType

  infer trs tys (TSigma a b) = do
    guardS "Σ a" =<< check trs tys VType a
    guardS "Σ b" =<< check trs tys (VClos trs (TPi a TType)) b
    pure VType

  infer trs tys (TId ty a b) = do
    guardS "Id ty" =<< check trs tys VType ty
    guardS "Id a"  =<< check trs tys (VClos trs ty) a
    guardS "Id b"  =<< check trs tys (VClos trs ty) b
    pure VType

  infer trs tys (TSuc n)  = do 
    guardS "suc n" =<< check trs tys VNat n
    pure VNat

  infer trs tys (TRefl ty tr) = do
    guardS "Refl ty" =<< check trs tys VType ty
    guardS "Refl tr" =<< check trs tys (VClos trs ty) tr
    pure (VClos trs (TId ty tr tr))

  infer trs tys (TNatInd c z s) = do
    guardS "ℕ C" =<< check trs tys (VClos trs (TPi TNat TType)) c
    guardS "ℕ z" =<< check trs tys (VApp (VClos trs c) (VClos [] TZero)) z
    guardS "ℕ s" =<< check trs tys (VClos trs (TPi TNat 
                                              (TPi (TApp (weakTr c) (TVar 0)) 
                                                   (TApp (weakTr2 c) (TSuc (TVar (FS FZ))))))) s
    pure (VClos trs (TPi TNat (TApp (weakTr c) (TVar 0))))

  infer trs tys (TJ ty a b c d) = do
    guardS "J ty" =<< check trs tys VType ty
    guardS "J a"  =<< check trs tys (VClos trs ty) a
    guardS "J b"  =<< check trs tys (VClos trs ty) b
    guardS "J c"  =<< check trs tys
                            (VClos trs (TPi ty {- a : A -}
                                            (TPi (weakTr ty) {- b : A -}
                                                 (TPi (TId (weakTr2 ty)
                                                           (TVar (FS FZ))
                                                           (TVar FZ)) {- Id A a b -}
                                                      TType)))) c
    guardS "J d"  =<< check trs tys (VClos trs (c `TApp` a `TApp` a `TApp` TRefl ty a)) d
    pure (VClos trs (TPi (TId ty a b) (weakTr c `TApp` weakTr a `TApp` weakTr b `TApp` TVar 0)))
                                                                 

  infer trs tys (TSigInd a b c f) = do
    guardS "Σ A" =<< check trs tys VType a
    guardS "Σ B" =<< check trs tys (VClos trs (TPi a TType)) b
    guardS "Σ C" =<< check trs tys (VClos trs (TPi (TSigma a b) TType)) c
    guardS "Σ f" =<< check trs tys (VClos trs (TPi a {-a:A-}
                                              (TPi (weakTr b `TApp` TVar 0) {-b:Ba-}
                                                   (weakTr2 c `TApp` TPair (TVar (FS FZ)) (TVar 0))))) f
    pure (VClos trs (TPi (TSigma a b) (weakTr c `TApp` TVar 0)))
                                      

  infer trs tys (TLet ty tr tri) = do
        guardS "let infer" =<< check trs tys (VClos trs ty) tr
        infer !(extT trs trs tr) !(extT tys trs ty) tri

  infer trs tys (TTopInd c st) = do
    guardS "⊤ C" =<< check trs tys (VClos trs (TPi TTop TType)) c
    guardS "⊤ ★" =<< check trs tys (VClos trs (TApp c TStar))   st
    pure (VClos trs (TPi TTop (TApp (weakTr c) (TVar 0))))

  infer trs tys (TBotInd c) = do
    guardS "⊥ C" =<< check trs tys (VClos trs (TPi TBot TType)) c
    pure (VClos trs (TPi TBot (TApp (weakTr c) (TVar 0))))

  infer trs tys x          = oops ("cannot infer type " ++ show x)

public export
typecheck : {auto deftrs : RefA DTR Value}
         -> {auto deftys : RefA DTY Value}
         -> {auto frst   : Ref NST Nat}
         -> Term 0 -> Term 0 -> IO (Either String Bool)
typecheck tr ty = resolve $ (&&) <$> check [] [] VType ty 
                                 <*> delay <$> check [] [] (VClos [] ty) tr