module Quox.Definition

import public Quox.No
import public Quox.Syntax
import Quox.Displace
import public Data.SortedMap
import public Quox.Loc
import Quox.Pretty
import Control.Eff
import Data.Singleton
import Decidable.Decidable


public export
data DefBody =
    Concrete (Term 0 0)
  | Postulate

namespace DefBody
  public export
  (.term0) : DefBody -> Maybe (Term 0 0)
  (Concrete t).term0 = Just t
  (Postulate).term0  = Nothing


public export
record Definition where
  constructor MkDef
  qty   : GQty
  type0 : Term 0 0
  body0 : DefBody
  loc_  : Loc

public export %inline
mkPostulate : GQty -> (type0 : Term 0 0) -> Loc -> Definition
mkPostulate qty type0 loc_ = MkDef {qty, type0, body0 = Postulate, loc_}

public export %inline
mkDef : GQty -> (type0, term0 : Term 0 0) -> Loc -> Definition
mkDef qty type0 term0 loc_ = MkDef {qty, type0, body0 = Concrete term0, loc_}

export Located Definition where def.loc = def.loc_
export Relocatable Definition where setLoc loc = {loc_ := loc}


parameters {d, n : Nat}
  public export %inline
  (.type) : Definition -> Term d n
  g.type = g.type0 // shift0 d // shift0 n

  public export %inline
  (.typeAt) : Definition -> Universe -> Term d n
  g.typeAt u = displace u g.type

  public export %inline
  (.term) : Definition -> Maybe (Term d n)
  g.term = g.body0.term0 <&> \t => t // shift0 d // shift0 n

  public export %inline
  (.termAt) : Definition -> Universe -> Maybe (Term d n)
  g.termAt u = displace u <$> g.term

  public export %inline
  toElim : Definition -> Universe -> Maybe $ Elim d n
  toElim def u = pure $ Ann !(def.termAt u) (def.typeAt u) def.loc

public export
(.typeWith) : Definition -> Singleton d -> Singleton n -> Term d n
g.typeWith (Val d) (Val n) = g.type

public export
(.typeWithAt) : Definition -> Singleton d -> Singleton n -> Universe -> Term d n
g.typeWithAt d n u = displace u $ g.typeWith d n

public export
(.termWith) : Definition -> Singleton d -> Singleton n -> Maybe (Term d n)
g.termWith (Val d) (Val n) = g.term


public export %inline
isZero : Definition -> Bool
isZero g = g.qty == GZero


public export
data DefEnvTag = DEFS

public export
Definitions : Type
Definitions = SortedMap Name Definition

public export
DefsReader : Type -> Type
DefsReader = ReaderL DEFS Definitions

public export
DefsState : Type -> Type
DefsState = StateL DEFS Definitions

public export %inline
lookupElim : {d, n : Nat} -> Name -> Universe -> Definitions -> Maybe (Elim d n)
lookupElim x u defs = toElim !(lookup x defs) u

public export %inline
lookupElim0 : Name -> Universe -> Definitions -> Maybe (Elim 0 0)
lookupElim0 = lookupElim


export
prettyDef : {opts : LayoutOpts} -> Name -> Definition -> Eff Pretty (Doc opts)
prettyDef name (MkDef qty type _ _) = withPrec Outer $ do
  qty   <- prettyQty qty.qty
  dot   <- dotD
  name  <- prettyFree name
  colon <- colonD
  type  <- prettyTerm [<] [<] type
  pure $ sep [hsep [hcat [qty, dot, name], colon], type]