module Quox.Typechecker

import public Quox.Syntax
import public Quox.Typing
import public Quox.Equal
import public Control.Monad.Either

%default total


private covering %inline
expectTYPE : MonadError Error m => Term d n -> m Universe
expectTYPE s =
  case (whnfT s).fst of
       TYPE l => pure l
       _      => throwError $ ExpectedTYPE s

private covering %inline
expectPi : MonadError Error m => Term d n ->
           m (Qty, Term d n, ScopeTerm d n)
expectPi ty =
  case (whnfT ty).fst of
       Pi qty _ arg res => pure (qty, arg, res)
       _                => throwError $ ExpectedPi ty

private %inline
expectEqualQ : MonadError Error m =>
               (expect, actual : Qty) -> m ()
expectEqualQ pi rh =
  unless (pi == rh) $ throwError $ ClashQ pi rh


private %inline
popQ : MonadError Error m => Qty -> QOutput (S n) -> m (QOutput n)
popQ pi (qctx :< rh) = expectEqualQ pi rh $> qctx


private %inline
tail : TyContext d (S n) -> TyContext d n
tail = {tctx $= tail, qctx $= tail}


private %inline
weakI : InferResult d n -> InferResult d (S n)
weakI = {type $= weakT, qout $= (:< zero)}

private
lookupBound : {n : Nat} -> Qty -> Var n -> TyContext d n -> InferResult d n
lookupBound pi VZ (MkTyContext {tctx = _ :< ty, _}) =
  InfRes {type = weakT ty, qout = zero :< pi}
lookupBound pi (VS i) ctx =
  weakI $ lookupBound pi i (tail ctx)


private %inline
subjMult : Qty -> Qty -> Subset Qty IsSubj
subjMult sg qty =
  if sg == Zero || qty == Zero
     then Element Zero %search
     else Element One  %search


public export
CanTC : (Type -> Type) -> Type
CanTC m = (MonadError Error m, MonadReader Definitions m)

parameters {auto _ : CanTC m}
 mutual
  -- [todo] it seems like the options here for dealing with substitutions are
  -- to either push them or parametrise the whole typechecker over ambient
  -- substitutions. both of them seem like the same amount of work for the
  -- computer but pushing is less work for the me

  export covering %inline
  check : {d, n : Nat} ->
          (ctx : TyContext d n) -> (sg : Qty) -> (0 _ : IsSubj sg) =>
          (subj : Term d n) -> (ty : Term d n) ->
          m (CheckResult n)
  check ctx sg subj ty = check' ctx sg (pushSubstsT subj) ty

  export covering %inline
  infer : {d, n : Nat} ->
          (ctx : TyContext d n) -> (sg : Qty) -> (0 _ : IsSubj sg) =>
          (subj : Elim d n) ->
          m (InferResult d n)
  infer ctx sg subj = infer' ctx sg (pushSubstsE subj)


  export covering
  check' : {d, n : Nat} ->
           (ctx : TyContext d n) -> (sg : Qty) -> (0 _ : IsSubj sg) =>
           (subj : NotCloTerm d n) -> (ty : Term d n) ->
           m (CheckResult n)

  check' ctx sg (Element (TYPE l) _) ty = do
    l' <- expectTYPE ty
    expectEqualQ zero sg
    unless (l < l') $ throwError $ BadUniverse l l'
    pure zero

  check' ctx sg (Element (Pi qty x arg res) _) ty = do
    l <- expectTYPE ty
    expectEqualQ zero sg
    ignore $ check ctx zero arg (TYPE l)
    case res of
         TUsed   res => ignore $ check (extendTy arg zero ctx) zero res (TYPE l)
         TUnused res => ignore $ check ctx                     zero res (TYPE l)
    pure zero

  check' ctx sg (Element (Lam x body) _) ty = do
    (qty, arg, res) <- expectPi ty
    -- [todo] do this properly?
    let body = fromScopeTerm body; res = fromScopeTerm res
    qout <- check (extendTy arg (sg * qty) ctx) sg body res
    popQ qty qout

  check' ctx sg (Element (E e) _) ty = do
    infres <- infer ctx sg e
    ignore $ check ctx zero ty (TYPE UAny)
    subT infres.type ty
    pure infres.qout

  export covering
  infer' : {d, n : Nat} ->
           (ctx : TyContext d n) -> (sg : Qty) -> (0 _ : IsSubj sg) =>
           (subj : NotCloElim d n) ->
           m (InferResult d n)

  infer' ctx sg (Element (F x) _) = do
    Just g <- asks $ lookup x  | Nothing => throwError $ NotInScope x
    when (isZero g) $ expectEqualQ sg Zero
    pure $ InfRes {type = g.type.def, qout = zero}

  infer' ctx sg (Element (B i) _) =
    pure $ lookupBound sg i ctx

  infer' ctx sg (Element (fun :@ arg) _) = do
    funres <- infer ctx sg fun
    (qty, argty, res) <- expectPi funres.type
    let Element sg' _ = subjMult sg qty
    argout <- check ctx sg' arg argty
    pure $ InfRes {type = fromScopeTerm res //. ((arg :# argty) ::: id),
                   qout = funres.qout + argout}

  infer' ctx sg (Element (tm :# ty) _) = do
    ignore $ check ctx zero ty (TYPE UAny)
    qout <- check ctx sg tm ty
    pure $ InfRes {type = ty, qout}