module Quox.Typechecker

import public Quox.Syntax
import public Quox.Typing
import Quox.Error

import Data.Nat
import Data.SortedMap
import Control.Monad.Reader
import Control.Monad.State
import Syntax.WithProof

%default total


public export
data Error : Type where
  ExpectedType : (subj, ty : Term d n) -> Error
  UniverseInconsistency : (lo, hi : Universe) -> Error
  GlobalNotInScope : Name -> Error

public export
record CheckResult
  {0 d, n : Nat}
  (0 ctx : WfContext d n)
  (0 subj, ty : Term d n)
where
  constructor MkCheckResult
  tmout, tyout : QOutput n
  hasType : HasTypeT {ctx = ctx.ctx, subj, ty, tmout, tyout}


public export
record InferResult
  {0 d, n : Nat}
  (0 ctx : WfContext d n)
  (0 subj : Elim d n)
where
  constructor MkInferResult
  tmout, tyout : QOutput n
  ty : Term d n
  hasType : HasTypeE {ctx = ctx.ctx, subj, ty, tmout, tyout}


public export
zeroCtx : Globals -> TyContext 0 0
zeroCtx globals = MkTyContext {globals, dctx = DNil, tctx = [<], qctx = [<]}

public export
zeroWfCtx : Globals -> WfContext 0 0
zeroWfCtx globals = MkWfContext (zeroCtx globals) NIL

public export
Check0 : (globals : Globals) -> (subj, ty : Term 0 0) -> Type
Check0 globals subj ty =
  HasTypeT {ctx = zeroCtx globals, subj, ty, tmout = [<], tyout = [<]}

public export
record Infer0 (globals : Globals) (subj : Elim 0 0) where
  constructor MkInfer0
  ty : Term 0 0
  hasType : HasTypeE {ctx = zeroCtx globals, subj, ty, tmout = [<], tyout = [<]}


export
0 isTypeZeroTyout : HasTypeT {ty = TYPE _, tyout, _} -> IsZero tyout
isTypeZeroTyout (TYPE  {zty, _}) = zty
isTypeZeroTyout (Pi    {zty, _}) = zty
isTypeZeroTyout (PushT {})   = ?isTypeZeroTyout_PushT


parameters {auto _ : MonadThrow Error m} mutual
  private
  checkUniverse : (lo, hi : Universe) -> m (lo `LT` hi)
  checkUniverse lo hi with (lo `isLT` hi)
    _ | Yes lt = pure lt
    _ | No  _  = throw $ UniverseInconsistency {lo, hi}

  private
  expectType : (subj, ty : Term d n) ->
               m (l : Universe ** ty = TYPE l)
  expectType _    (TYPE l) = pure (l ** Refl)
  expectType subj ty       = throw $ ExpectedType {subj, ty}

  private
  checkType : (ctx : WfContext d n) ->
              (l : Universe) -> (ty : Term d n) ->
              m (CheckResult ctx (TYPE l) ty)
  checkType (MkWfContext {ctx, wf}) l ty = do
    (l' ** eq) <- expectType (TYPE l) ty; let Refl = eq
    lt <- checkUniverse l l'
    let Element _ zprf = zero ctx.qctx
    pure $ MkCheckResult {hasType = TYPE {wf, lt, ztm = zprf, zty = zprf}, _}

  private
  checkPi : (ctx : WfContext d n) ->
            (qtm, qty : Qty) ->
            (x : Name) ->
            (arg : Term d n) ->
            (res : Term d (S n)) ->
            (ty  : Term d n) ->
            m (CheckResult ctx (Pi qtm qty x arg res) ty)
  checkPi (MkWfContext {ctx, wf}) qtm qty x arg res ty = do
    let whole = Pi qtm qty x arg res
    (l' ** eq) <- expectType whole ty; let Refl = eq
    ?checkPi_rhs


  private
  inferFree : (ctx : WfContext {}) -> (x : Name) -> m (InferResult ctx (F x))
  inferFree (MkWfContext {ctx, wf}) x =
    case @@ lookup x ctx.globals of
      (Just g ** eq) =>
        let Element _ zprf = zero ctx.tctx
            hasType = F {wf, g, eq, ztm = zprf, zty = zprf}
        in
        pure $ MkInferResult {hasType, _}
      _ => throw $ GlobalNotInScope x

  private
  inferBound : (ctx : WfContext _ n) -> (i : Var n) -> m (InferResult ctx (B i))
  inferBound (MkWfContext {ctx, wf}) i =
    pure $ MkInferResult {hasType = B {wf, look = lookup ctx i, eq = Refl}, _}

  private
  inferAnn : (ctx : WfContext d n) -> (tm, ty : Term d n) ->
             m (InferResult ctx (tm :# ty))
  inferAnn ctx tm ty = do
    ignore $ check ctx ty (TYPE UAny)
    MkCheckResult {hasType, _} <- check ctx tm ty
    pure $ MkInferResult {hasType = Ann hasType, _}


  export
  check : (ctx : WfContext d n) -> (subj, ty : Term d n) ->
          m (CheckResult ctx subj ty)
  check ctx subj ty = case subj of
    TYPE l               => checkType ctx l ty
    Pi qtm qty x arg res => checkPi ctx qtm qty x arg res ty
    Lam x body           => ?check_rhs_3
    E e                  => ?check_rhs_4
    CloT  s th           => ?check_rhs_5
    DCloT s th           => ?check_rhs_6

  export
  infer : (ctx : WfContext d n) -> (subj : Elim d n) ->
          m (InferResult ctx subj)
  infer ctx subj = case subj of
    F x        => inferFree ctx x
    B i        => inferBound ctx i
    fun :@ arg => ?infer_rhs_3
    tm :# ty   => inferAnn ctx tm ty
    CloE  e th => ?infer_rhs_5
    DCloE e th => ?infer_rhs_6


parameters {auto _ : MonadThrow Error m} (globals : Globals)
  public export
  check0 : (subj, ty : Term 0 0) -> m (Check0 globals subj ty)
  check0 subj ty = pure $
    case !(check (zeroWfCtx globals) subj ty) of
      MkCheckResult [<] [<] prf => prf

  public export
  infer0 : (subj : Elim 0 0) -> (ty : Term 0 0) -> m (Infer0 globals subj)
  infer0 subj ty = pure $
    case !(infer (zeroWfCtx globals) subj) of
      MkInferResult [<] [<] ty prf => MkInfer0 ty prf