module Quox.Typing.Context

import Quox.Syntax
import Quox.Context
import Quox.Pretty
import public Data.Singleton

%default total


public export
QContext : Nat -> Type
QContext = Context' Qty

public export
TContext : TermLike
TContext d = Context (Term d)

public export
QOutput : Nat -> Type
QOutput = Context' Qty

public export
DimAssign : Nat -> Type
DimAssign = Context' DimConst


public export
record TyContext d n where
  constructor MkTyContext
  {auto dimLen  : Singleton d}
  {auto termLen : Singleton n}
  dctx   : DimEq d
  dnames : NContext d
  tctx   : TContext d n
  tnames : NContext n
  qtys   : QContext n -- only used for printing
%name TyContext ctx


public export
record EqContext n where
  constructor MkEqContext
  {dimLen : Nat}
  {auto termLen : Singleton n}
  dassign : DimAssign dimLen -- only used for printing
  dnames  : NContext dimLen -- only used for printing
  tctx    : TContext 0 n
  tnames  : NContext n
  qtys    : QContext n -- only used for printing
%name EqContext ctx


public export
record WhnfContext d n where
  constructor MkWhnfContext
  dnames : NContext d
  tnames : NContext n
  tctx   : TContext d n
%name WhnfContext ctx

namespace TContext
  export %inline
  pushD : TContext d n -> TContext (S d) n
  pushD tel = map (// shift 1) tel

  export %inline
  zeroFor : Context tm n -> QOutput n
  zeroFor ctx = Zero <$ ctx

private
extendLen : Telescope a from to -> Singleton from -> Singleton to
extendLen [<]        x = x
extendLen (tel :< _) x = [|S $ extendLen tel x|]


public export
CtxExtension : Nat -> Nat -> Nat -> Type
CtxExtension d = Telescope ((Qty, BaseName,) . Term d)

namespace TyContext
  public export %inline
  empty : TyContext 0 0
  empty =
    MkTyContext {dctx = new, dnames = [<], tctx = [<], tnames = [<], qtys = [<]}

  public export %inline
  null : TyContext d n -> Bool
  null ctx = null ctx.dnames && null ctx.tnames

  export %inline
  extendTyN : CtxExtension d from to -> TyContext d from -> TyContext d to
  extendTyN xss (MkTyContext {termLen, dctx, dnames, tctx, tnames, qtys}) =
    let (qs, xs, ss) = unzip3 xss in
    MkTyContext {
      dctx, dnames,
      termLen = extendLen xss termLen,
      tctx    = tctx . ss,
      tnames  = tnames . xs,
      qtys    = qtys . qs
    }

  export %inline
  extendTy : Qty -> BaseName -> Term d n -> TyContext d n -> TyContext d (S n)
  extendTy q x s = extendTyN [< (q, x, s)]

  export %inline
  extendDim : BaseName -> TyContext d n -> TyContext (S d) n
  extendDim x (MkTyContext {dimLen, dctx, dnames, tctx, tnames, qtys}) =
    MkTyContext {
      dctx   = dctx :<? Nothing,
      dnames = dnames :< x,
      dimLen = [|S dimLen|],
      tctx   = pushD tctx,
      tnames, qtys
    }

  export %inline
  eqDim : Dim d -> Dim d -> TyContext d n -> TyContext d n
  eqDim p q = {dctx $= set p q, dimLen $= id, termLen $= id}

  export
  toWhnfContext : TyContext d n -> WhnfContext d n
  toWhnfContext (MkTyContext {dnames, tnames, tctx, _}) =
    MkWhnfContext {dnames, tnames, tctx}


namespace QOutput
  export %inline
  (+) : QOutput n -> QOutput n -> QOutput n
  (+) = zipWith (+)

  export %inline
  (*) : Qty -> QOutput n -> QOutput n
  (*) pi = map (pi *)

  export %inline
  zeroFor : TyContext _ n -> QOutput n
  zeroFor ctx = zeroFor ctx.tctx


export
makeDAssign : DSubst d 0 -> DimAssign d
makeDAssign (Shift SZ)   = [<]
makeDAssign (K e ::: th) = makeDAssign th :< e

export
makeEqContext' : {d : Nat} -> TyContext d n -> DSubst d 0 -> EqContext n
makeEqContext' ctx th = MkEqContext {
  termLen = ctx.termLen,
  dassign = makeDAssign th,
  dnames  = ctx.dnames,
  tctx    = map (// th) ctx.tctx,
  tnames  = ctx.tnames,
  qtys    = ctx.qtys
}

export
makeEqContext : TyContext d n -> DSubst d 0 -> EqContext n
makeEqContext ctx@(MkTyContext {dnames, _}) th =
  let (d' ** Refl) = lengthPrf0 dnames in makeEqContext' ctx th

namespace EqContext
  public export %inline
  empty : EqContext 0
  empty = MkEqContext {
    dassign = [<], dnames = [<], tctx = [<], tnames = [<], qtys = [<]
  }

  public export %inline
  null : EqContext n -> Bool
  null ctx = null ctx.dnames && null ctx.tnames

  export %inline
  extendTyN : Telescope (\n => (Qty, BaseName, Term 0 n)) from to ->
              EqContext from -> EqContext to
  extendTyN xss (MkEqContext {termLen, dassign, dnames, tctx, tnames, qtys}) =
    let (qs, xs, ss) = unzip3 xss in
    MkEqContext {
      termLen = extendLen xss termLen,
      tctx    = tctx . ss,
      tnames  = tnames . xs,
      qtys    = qtys . qs,
      dassign, dnames
    }

  export %inline
  extendTy : Qty -> BaseName -> Term 0 n -> EqContext n -> EqContext (S n)
  extendTy q x s = extendTyN [< (q, x, s)]

  export %inline
  extendDim : BaseName -> DimConst -> EqContext n -> EqContext n
  extendDim x e (MkEqContext {dassign, dnames, tctx, tnames, qtys}) =
    MkEqContext {dassign = dassign :< e, dnames = dnames :< x,
                 tctx, tnames, qtys}

  export
  toTyContext : (e : EqContext n) -> TyContext e.dimLen n
  toTyContext (MkEqContext {dimLen, dassign, dnames, tctx, tnames, qtys}) =
    MkTyContext {
      dctx = fromGround dassign,
      tctx = map (// shift0 dimLen) tctx,
      dnames, tnames, qtys
    }

  export
  toWhnfContext : (ectx : EqContext n) -> WhnfContext 0 n
  toWhnfContext (MkEqContext {tnames, tctx, _}) =
    MkWhnfContext {dnames = [<], tnames, tctx}

namespace WhnfContext
  public export %inline
  empty : WhnfContext 0 0
  empty = MkWhnfContext [<] [<] [<]

  export
  extendDimN : {s : Nat} -> NContext s -> WhnfContext d n ->
               WhnfContext (s + d) n
  extendDimN ns (MkWhnfContext {dnames, tnames, tctx}) =
    MkWhnfContext {
      dnames = dnames ++ toSnocVect' ns,
      tctx   = dweakT s <$> tctx,
      tnames
    }

  export
  extendDim : BaseName -> WhnfContext d n -> WhnfContext (S d) n
  extendDim i = extendDimN [< i]


private
data CtxBinder a = MkCtxBinder BaseName a

PrettyHL a => PrettyHL (CtxBinder a) where
  prettyM (MkCtxBinder x t) = pure $
    sep [hsep [!(pretty0M $ TV x), colonD], !(pretty0M t)]

parameters (unicode : Bool)
  private
  pipeD : Doc HL
  pipeD = hl Syntax "|"

  export covering
  prettyTContext : NContext d ->
                   QContext n -> NContext n ->
                   TContext d n -> Doc HL
  prettyTContext ds qs xs ctx = separate comma $ toList $ go qs xs ctx where
    go : QContext m -> NContext m -> TContext d m -> SnocList (Doc HL)
    go [<]       [<]       [<] = [<]
    go (qs :< q) (xs :< x) (ctx :< t) =
      let bind = MkWithQty q $ MkCtxBinder x t in
      go qs xs ctx :<
      runPrettyWith unicode (toSnocList' ds) (toSnocList' xs) (pretty0M bind)

  export covering
  prettyTyContext : TyContext d n -> Doc HL
  prettyTyContext (MkTyContext dctx dnames tctx tnames qtys) =
    case dctx of
      C [<] => prettyTContext dnames qtys tnames tctx
      _     => sep [prettyDimEq dnames dctx <++> pipeD,
                    prettyTContext dnames qtys tnames tctx]

  export covering
  prettyEqContext : EqContext n -> Doc HL
  prettyEqContext (MkEqContext dassign dnames tctx tnames qtys) =
    case dassign of
      [<] => prettyTContext [<] qtys tnames tctx
      _   => sep [prettyDimEq dnames (fromGround dassign) <++> pipeD,
                  prettyTContext [<] qtys tnames tctx]