module Quox.Syntax.Term.Pretty

import Quox.Syntax.Term.Base
import Quox.Syntax.Term.Split
import Quox.Syntax.Term.Subst
import Quox.Pretty

import Data.Vect

%default total


private %inline
arrowD, lamD, eqndD, dlamD, annD : Pretty.HasEnv m => m (Doc HL)
arrowD = hlF Syntax $ ifUnicode "→"  "->"
lamD   = hlF Syntax $ ifUnicode "λ"  "fun"
eqndD  = hlF Syntax $ ifUnicode "≡"  "=="
dlamD  = hlF Syntax $ ifUnicode "λᴰ" "dfun"
annD   = hlF Syntax $ ifUnicode "∷"  "::"

private %inline
typeD, eqD, colonD : Doc HL
typeD  = hl Syntax "Type"
eqD    = hl Syntax "Eq"
colonD = hl Syntax ":"

mutual
  export covering
  PrettyHL q => PrettyHL (Term q d n) where
    prettyM (TYPE l) =
      parensIfM App $ typeD <//> !(withPrec Arg $ prettyM l)
    prettyM (Pi qty x s t) =
      parensIfM Outer $ hang 2 $
        !(prettyBinder [qty] x s) <++> !arrowD
          <//> !(under T x $ prettyM t)
    prettyM (Lam x t) =
      let GotLams {names, body, _} = getLams' [x] t.term Refl in
      prettyLams T (toList names) body
    prettyM (Eq _ (DUnused ty) l r) =
      parensIfM Eq !(withPrec InEq $ pure $
        sep [!(prettyM l) <++> !eqndD, !(prettyM r) <++> colonD, !(prettyM ty)])
    prettyM (Eq i (DUsed ty) l r) =
      parensIfM App $
        eqD <++>
        sep [bracks !(withPrec Outer $ pure $ hang 2 $
               sep [hl DVar !(prettyM i) <++> !arrowD,
                    !(under D i $ prettyM ty)]),
             !(withPrec Arg $ prettyM l),
             !(withPrec Arg $ prettyM r)]
    prettyM (DLam i t) =
      let GotDLams {names, body, _} = getDLams' [i] t.term Refl in
      prettyLams D (toList names) body
    prettyM (E e) = bracks <$> prettyM e
    prettyM (CloT s th) =
      parensIfM SApp . hang 2 =<<
        [|withPrec SApp (prettyM s) </> prettyTSubst th|]
    prettyM (DCloT s th) =
      parensIfM SApp . hang 2 =<<
        [|withPrec SApp (prettyM s) </> prettyDSubst th|]

  export covering
  PrettyHL q => PrettyHL (Elim q d n) where
    prettyM (F x) =
      hl' Free <$> prettyM x
    prettyM (B i) =
      prettyVar TVar TVarErr (!ask).tnames i
    prettyM (e :@ s) =
      let GotArgs {fun, args, _} = getArgs' e [s] in
      prettyApps fun args
    prettyM (e :% d) =
      let GotDArgs {fun, args, _} = getDArgs' e [d] in
      prettyApps fun args
    prettyM (s :# a) =
      parensIfM Ann $ hang 2 $
        !(withPrec AnnL $ prettyM s) <++> !annD
          <//> !(withPrec Ann $ prettyM a)
    prettyM (CloE e th)  =
      parensIfM SApp . hang 2 =<<
        [|withPrec SApp (prettyM e) </> prettyTSubst th|]
    prettyM (DCloE e th) =
      parensIfM SApp . hang 2 =<<
        [|withPrec SApp (prettyM e) </> prettyDSubst th|]

  export covering
  {s : Nat} -> PrettyHL q => PrettyHL (ScopeTermN s q d n) where
    prettyM body = prettyM body.term

  export covering
  prettyTSubst : Pretty.HasEnv m => PrettyHL q =>
                 TSubst q d from to -> m (Doc HL)
  prettyTSubst s = prettySubstM prettyM (!ask).tnames TVar "[" "]" s

  export covering
  prettyBinder : Pretty.HasEnv m => PrettyHL q =>
                 List q -> Name -> Term q d n -> m (Doc HL)
  prettyBinder pis x a =
    pure $ parens $ hang 2 $
      hsep [hl TVar !(prettyM x),
            sep [!(prettyQtyBinds pis),
                 hsep [colonD, !(withPrec Outer $ prettyM a)]]]

  export covering
  prettyLams : Pretty.HasEnv m => PrettyHL q =>
               BinderSort -> List Name -> Term q _ _ -> m (Doc HL)
  prettyLams sort names body = do
    lam    <- case sort of T => lamD; D => dlamD
    header <- sequence $ [hl TVar <$> prettyM x | x <- names] ++ [arrowD]
    body   <- unders sort names $ prettyM body
    parensIfM Outer $ sep (lam :: header) <//> body

  export covering
  prettyApps : Pretty.HasEnv m => PrettyHL f => PrettyHL a =>
               f -> List a -> m (Doc HL)
  prettyApps fun args =
    parensIfM App =<< withPrec Arg
      [|prettyM fun <//> (align . sep <$> traverse prettyM args)|]