module Quox.Typechecker

import public Quox.Typing
import public Quox.Equal

import Data.List
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 =>
        QOutput q s -> 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]


parameters {auto _ : IsQty q} {auto _ : CanTC q m}
 mutual
  ||| "Ψ | Γ ⊢ σ · 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.
  |||
  ||| if the dimension context is inconsistent, then return `Nothing`, without
  ||| doing any further work.
  export covering %inline
  check : (ctx : TyContext q d n) -> SQty q -> Term q d n -> Term q d n ->
          m (CheckResult ctx.dctx q n)
  check ctx sg subj ty = ifConsistent ctx.dctx $ checkC ctx sg subj 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

  ||| `check`, assuming the dimension context is consistent
  export covering %inline
  checkC : (ctx : TyContext q d n) -> SQty q -> Term q d n -> Term q d n ->
           m (CheckResult' q n)
  checkC ctx sg subj ty =
    wrapErr (WhileChecking ctx sg.fst subj ty) $
      let Element subj nc = pushSubsts subj in
      check' ctx sg subj ty

  ||| "Ψ | Γ ⊢₀ s ⇐ ★ᵢ"
  |||
  ||| `checkType ctx subj ty` checks a type (in a zero context). sometimes the
  ||| universe doesn't matter, only that a term is _a_ type, so it is optional.
  export covering %inline
  checkType : TyContext q d n -> Term q d n -> Maybe Universe -> m ()
  checkType ctx subj l = ignore $ ifConsistent ctx.dctx $ checkTypeC ctx subj l

  export covering %inline
  checkTypeC : TyContext q d n -> Term q d n -> Maybe Universe -> m ()
  checkTypeC ctx subj l =
    wrapErr (WhileCheckingTy ctx subj l) $ checkTypeNoWrap ctx subj l

  export covering %inline
  checkTypeNoWrap : TyContext q d n -> Term q d n -> Maybe Universe -> m ()
  checkTypeNoWrap ctx subj l =
    let Element subj nc = pushSubsts subj in
    checkType' ctx subj l

  ||| "Ψ | Γ ⊢ σ · 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.
  |||
  ||| if the dimension context is inconsistent, then return `Nothing`, without
  ||| doing any further work.
  export covering %inline
  infer : (ctx : TyContext q d n) -> SQty q -> Elim q d n ->
          m (InferResult ctx.dctx q d n)
  infer ctx sg subj = ifConsistent ctx.dctx $ inferC ctx sg subj

  ||| `infer`, assuming the dimension context is consistent
  export covering %inline
  inferC : (ctx : TyContext q d n) -> SQty q -> Elim q d n ->
           m (InferResult' q d n)
  inferC ctx sg subj =
    wrapErr (WhileInferring ctx sg.fst subj) $
      let Element subj nc = pushSubsts subj in
      infer' ctx sg subj


  private covering
  toCheckType : TyContext q d n -> SQty q ->
                (subj : Term q d n) -> (0 nc : NotClo subj) => Term q d n ->
                m (CheckResult' q n)
  toCheckType ctx sg t ty = do
    u <- expectTYPE !ask ctx ty
    expectEqualQ zero sg.fst
    checkTypeNoWrap ctx t (Just u)
    pure $ zeroFor ctx

  private 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 t@(TYPE _) ty = toCheckType ctx sg t ty

  check' ctx sg t@(Pi {}) ty = toCheckType ctx sg t ty

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

  check' ctx sg t@(Sig {}) ty = toCheckType ctx sg t ty

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

  check' ctx sg t@(Enum _) ty = toCheckType ctx sg t ty

  check' ctx sg (Tag t) ty = do
    tags <- expectEnum !ask ctx ty
    -- if t ∈ ts
    unless (t `elem` tags) $ throwError $ TagNotIn t tags
    -- then Ψ | Γ ⊢ σ · t ⇐ {ts} ⊳ 𝟎
    pure $ zeroFor ctx

  check' ctx sg t@(Eq {}) ty = toCheckType ctx sg t ty

  check' ctx sg (DLam body) ty = do
    (ty, l, r) <- expectEq !ask ctx ty
    -- if Ψ, i | Γ ⊢ σ · t ⇐ A ⊳ Σ
    qout <- checkC (extendDim body.name 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 Nat ty = toCheckType ctx sg Nat ty

  check' ctx sg Zero ty = do
    expectNat !ask ctx ty
    pure $ zeroFor ctx

  check' ctx sg (Succ n) ty = do
    expectNat !ask ctx ty
    checkC ctx sg n Nat

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

  private covering
  checkType' : TyContext q d n ->
               (subj : Term q d n) -> (0 nc : NotClo subj) =>
               Maybe Universe -> m ()

  checkType' ctx (TYPE k) u = do
    -- if 𝓀 < ℓ then Ψ | Γ ⊢₀ Type 𝓀 ⇐ Type ℓ
    case u of
      Just l  => unless (k < l) $ throwError $ BadUniverse k l
      Nothing => pure ()

  checkType' ctx (Pi qty arg res) u = do
    -- if Ψ | Γ ⊢₀ A ⇐ Type ℓ
    checkTypeC ctx arg u
    -- if Ψ | Γ, x : A ⊢₀ B ⇐ Type ℓ
    case res.body of
      Y res' => checkTypeC (extendTy zero res.name arg ctx) res' u
      N res' => checkTypeC ctx                              res' u
    -- then Ψ | Γ ⊢₀ (π·x : A) → B ⇐ Type ℓ

  checkType' ctx t@(Lam {}) u =
    throwError $ NotType ctx t

  checkType' ctx (Sig fst snd) u = do
    -- if Ψ | Γ ⊢₀ A ⇐ Type ℓ
    checkTypeC ctx fst u
    -- if Ψ | Γ, x : A ⊢₀ B ⇐ Type ℓ
    case snd.body of
      Y snd' => checkTypeC (extendTy zero snd.name fst ctx) snd' u
      N snd' => checkTypeC ctx                              snd' u
    -- then Ψ | Γ ⊢₀ (x : A) × B ⇐ Type ℓ

  checkType' ctx t@(Pair {}) u =
    throwError $ NotType ctx t

  checkType' ctx (Enum _) u = pure ()
    -- Ψ | Γ ⊢₀ {ts} ⇐ Type ℓ

  checkType' ctx t@(Tag {}) u =
    throwError $ NotType ctx t

  checkType' ctx (Eq t l r) u = do
    -- if Ψ, i | Γ ⊢₀ A ⇐ Type ℓ
    case t.body of
      Y t' => checkTypeC (extendDim t.name ctx) t' u
      N t' => checkTypeC ctx                    t' 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 ℓ

  checkType' ctx t@(DLam {}) u =
    throwError $ NotType ctx t

  checkType' ctx Nat u = pure ()
  checkType' ctx Zero u = throwError $ NotType ctx Zero
  checkType' ctx t@(Succ _) u = throwError $ NotType ctx t

  checkType' ctx (E e) u = do
    -- if Ψ | Γ ⊢ σ · e ⇒ A' ⊳ Σ
    infres <- inferC ctx szero e
    -- if Ψ | Γ ⊢ A' <: A
    case u of
      Just u  => subtype ctx infres.type (TYPE u)
      Nothing => ignore $ expectTYPE !ask ctx infres.type
    -- then Ψ | Γ ⊢ σ · e ⇐ A ⊳ Σ


  private 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 = injectT ctx g.type, qout = zeroFor ctx}
   where
    lookupFree : Name -> m (Definition q)
    lookupFree x = lookupFree' !ask x

  infer' ctx sg (B i) =
    -- if x : A ∈ Γ
    -- then Ψ | Γ ⊢ σ · x ⇒ A ⊳ (𝟎, σ·x, 𝟎)
    pure $ lookupBound sg.fst i ctx.tctx
   where
    lookupBound : q -> Var n -> TContext q d n -> InferResult' q d n
    lookupBound pi VZ (ctx :< ty) =
      InfRes {type = weakT ty, qout = zeroFor ctx :< pi}
    lookupBound pi (VS i) (ctx :< _) =
      let InfRes {type, qout} = lookupBound pi i ctx in
      InfRes {type = weakT type, qout = qout :< zero}

  infer' ctx sg (fun :@ arg) = do
    -- if Ψ | Γ ⊢ σ · f ⇒ (π·x : A) → B ⊳ Σ₁
    funres <- inferC ctx sg fun
    (qty, argty, res) <- expectPi !ask ctx funres.type
    -- if Ψ | Γ ⊢ σ ⨴ π · s ⇐ A ⊳ Σ₂
    argout <- checkC 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 Ψ | Γ ⊢ σ · pair ⇒ (x : A) × B ⊳ Σ₁
    pairres <- inferC ctx sg pair
    -- if Ψ | Γ, p : (x : A) × B ⊢₀ ret ⇐ Type
    checkTypeC (extendTy zero ret.name pairres.type ctx) ret.term Nothing
    (tfst, tsnd) <- expectSig !ask ctx pairres.type
    -- if Ψ | Γ, x : A, y : B ⊢ σ · body ⇐
    --    ret[(x, y) ∷ (x : A) × B/p] ⊳ Σ₂, ρ₁·x, ρ₂·y
    -- with ρ₁, ρ₂ ≤ πσ
    let [< x, y]  = body.names
        pisg    = pi * sg.fst
        bodyctx = extendTyN [< (pisg, x, tfst), (pisg, y, tsnd.term)] ctx
        bodyty  = substCasePairRet pairres.type ret
    bodyout <- checkC bodyctx sg body.term bodyty
    -- then Ψ | Γ ⊢ σ · case ⋯ ⇒ ret[pair/p] ⊳ πΣ₁ + Σ₂
    pure $ InfRes {
      type = sub1 ret pair,
      qout = pi * pairres.qout + !(popQs [< pisg, pisg] bodyout)
    }

  infer' ctx sg (CaseEnum pi t ret arms) {d, n} = do
    -- if 1 ≤ π
    expectCompatQ one pi
    -- if Ψ | Γ ⊢ σ · t ⇒ {ts} ⊳ Σ₁
    tres <- inferC ctx sg t
    ttags <- expectEnum !ask ctx tres.type
    -- if Ψ | Γ, x : {ts} ⊢₀ A ⇐ Type
    checkTypeC (extendTy zero ret.name tres.type ctx) ret.term Nothing
    -- if for each "a ⇒ s" in arms,
    --   Ψ | Γ ⊢ σ · s ⇐ A[a ∷ {ts}/x] ⊳ Σ₂
    -- for fixed Σ₂
    let arms = SortedMap.toList arms
    let armTags = SortedSet.fromList $ map fst arms
    unless (ttags == armTags) $ throwError $ BadCaseEnum ttags armTags
    armres <- for arms $ \(a, s) =>
      checkC ctx sg s (sub1 ret (Tag a :# tres.type))
    armout <- allEqual (zip armres arms)
    -- then Ψ | Γ ⊢ σ · case ⋯ ⇒ ret[t/x] ⊳ πΣ₁ + Σ₂
    pure $ InfRes {
      type = sub1 ret t,
      qout = pi * tres.qout + armout
    }
   where
    allEqual : List (QOutput q n, TagVal, Term q d n) -> m (QOutput q n)
    allEqual [] = pure $ zeroFor ctx
    allEqual lst@((x, _) :: xs) =
      if all (\y => x == fst y) xs then pure x
      else throwError $ BadCaseQtys ctx $
             map (\(qs, t, s) => (qs, Tag t, s)) lst

  infer' ctx sg (CaseNat pi pi' n ret zer suc) = do
    -- if 1 ≤ π
    expectCompatQ one pi
    -- if Ψ | Γ ⊢ σ · n ⇒ ℕ ⊳ Σn
    nres <- inferC ctx sg n
    expectNat !ask ctx nres.type
    -- if Ψ | Γ, n : ℕ ⊢₀ A ⇐ Type
    checkTypeC (extendTy zero ret.name Nat ctx) ret.term Nothing
    -- if Ψ | Γ ⊢ σ · zer ⇐ A[0 ∷ ℕ/n] ⊳ Σz
    zerout <- checkC ctx sg zer (sub1 ret (Zero :# Nat))
    -- if Ψ | Γ, n : ℕ, ih : A ⊢ σ · suc ⇐ A[succ p ∷ ℕ/n] ⊳ Σs, ρ₁.p, ρ₂.ih
    -- with Σz = Σs, (ρ₁ + ρ₂) ≤ πσ
    let [< p, ih] = suc.names
        pisg      = pi * sg.fst
        sucCtx    = extendTyN [< (pisg, p, Nat), (pi', ih, ret.term)] ctx
        sucType   = substCaseNatRet ret
    sucout :< qp :< qih <- checkC sucCtx sg suc.term sucType
    unless (zerout == sucout) $ do
      let sucp = Succ $ FT $ unq p
          suc  = subN suc [< F $ unq p, F $ unq ih]
      throwError $ BadCaseQtys ctx [(zerout, Zero, zer), (sucout, sucp, suc)]
    expectCompatQ qih pi'
    -- [fixme] better error here
    expectCompatQ (qp + qih) pisg
    -- then Ψ | Γ ⊢ case ⋯ ⇒ A[n] ⊳ πΣn + Σz
    pure $ InfRes {
      type = sub1 ret n,
      qout = pi * nres.qout + zerout
    }

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

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