module Quox.Typechecker

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

%default total


public export
0 CanTC' : (q : Type) -> (q -> Type) -> (Type -> Type) -> Type
CanTC' q isGlobal m = (HasErr q m, MonadReader (Definitions' q isGlobal) m)

public export
0 CanTC : (q : Type) -> IsQty q => (Type -> Type) -> Type
CanTC q = CanTC' q IsGlobal


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

private covering %inline
expectPi : HasErr q m => Term q d n ->
           m (q, Term q d n, ScopeTerm q d n)
expectPi ty =
  case whnfT ty of
    Element (Pi qty _ arg res) _ => pure (qty, arg, res)
    _ => throwError $ ExpectedPi ty

private covering %inline
expectEq : HasErr q m => Term q d n ->
           m (DScopeTerm q d n, Term q d n, Term q d n)
expectEq ty =
  case whnfT ty of
    Element (Eq _ ty l r) _ => pure (ty, l, r)
    _ => throwError $ ExpectedEq ty

private %inline
expectEqualQ : HasErr q m => Eq q =>
               (expect, actual : q) -> m ()
expectEqualQ pi rh =
  unless (pi == rh) $ throwError $ ClashQ pi rh


private %inline
popQ : HasErr q m => Eq q => q -> QOutput q (S n) -> m (QOutput q n)
popQ pi (qctx :< rh) = expectEqualQ pi rh $> qctx


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


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

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

private
lookupFree : IsQty q => CanTC q m => Name -> m (Definition q)
lookupFree x =
  case lookup x !ask of
    Just d  => pure d
    Nothing => throwError $ NotInScope x

private %inline
subjMult : IsQty q => (sg : SQty q) -> q -> SQty q
subjMult sg qty = if isYes $ isZero qty then szero else sg


export
makeDimEq : DContext d -> DimEq d
makeDimEq DNil           = zeroEq
makeDimEq (DBind dctx)   = makeDimEq dctx :<? Nothing
makeDimEq (DEq p q dctx) = set p q $ makeDimEq dctx


parameters {auto _ : IsQty q} {auto _ : CanTC q 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

  ||| `check ctx sg subj ty` checks that in the context `ctx`, the term
  ||| `subj` has the type `ty`, with quantity `sg`. if so, returns the
  ||| quantities of all bound variables that it used.
  export covering %inline
  check : TyContext q d n -> SQty q -> Term q d n -> Term q d n ->
          m (CheckResult q n)
  check ctx sg subj ty =
    let Element subj nc = pushSubstsT subj in
    check' ctx sg subj nc ty

  ||| `infer ctx sg subj` infers the type of `subj` in the context `ctx`,
  ||| and returns its type and the bound variables it used.
  export covering %inline
  infer : TyContext q d n -> SQty q -> Elim q d n -> m (InferResult q d n)
  infer ctx sg subj =
    let Element subj nc = pushSubstsE subj in
    infer' ctx sg subj nc


  export covering
  check' : TyContext q d n -> SQty q ->
           (subj : Term q d n) -> (0 nc : NotCloT subj) -> Term q d n ->
           m (CheckResult q n)

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

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

  check' ctx sg (Lam x body) _ ty = do
    (qty, arg, res) <- expectPi ty
    qout <- check (extendTy arg (sg.fst * qty) ctx) sg body.term res.term
    popQ qty qout

  check' ctx sg (Eq i t l r) _ ty = do
    u <- expectTYPE ty
    expectEqualQ zero sg.fst
    case t of
      DUsed   t => ignore $ check (extendDim ctx) sg t (TYPE u)
      DUnused t => ignore $ check ctx sg t (TYPE u)
    ignore $ check ctx sg t.zero l
    ignore $ check ctx sg t.one  r
    pure $ zeroFor ctx

  check' ctx sg (DLam i body) _ ty = do
    (ty, l, r) <- expectEq ty
    qout <- check (extendDim ctx) sg body.term ty.term
    let eqs = makeDimEq ctx.dctx
    equalTWith eqs body.zero l
    equalTWith eqs body.one  r
    pure qout

  check' ctx sg (E e) _ ty = do
    infres <- infer ctx sg e
    ignore $ check ctx szero ty (TYPE UAny)
    subTWith (makeDimEq ctx.dctx) infres.type ty
    pure infres.qout

  export covering
  infer' : TyContext q d n -> SQty q ->
           (subj : Elim q d n) -> (0 nc : NotCloE subj) ->
           m (InferResult q d n)

  infer' ctx sg (F x) _ = do
    g <- lookupFree x
    when (isYes $ isZero g) $ expectEqualQ sg.fst zero
    pure $ InfRes {type = g.type.get, qout = zeroFor ctx}

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

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

  infer' ctx sg (fun :% dim) _ = do
    InfRes {type, qout} <- infer ctx sg fun
    (ty, _, _) <- expectEq type
    pure $ InfRes {type = dsub1 ty dim, qout}

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