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
typeD, arrowD, timesD, lamD, eqndD, dlamD, annD :
  Pretty.HasEnv m => m (Doc HL)
typeD   = hlF Syntax $ ifUnicode "★" "Type"
arrowD  = hlF Syntax $ ifUnicode "→" "->"
timesD  = hlF Syntax $ ifUnicode "×" "**"
lamD    = hlF Syntax $ ifUnicode "λ" "fun"
eqndD   = hlF Syntax $ ifUnicode "≡" "=="
dlamD   = hlF Syntax $ ifUnicode "δ" "dfun"
annD    = hlF Syntax $ ifUnicode "∷" "::"

private %inline
eqD, colonD, commaD, dotD, caseD, returnD, ofD : Doc HL
eqD     = hl Syntax "Eq"
colonD  = hl Syntax ":"
commaD  = hl Syntax ","
dotD    = hl Delim  "."
caseD   = hl Syntax "case"
ofD     = hl Syntax "of"
returnD = hl Syntax "return"


export covering
prettyBindType : PrettyHL a => PrettyHL b => PrettyHL q =>
                 Pretty.HasEnv m =>
                 List q -> BaseName -> a -> Doc HL -> b -> m (Doc HL)
prettyBindType qtys x s arr t =
    parensIfM Outer $ hang 2 $
      parens (hang 2 $ !(prettyQtyBinds qtys (TV x)) <++> colonD
                <//> !(prettyM s)) <++> arr
        <//> !(under T x $ prettyM t)

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

export covering
prettyApps : PrettyHL f => PrettyHL a => Pretty.HasEnv m =>
             Maybe (Doc HL) -> f -> List a -> m (Doc HL)
prettyApps pfx fun args =
  parensIfM App =<< withPrec Arg
    [|prettyM fun <//> (align . sep <$> traverse prettyArg args)|]
 where
  prettyArg : a -> m (Doc HL)
  prettyArg x = maybe id (\p => map (hl Delim p <+>)) pfx $ prettyM x

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

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

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

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

-- [fixme] put delimiters around tags that aren't simple names
export
prettyTag : TagVal -> Doc HL
prettyTag t = hl Tag $ "`" <+> fromString t


mutual
  export covering
  PrettyHL q => PrettyHL (Term q d n) where
    prettyM (TYPE l) =
      parensIfM App $ !typeD <+> hl Syntax !(prettyUnivSuffix l)
    prettyM (Pi qty s (S [x] t)) =
      prettyBindType [qty] x s !arrowD t
    prettyM (Lam (S x t)) =
      let GotLams {names, body, _} = getLams' x t.term Refl in
      prettyLams T (toList names) body
    prettyM (Sig s (S [x] t)) =
      prettyBindType {q} [] x s !timesD t
    prettyM (Pair s t) =
      let GotPairs {init, last, _} = getPairs' [< s] t in
      prettyTuple $ toList $ init :< last
    prettyM (Enum tags) =
      pure $ delims "`{" "}" . aseparate comma $ map prettyTag $
        Prelude.toList tags
    prettyM (Tag t) =
      pure $ prettyTag t
    prettyM (Eq (S _ (N ty)) l r) =
      parensIfM Eq !(withPrec InEq $ pure $
        sep [!(prettyM l) <++> !eqndD,
             !(prettyM r) <++> colonD, !(prettyM ty)])
    prettyM (Eq (S [i] (Y ty)) l r) =
      parensIfM App $
        eqD <++>
        sep [bracks !(withPrec Outer $ pure $ hang 2 $
               sep [hl DVar !(prettyM i) <+> dotD,
                    !(under D i $ prettyM ty)]),
             !(withPrec Arg $ prettyM l),
             !(withPrec Arg $ prettyM r)]
    prettyM (DLam (S 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 Nothing fun args
    prettyM (CasePair pi p (S [r] ret) (S [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 (CaseEnum pi t (S [r] ret) arms) =
      prettyCase pi t r ret
        [([], prettyTag t, b) | (t, b) <- SortedMap.toList arms]
    prettyM (e :% d) =
      let GotDArgs {fun, args, _} = getDArgs' e [d] in
      prettyApps (Just "@") 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 (ScopedBody s (Term 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