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 : (q : Nat) -> Type where
  MonoDef   : Term 0 0 0 -> DefBody 0
  PolyDef   : (0 nz : IsSucc q) => Term q 0 0 -> DefBody q
  Postulate : DefBody q

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


public export
record Definition where
  constructor MkDef
  qty    : GQty
  {qlen  : Nat}
  qargs  : BContext qlen
  type0  : Term qlen 0 0
  body0  : DefBody qlen
  scheme : Maybe String
  isMain : Bool
  loc_   : Loc

public export %inline
mkPostulate : GQty -> BContext q -> Term q 0 0 ->
              Maybe String -> Bool -> Loc -> Definition
mkPostulate qty qargs type0 scheme isMain loc_ =
  let Val q = lengthPrf0 qargs in
  MkDef {qty, qargs, type0, body0 = Postulate, scheme, isMain, loc_}

public export %inline
mkDef : GQty -> BContext q -> (type0, term0 : Term q 0 0) -> Maybe String ->
        Bool -> Loc -> Definition
mkDef qty qargs type0 term0 scheme isMain loc_ =
  case (lengthPrf0 qargs) of
    Val 0 =>
      MkDef {qty, qargs, type0, body0 = MonoDef term0, scheme, isMain, loc_}
    Val (S _) =>
      MkDef {qty, qargs, type0, body0 = PolyDef term0, scheme, isMain, loc_}

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


public export
record Poly (0 tm : TermLike) d n where
  constructor P
  qlen : Nat
  type : tm qlen d n


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

  public export %inline
  (.typeAt) : Definition -> Universe -> Poly Term d n
  def.typeAt u = {type $= displace u} def.type

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

  public export %inline
  (.termAt) : Definition -> Universe -> Maybe (Poly Term d n)
  def.termAt u = {type $= displace u} <$> def.term

  public export %inline
  toElim : Definition -> Universe -> Maybe (Poly Elim d n)
  toElim def u = do
    tm <- def.body0.term0; let ty = def.type0
    pure $ P def.qlen $ Ann tm ty def.loc // shift0 d // shift0 n

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

public export
(.typeWithAt) : Definition -> Singleton d -> Singleton n ->
                Universe -> Poly Term d n
def.typeWithAt (Val d) (Val n) u = def.typeAt u

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


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


public export
NDefinition : Type
NDefinition = (Name, Definition)

public export
Definitions : Type
Definitions = SortedMap Name Definition

public export
data DefEnvTag = DEFS

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 (Poly Elim d n)
lookupElim x u defs = toElim !(lookup x defs) u

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


export
prettyQBinders : {opts : LayoutOpts} -> BContext q -> Eff Pretty (Doc opts)
prettyQBinders [<] = pure empty
prettyQBinders qnames =
  qbrackets . separateTight !commaD =<< traverse prettyQBind (toList' qnames)

export
prettyDef : {opts : LayoutOpts} -> Name -> Definition -> Eff Pretty (Doc opts)
prettyDef name def = withPrec Outer $ do
  qty   <- prettyQConst def.qty.qconst
  dot   <- dotD
  name  <- prettyFree name
  qargs <- prettyQBinders def.qargs
  colon <- colonD
  type  <- prettyTerm (fromQNames def.qargs) def.type0
  hangDSingle (hsep [hcat [qty, dot, name, qargs], colon]) type