module Quox.Typechecker

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

%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
popQs : HasErr q m => IsQty q =>
        SnocVect s q -> QOutput q (s + n) -> m (QOutput q n)
popQs [<]         qout         = pure qout
popQs (pis :< pi) (qout :< rh) = do expectCompatQ rh pi; popQs pis qout

private %inline
popQ : HasErr q m => IsQty q => q -> QOutput q (S n) -> m (QOutput q n)
popQ pi = popQs [< pi]



private %inline
tail : TyContext q d (S n) -> TyContext q d n
tail = {tctx $= 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 (MkTyContext {tctx = tctx :< ty, _}) =
  InfRes {type = weakT ty, qout = (zero <$ tctx) :< pi}
lookupBound pi (VS i) ctx =
  weakI $ lookupBound pi i (tail ctx)

private %inline
lookupFree : CanTC' q g m => Name -> m (Definition' q g)
lookupFree x = lookupFree' !ask x


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

  ||| "Ψ | Γ ⊢ σ · s ⇐ A ⊳ Σ"
  |||
  ||| `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 = pushSubsts subj in
    check' ctx sg subj nc ty

  ||| "Ψ | Γ ⊢₀ s ⇐ A"
  |||
  ||| `check0 ctx subj ty` checks a term (as `check`) in a zero context.
  export covering %inline
  check0 : TyContext q d n -> Term q d n -> Term q d n -> m ()
  check0 ctx tm ty = ignore $ check ctx szero tm ty
    -- the output will always be 𝟎 because the subject quantity is 0

  ||| "Ψ | Γ ⊢ σ · e ⇒ A ⊳ Σ"
  |||
  ||| `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 = pushSubsts subj in
    infer' ctx sg subj nc


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

  check' ctx sg (TYPE k) _ ty = do
    -- if 𝓀 < ℓ then Ψ | Γ ⊢₀ Type 𝓀 ⇐ Type ℓ
    l <- expectTYPE !ask ty
    expectEqualQ zero sg.fst
    unless (k < l) $ throwError $ BadUniverse k l
    pure $ zeroFor ctx

  check' ctx sg (Pi qty _ arg res) _ ty = do
    l <- expectTYPE !ask ty
    expectEqualQ zero sg.fst
    -- if Ψ | Γ ⊢₀ A ⇐ Type ℓ
    check0 ctx arg (TYPE l)
    -- if Ψ | Γ, x : A ⊢₀ B ⇐ Type ℓ
    case res of
      TUsed   res => check0 (extendTy arg ctx) res (TYPE l)
      TUnused res => check0 ctx res (TYPE l)
    -- then Ψ | Γ ⊢₀ (π·x : A) → B ⇐ Type ℓ
    pure $ zeroFor ctx

  check' ctx sg (Lam _ body) _ ty = do
    (qty, arg, res) <- expectPi !ask ty
    -- if Ψ | Γ, x : A ⊢ σ · t ⇐ B ⊳ Σ, ρ·x
    -- with ρ ≤ σπ
    let qty' = sg.fst * qty
    qout <- check (extendTy arg ctx) sg body.term res.term
    -- then Ψ | Γ ⊢ σ · (λx ⇒ t) ⇐ (π·x : A) → B ⊳ Σ
    popQ qty' qout

  check' ctx sg (Sig _ fst snd) _ ty = do
    l <- expectTYPE !ask ty
    expectEqualQ zero sg.fst
    -- if Ψ | Γ ⊢₀ A ⇐ Type ℓ
    check0 ctx fst (TYPE l)
    -- if Ψ | Γ, x : A ⊢₀ B ⇐ Type ℓ
    case snd of
      TUsed   snd => check0 (extendTy fst ctx) snd (TYPE l)
      TUnused snd => check0 ctx snd (TYPE l)
    -- then Ψ | Γ ⊢₀ (x : A) × B ⇐ Type ℓ
    pure $ zeroFor ctx

  check' ctx sg (Pair fst snd) _ ty = do
    (tfst, tsnd) <- expectSig !ask ty
    -- if Ψ | Γ ⊢ σ · s ⇐ A ⊳ Σ₁
    qfst <- check ctx sg fst tfst
    let tsnd = sub1 tsnd (fst :# tfst)
    -- if Ψ | Γ ⊢ σ · t ⇐ B[s] ⊳ Σ₂
    qsnd <- check ctx sg snd tsnd
    -- then Ψ | Γ ⊢ σ · (s, t) ⇐ (x : A) × B ⊳ Σ₁ + Σ₂
    pure $ qfst + qsnd

  check' ctx sg (Eq i t l r) _ ty = do
    u <- expectTYPE !ask ty
    expectEqualQ zero sg.fst
    -- if Ψ, i | Γ ⊢₀ A ⇐ Type ℓ
    case t of
      DUsed   t => check0 (extendDim ctx) t (TYPE u)
      DUnused t => check0 ctx t (TYPE u)
    -- if Ψ | Γ ⊢₀ l ⇐ A‹0›
    check0 ctx t.zero l
    -- if Ψ | Γ ⊢₀ r ⇐ A‹1›
    check0 ctx t.one  r
    -- then Ψ | Γ ⊢₀ Eq [i ⇒ A] l r ⇐ Type ℓ
    pure $ zeroFor ctx

  check' ctx sg (DLam i body) _ ty = do
    (ty, l, r) <- expectEq !ask ty
    -- if Ψ, i | Γ ⊢ σ · t ⇐ A ⊳ Σ
    qout <- check (extendDim ctx) sg body.term ty.term
    -- if Ψ | Γ ⊢ t‹0› = l : A‹0›
    equal ctx ty.zero body.zero l
    -- if Ψ | Γ ⊢ t‹1› = r : A‹1›
    equal ctx ty.one body.one  r
    -- then Ψ | Γ ⊢ σ · (λᴰi ⇒ t) ⇐ Eq [i ⇒ A] l r ⊳ Σ
    pure qout

  check' ctx sg (E e) _ ty = do
    -- if Ψ | Γ ⊢ σ · e ⇒ A' ⊳ Σ
    infres <- infer ctx sg e
    -- if Ψ | Γ ⊢ A' <: A
    subtype ctx infres.type ty
    -- then Ψ | Γ ⊢ σ · e ⇐ A ⊳ Σ
    pure infres.qout

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

  infer' ctx sg (F x) _ = do
    -- if π·x : A {≔ s} in global context
    g <- lookupFree x
    -- if σ ≤ π
    expectCompatQ sg.fst g.qty
    -- then Ψ | Γ ⊢ σ · x ⇒ A ⊳ 𝟎
    pure $ InfRes {type = g.type.get, qout = zeroFor ctx}

  infer' ctx sg (B i) _ =
    -- if x : A ∈ Γ
    -- then Ψ | Γ ⊢ σ · x ⇒ A ⊳ (𝟎, σ·x, 𝟎)
    pure $ lookupBound sg.fst i ctx

  infer' ctx sg (fun :@ arg) _ = do
    -- if Ψ | Γ ⊢ σ · f ⇒ (π·x : A) → B ⊳ Σ₁
    funres <- infer ctx sg fun
    (qty, argty, res) <- expectPi !ask funres.type
    -- if Ψ | Γ ⊢ σ ⨴ π · s ⇐ A ⊳ Σ₂
    argout <- check ctx (subjMult sg qty) arg argty
    -- then Ψ | Γ ⊢ σ · f s ⇒ B[s] ⊳ Σ₁ + Σ₂
    pure $ InfRes {
      type = sub1 res $ arg :# argty,
      qout = funres.qout + argout
    }

  infer' ctx sg (CasePair pi pair _ ret _ _ body) _ = do
    -- if 1 ≤ π
    expectCompatQ one pi
    -- if Ψ | Γ ⊢ 1 · pair ⇒ (x : A) × B ⊳ Σ₁
    pairres <- infer ctx sone pair
    check0 (extendTy pairres.type ctx) ret.term (TYPE UAny)
    (tfst, tsnd) <- expectSig !ask pairres.type
    -- if Ψ | Γ, x : A, y : B ⊢ σ · body ⇐ ret[(x, y)] ⊳ Σ₂, ρ₁·x, ρ₂·y
    -- with ρ₁, ρ₂ ≤ π
    let bodyctx = extendTyN [< tfst, tsnd.term] ctx
        bodyty  = substCasePairRet pairres.type ret
    bodyout <- check bodyctx sg body.term bodyty >>= popQs [< pi, pi]
    -- then Ψ | Γ ⊢ σ · case ⋯ ⇒ ret[pair] ⊳ πΣ₁ + Σ₂
    pure $ InfRes {
      type = sub1 ret pair,
      qout = pi * pairres.qout + bodyout
    }

  infer' ctx sg (fun :% dim) _ = do
    -- if Ψ | Γ ⊢ σ · f ⇒ Eq [i ⇒ A] l r ⊳ Σ
    InfRes {type, qout} <- infer ctx sg fun
    (ty, _, _) <- expectEq !ask type
    -- then Ψ | Γ ⊢ σ · f p ⇒ A‹p› ⊳ Σ
    pure $ InfRes {type = dsub1 ty dim, qout}

  infer' ctx sg (term :# type) _ = do
    -- if Ψ | Γ ⊢₀ A ⇐ Type ℓ
    check0 ctx type (TYPE UAny)
    -- if Ψ | Γ ⊢ σ · s ⇐ A ⊳ Σ
    qout <- check ctx sg term type
    -- then Ψ | Γ ⊢ σ · (s ∷ A) ⇒ A ⊳ Σ
    pure $ InfRes {type, qout}