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, timesD, darrowD, lamD, eqndD, dlamD, annD :
  Pretty.HasEnv m => m (Doc HL)
arrowD  = hlF Syntax $ ifUnicode "→"  "->"
timesD  = hlF Syntax $ ifUnicode "×"  "**"
darrowD = 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, commaD, caseD, returnD, ofD : Doc HL
typeD   = hl Syntax "Type"
eqD     = hl Syntax "Eq"
colonD  = hl Syntax ":"
commaD  = hl Syntax ","
caseD   = hl Syntax "case"
ofD     = hl Syntax "of"
returnD = hl Syntax "return"

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) =
      prettyBindType [qty] x s !arrowD t
    prettyM (Lam x t) =
      let GotLams {names, body, _} = getLams' [x] t.term Refl in
      prettyLams T (toList names) body
    prettyM (Sig x s t) =
      prettyBindType [] x s !timesD t
    prettyM (Pair s t) =
      let GotPairs {init, last, _} = getPairs t in
      prettyTuple $ s :: init ++ [last]
    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) <++> !darrowD,
                    !(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 (CasePair pi p r ret x y body) = do
      pat <- (parens . separate commaD . map (hl TVar)) <$>
             traverse prettyM [x, y]
      prettyCase pi p r ret [([x, y], pat, body)]
    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
  prettyBindType : Pretty.HasEnv m => PrettyHL q =>
                   List q -> Name -> Term q d n -> Doc HL ->
                   ScopeTerm q d n -> m (Doc HL)
  prettyBindType qtys x s arr t =
      parensIfM Outer $ hang 2 $
        !(prettyBinder qtys x s) <++> arr
          <//> !(under T x $ prettyM t)

  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] ++ [darrowD]
    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)|]

  export covering
  prettyTuple : Pretty.HasEnv m => PrettyHL a => List a -> m (Doc HL)
  prettyTuple = map (parens . align . separate commaD) . traverse prettyM

  export covering
  prettyArm : Pretty.HasEnv m => PrettyHL a =>
              (List Name, Doc HL, a) -> m (Doc HL)
  prettyArm (xs, pat, body) =
    pure $ hang 2 $ sep
      [hsep [pat, !darrowD],
       !(withPrec Outer $ unders T xs $ prettyM body)]

  export covering
  prettyArms : Pretty.HasEnv m => PrettyHL a =>
               List (List Name, Doc HL, a) -> m (Doc HL)
  prettyArms = map (braces . align . sep) . traverse prettyArm

  export covering
  prettyCase : Pretty.HasEnv m =>
               PrettyHL q => PrettyHL a => PrettyHL b => PrettyHL c =>
               q -> a -> Name -> b -> List (List Name, Doc HL, c) -> m (Doc HL)
  prettyCase pi elim r ret arms =
    pure $ align $ sep $
      [hsep [caseD,   !(prettyM elim), !(prettyQtyBinds [pi])],
       hsep [returnD, !(prettyM r), !darrowD, !(under T r $ prettyM ret)],
       hsep [ofD,     !(prettyArms arms)]]