module Quox.Typing.Context

import Quox.Syntax
import Quox.Context
import Quox.Pretty
import public Quox.SingletonExtra
import Derive.Prelude

%default total
%language ElabReflection


public export
QContext : (q, n : Nat) -> Type
QContext q n = Context' (Qty q) n

public export
record LocalVar q d n where
  constructor MkLocal
  type : Term q d n
  term : Maybe (Term q d n) -- if from a `let`
-- %runElab deriveIndexed "LocalVar" [Show]

export
[LocateVarType] Located (LocalVar q d n) where x.loc = x.type.loc

namespace LocalVar
  export %inline
  letVar : (type, term : Term q d n) -> LocalVar q d n
  letVar type term = MkLocal {type, term = Just term}

  export %inline
  lamVar : (type : Term q d n) -> LocalVar q d n
  lamVar type = MkLocal {type, term = Nothing}

  export %inline
  mapVar : (Term q d n -> Term q' d' n') -> LocalVar q d n -> LocalVar q' d' n'
  mapVar f = {type $= f, term $= map f}

  export %inline
  subQ : {q1, q2 : Nat} -> QSubst q1 q2 -> LocalVar q1 d n -> LocalVar q2 d n
  subQ th = mapVar (// th)

  export %inline
  weakQ : {q : Nat} -> LocalVar q d n -> LocalVar (S q) d n
  weakQ = subQ $ shift 1

  export %inline
  subD : DSubst d1 d2 -> LocalVar q d1 n -> LocalVar q d2 n
  subD th = mapVar (// th)

  export %inline
  weakD : LocalVar q d n -> LocalVar q (S d) n
  weakD = subD $ shift 1

export %inline CanShift  (LocalVar q d) where l // by = mapVar (// by) l
export %inline CanQSubst LocalVar       where l // th = mapVar (// th) l
export %inline CanDSubst LocalVar       where l // th = mapVar (// th) l
export %inline CanTSubst LocalVar       where l // th = mapVar (// th) l

public export
TContext : TermLike
TContext q d = Context (LocalVar q d)

public export
QOutput : (q, n : Nat) -> Type
QOutput = QContext

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


public export
record TyContext q d n where
  constructor MkTyContext
  {auto qtyLen  : Singleton q}
  {auto dimLen  : Singleton d}
  {auto termLen : Singleton n}
  qnames : BContext q -- only used for printing
  dctx   : DimEq d
  dnames : BContext d -- only used for printing
  tctx   : TContext q d n
  tnames : BContext n -- only used for printing
  qtys   : QContext q n -- only used for printing
%name TyContext ctx
-- %runElab deriveIndexed "TyContext" [Show]


public export
record EqContext q n where
  constructor MkEqContext
  {auto qtyLen : Singleton q}
  {dimLen : Nat}
  {auto termLen : Singleton n}
  qnames  : BContext q -- only used for printing
  dassign : DimAssign dimLen -- only used for printing
  dnames  : BContext dimLen -- only used for printing
  tctx    : TContext q 0 n
  tnames  : BContext n -- only used for printing
  qtys    : QContext q n -- only used for printing
%name EqContext ctx
-- %runElab deriveIndexed "EqContext" [Show]


public export
record WhnfContext q d n where
  constructor MkWhnfContext
  {auto qtyLen  : Singleton q}
  {auto dimLen  : Singleton d}
  {auto termLen : Singleton n}
  qnames : BContext q -- only used for printing
  dnames : BContext d
  tnames : BContext n
  tctx   : TContext q d n
%name WhnfContext ctx
-- %runElab deriveIndexed "WhnfContext" [Show]

namespace TContext
  export %inline
  zeroFor : {q : Nat} -> Located1 tm => Context tm n -> QOutput q n
  zeroFor = map $ \x => zero x.loc

public export
extendLen : Telescope a n1 n2 -> Singleton n1 -> Singleton n2
extendLen [<]        x = x
extendLen (tel :< _) x = [|S $ extendLen tel x|]


public export
CtxExtension : (q, d, n1, n2 : Nat) -> Type
CtxExtension q d = Telescope $ \n => (Qty q, BindName, Term q d n)

public export
CtxExtension0 : (q, d, n1, n2 : Nat) -> Type
CtxExtension0 q d = Telescope $ \n => (BindName, Term q d n)

public export
CtxExtensionLet : (q, d, n1, n2 : Nat) -> Type
CtxExtensionLet q d = Telescope $ \n => (Qty q, BindName, LocalVar q d n)

public export
CtxExtensionLet0 : (q, d, n1, n2 : Nat) -> Type
CtxExtensionLet0 q d = Telescope $ \n => (BindName, LocalVar q d n)

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

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

  export %inline
  extendTyLetN : CtxExtensionLet q d n1 n2 ->
                 TyContext q d n1 -> TyContext q d n2
  extendTyLetN xss
      (MkTyContext {termLen, qnames, dctx, dnames, tctx, tnames, qtys}) =
    let (qs, xs, ls) = unzip3 xss in
    MkTyContext {
      qnames, dctx, dnames,
      termLen = extendLen xss termLen,
      tctx    = tctx . ls,
      tnames  = tnames . xs,
      qtys    = qtys . qs
    }

  export %inline
  extendTyN : CtxExtension q d n1 n2 -> TyContext q d n1 -> TyContext q d n2
  extendTyN = extendTyLetN . map (\(q, x, s) => (q, x, lamVar s))

  export %inline
  extendTyLetN0 : CtxExtensionLet0 q d n1 n2 ->
                  TyContext q d n1 -> TyContext q d n2
  extendTyLetN0 xss ctx =
    let Val q = ctx.qtyLen in
    extendTyLetN (map (\(x, t) => (zero x.loc, x, t)) xss) ctx

  export %inline
  extendTyN0 : CtxExtension0 q d n1 n2 ->
               TyContext q d n1 -> TyContext q d n2
  extendTyN0 xss ctx =
    let Val q = ctx.qtyLen in
    extendTyN (map (\(x, t) => (zero x.loc, x, t)) xss) ctx

  export %inline
  extendTyLet : Qty q -> BindName -> Term q d n -> Term q d n ->
                TyContext q d n -> TyContext q d (S n)
  extendTyLet q x s e = extendTyLetN [< (q, x, letVar s e)]

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

  export %inline
  extendTy0 : BindName -> Term q d n ->
              TyContext q d n -> TyContext q d (S n)
  extendTy0 x t ctx = let Val q = ctx.qtyLen in extendTy (zero x.loc) x t ctx

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

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

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

  public export
  (.names) : TyContext q d n -> NameContexts q d n
  (MkTyContext {qnames, dnames, tnames, _}).names =
    MkNameContexts {qnames, dnames, tnames}


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

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

  export %inline
  zeroFor : TyContext q _ n -> QOutput q n
  zeroFor ctx =
    let Val q = ctx.qtyLen in
    zeroFor ctx.tctx @{LocateVarType}


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

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

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

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

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

  export %inline
  extendTyLetN : CtxExtensionLet q 0 n1 n2 -> EqContext q n1 -> EqContext q n2
  extendTyLetN xss
      (MkEqContext {termLen, qnames, dassign, dnames, tctx, tnames, qtys}) =
    let (qs, xs, ls) = unzip3 xss in
    MkEqContext {
      termLen = extendLen xss termLen,
      tctx    = tctx . ls,
      tnames  = tnames . xs,
      qtys    = qtys . qs,
      dassign, dnames, qnames
    }

  export %inline
  extendTyN : CtxExtension q 0 n1 n2 -> EqContext q n1 -> EqContext q n2
  extendTyN = extendTyLetN . map (\(q, x, s) => (q, x, lamVar s))

  export %inline
  extendTyLetN0 : CtxExtensionLet0 q 0 n1 n2 -> EqContext q n1 -> EqContext q n2
  extendTyLetN0 xss ctx =
    let Val q = ctx.qtyLen in
    extendTyLetN (map (\(x, t) => (zero x.loc, x, t)) xss) ctx

  export %inline
  extendTyN0 : CtxExtension0 q 0 n1 n2 -> EqContext q n1 -> EqContext q n2
  extendTyN0 xss ctx =
    let Val q = ctx.qtyLen in
    extendTyN (map (\(x, t) => (zero x.loc, x, t)) xss) ctx

  export %inline
  extendTyLet : Qty q -> BindName -> Term q 0 n -> Term q 0 n ->
                EqContext q n -> EqContext q (S n)
  extendTyLet q x s e = extendTyLetN [< (q, x, letVar s e)]

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

  export %inline
  extendTy0 : BindName -> Term q 0 n -> EqContext q n -> EqContext q (S n)
  extendTy0 x t ctx = let Val q = ctx.qtyLen in extendTy (zero x.loc) x t ctx

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

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

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

  export
  injElim : WhnfContext q d n -> Elim q 0 0 -> Elim q d n
  injElim ctx e =
    let Val d = ctx.dimLen; Val n = ctx.termLen in
    e // shift0 d // shift0 n

  public export
  (.names) : (e : EqContext q n) -> NameContexts q e.dimLen n
  (MkEqContext {qnames, dnames, tnames, _}).names =
    MkNameContexts {qnames, dnames, tnames}

  public export
  (.names0) : EqContext q n -> NameContexts q 0 n
  (MkEqContext {qnames, tnames, _}).names0 =
    MkNameContexts {qnames, dnames = [<], tnames}

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

  export
  extendTy' : BindName -> LocalVar q d n ->
              WhnfContext q d n -> WhnfContext q d (S n)
  extendTy' x var (MkWhnfContext {termLen, qnames, dnames, tnames, tctx}) =
    MkWhnfContext {
      dnames, qnames,
      termLen = [|S termLen|],
      tnames  = tnames :< x,
      tctx    = tctx   :< var
    }

  export %inline
  extendTy : BindName -> Term q d n ->
             WhnfContext q d n -> WhnfContext q d (S n)
  extendTy x ty ctx = extendTy' x (lamVar ty) ctx

  export %inline
  extendTyLet : BindName -> (type, term : Term q d n) ->
                WhnfContext q d n -> WhnfContext q d (S n)
  extendTyLet x type term ctx = extendTy' x (letVar {type, term}) ctx

  export
  extendDimN : {s : Nat} -> BContext s -> WhnfContext q d n ->
               WhnfContext q (s + d) n
  extendDimN ns (MkWhnfContext {qnames, dnames, tnames, tctx, dimLen}) =
    MkWhnfContext {
      dimLen = [|Val s + dimLen|],
      dnames = dnames ++ toSnocVect' ns,
      tctx   = map (subD $ shift s) tctx,
      qnames, tnames
    }

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

  public export
  (.names) : WhnfContext q d n -> NameContexts q d n
  (MkWhnfContext {qnames, dnames, tnames, _}).names =
    MkNameContexts {qnames, dnames, tnames}


private
prettyTContextElt : {opts : _} ->
                    NameContexts q d n ->
                    Doc opts -> BindName -> LocalVar q d n ->
                    Eff Pretty (Doc opts)
prettyTContextElt names q x s = do
  let Val _ = lengthPrf0 names.qnames
  dot <- dotD
  x  <- prettyTBind x; colon <- colonD
  ty <- withPrec Outer $ prettyTerm names s.type; eq <- cstD
  tm <- traverse (withPrec Outer . prettyTerm names) s.term
  d  <- askAt INDENT
  let qx = hcat [q, dot, x]
  pure $ case tm of
    Nothing =>
      ifMultiline (hsep [qx, colon, ty]) (vsep [qx, indent d $ colon <++> ty])
    Just tm =>
      ifMultiline (hsep [qx, colon, ty, eq, tm])
                  (vsep [qx, indent d $ colon <++> ty, indent d $ eq <++> tm])

private
prettyTContext' : {opts : _} ->
                  NameContexts q d n -> Context' (Doc opts) n ->
                  TContext q d n -> Eff Pretty (SnocList (Doc opts))
prettyTContext' _ [<] [<] = pure [<]
prettyTContext' names (qtys :< q) (tys :< t) =
  let names' = {tnames $= tail, termLen $= map pred} names in
  [|prettyTContext' names' qtys tys :<
    prettyTContextElt names' q (head names.tnames) t|]

export
prettyTContext : {opts : _} ->
                 NameContexts q d n -> QContext q n ->
                 TContext q d n -> Eff Pretty (Doc opts)
prettyTContext names qtys tys = do
  qtys <- traverse (prettyQty names.qnames) qtys
  sepSingleTight !commaD . toList <$> prettyTContext' names qtys tys

export
prettyTyContext : {opts : _} -> TyContext q d n -> Eff Pretty (Doc opts)
prettyTyContext ctx = case ctx.dctx of
  C [<] => prettyTContext ctx.names ctx.qtys ctx.tctx
  _     => pure $
    sepSingle [!(prettyDimEq ctx.dnames ctx.dctx) <++> !pipeD,
               !(prettyTContext ctx.names ctx.qtys ctx.tctx)]

export
prettyEqContext : {opts : _} -> EqContext q n -> Eff Pretty (Doc opts)
prettyEqContext ctx = prettyTyContext $ toTyContext ctx

export
prettyWhnfContext : {opts : _} -> WhnfContext q d n -> Eff Pretty (Doc opts)
prettyWhnfContext ctx =
  let Val n = ctx.termLen in
  sepSingleTight !commaD . toList <$>
    prettyTContext' ctx.names (replicate n "_") ctx.tctx