module Quox.Syntax.Term.Pretty

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

import Data.Vect

%default total


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


export
prettyUnivSuffix : Pretty.HasEnv m => Universe -> m (Doc HL)
prettyUnivSuffix l =
  ifUnicode (pretty $ pack $ map sub $ unpack $ show l) (pretty l)
 where
  sub : Char -> Char
  sub c = case c of
    '0' => '₀'; '1' => '₁'; '2' => '₂'; '3' => '₃'; '4' => '₄'
    '5' => '₅'; '6' => '₆'; '7' => '₇'; '8' => '₈'; '9' => '₉'; _ => c

export
prettyUniverse : Universe -> Doc HL
prettyUniverse = hl Syntax . pretty

export
prettyBind : PrettyHL a => PrettyHL q => Pretty.HasEnv m =>
             List q -> BaseName -> a -> m (Doc HL)
prettyBind qtys x s = do
  var <- prettyQtyBinds qtys $ TV x
  s   <- withPrec Outer $ prettyM s
  pure $ var <++> colonD <%%> hang 2 s

export
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 = do
  bind <- prettyBind qtys x s
  t    <- withPrec AnnR $ under T x $ prettyM t
  parensIfM AnnR $ hang 2 $ parens bind <++> arr <%%> t

export
prettyArm : PrettyHL a => Pretty.HasEnv m =>
            BinderSort -> SnocList BaseName -> Doc HL -> a -> m (Doc HL)
prettyArm sort xs pat body = do
  body <- withPrec Outer $ unders sort xs $ prettyM body
  pure $ hang 2 $ sep [pat <++> !darrowD, body]

export
prettyLams : PrettyHL a => Pretty.HasEnv m =>
             Maybe (Doc HL) -> BinderSort -> SnocList BaseName -> a ->
             m (Doc HL)
prettyLams lam sort names body = do
  let var = case sort of T => TVar; D => DVar
  header <- sequence $ [hlF var $ prettyM x | x <- toList names]
  let header = sep $ maybe header (:: header) lam
  parensIfM Outer =<< prettyArm sort names header body

export
prettyApps : PrettyHL f => PrettyHL a => Pretty.HasEnv m =>
             Maybe (Doc HL) -> f -> List a -> m (Doc HL)
prettyApps pfx fun args = do
  fun  <- withPrec App $ prettyM fun
  args <- traverse (withPrec Arg . prettyArg) args
  parensIfM App $ hang 2 $ sep $ fun :: args
 where
  prettyArg : a -> m (Doc HL)
  prettyArg = case pfx of
    Nothing  => prettyM
    Just pfx => \x => pure $ hl Delim pfx <+> !(prettyM x)

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

export
prettyArms : PrettyHL a => Pretty.HasEnv m =>
             List (SnocList BaseName, Doc HL, a) -> m (Doc HL)
prettyArms =
  map (braces . asep) . traverse (\(xs, l, r) => prettyArm T xs l r)

export
prettyCase : PrettyHL a => PrettyHL b => PrettyHL c => PrettyHL q =>
             Pretty.HasEnv m =>
             q -> a -> BaseName -> b -> List (SnocList BaseName, Doc HL, c) ->
             m (Doc HL)
prettyCase pi elim r ret arms = do
  elim <- prettyQtyBinds [pi] elim
  ret  <- prettyLams Nothing T [< r] ret
  arms <- prettyArms arms
  pure $ asep [caseD <++> elim, returnD <++> ret, ofD <++> arms]

export
escapeString : String -> String
escapeString = concatMap esc1 . unpack where
  esc1 : Char -> String
  esc1 '"'  = #"\""#
  esc1 '\\' = #"\\"#
  esc1 '\n' = #"\n"#
  esc1 c    = singleton c

export
quoteTag : TagVal -> Doc HL
quoteTag tag =
  if isName tag then fromString tag else
  hcat ["\"", fromString $ escapeString tag, "\""]

export
prettyTag : TagVal -> Doc HL
prettyTag t = hl Tag $ "'" <+> quoteTag t

export
prettyTagBare : TagVal -> Doc HL
prettyTagBare t = hl Tag $ quoteTag t


parameters (showSubsts : Bool)
 mutual
  export covering
  [TermSubst] PrettyHL q => PrettyHL (Term q d n) using TermSubst ElimSubst
  where
    prettyM (TYPE l) =
      parensIfM App $ !typeD <+> hl Syntax !(prettyUnivSuffix l)
    prettyM (Pi qty s (S [< x] t)) =
      prettyBindType [qty] x s !arrowD t.term
    prettyM (Lam (S x t)) =
      let GotLams {names, body, _} = getLams' x t.term Refl in
      prettyLams (Just !lamD) T (toSnocList' names) body
    prettyM (Sig s (S _ (N t))) = do
      s <- withPrec InTimes $ prettyM s
      t <- withPrec Times   $ prettyM t
      parensIfM Times $ asep [s <++> !timesD, t]
    prettyM (Sig s (S [< x] (Y 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 prettyTagBare $
        Prelude.toList tags
    prettyM (Tag t) =
      pure $ prettyTag t
    prettyM (Eq (S _ (N ty)) l r) = do
      l  <- withPrec InEq $ prettyM l
      r  <- withPrec InEq $ prettyM r
      ty <- withPrec InEq $ prettyM ty
      parensIfM Eq $ asep [l <++> !eqndD, r <++> colonD, ty]
    prettyM (Eq (S [< i] (Y ty)) l r) = do
      ty <- bracks <$> withPrec Outer (prettyLams Nothing D [< i] ty)
      l  <- withPrec Arg $ prettyM l
      r  <- withPrec Arg $ prettyM r
      parensIfM App $ eqD <++> asep [ty, l, r]
    prettyM (DLam (S i t)) =
      let GotDLams {names, body, _} = getDLams' i t.term Refl in
      prettyLams (Just !dlamD) D (toSnocList' names) body
    prettyM (E e) = prettyM e
    prettyM (CloT s th) =
      if showSubsts then
        parensIfM SApp . hang 2 =<<
          [|withPrec SApp (prettyM s) <%> prettyTSubst th|]
      else
        prettyM $ pushSubstsWith' id th s
    prettyM (DCloT s th) =
      if showSubsts then
        parensIfM SApp . hang 2 =<<
          [|withPrec SApp (prettyM s) <%> prettyDSubst th|]
      else
        prettyM $ pushSubstsWith' th id s

  export covering
  [ElimSubst] PrettyHL q => PrettyHL (Elim q d n) using TermSubst ElimSubst
  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 <$> traverse (hlF TVar . prettyM) [x, y]
      prettyCase pi p r ret.term [([< x, y], pat, body.term)]
    prettyM (CaseEnum pi t (S [< r] ret) arms) =
      prettyCase pi t r ret.term
        [([<], 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) = do
      s <- withPrec AnnL $ prettyM s
      a <- withPrec AnnR $ prettyM a
      parensIfM AnnR $ hang 2 $ s <++> !annD <%%> a
    prettyM (CloE e th)  =
      if showSubsts then
        parensIfM SApp . hang 2 =<<
          [|withPrec SApp (prettyM e) <%> prettyTSubst th|]
      else
        prettyM $ pushSubstsWith' id th e
    prettyM (DCloE e th) =
      if showSubsts then
        parensIfM SApp . hang 2 =<<
          [|withPrec SApp (prettyM e) <%> prettyDSubst th|]
      else
        prettyM $ pushSubstsWith' th id e

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

export covering %inline
PrettyHL q => PrettyHL (Term q d n) where
  prettyM = prettyM @{TermSubst False}

export covering %inline
PrettyHL q => PrettyHL (Elim q d n) where
  prettyM = prettyM @{ElimSubst False}


export covering
prettyTerm : PrettyHL q => (unicode : Bool) ->
             (dnames : NContext d) -> (tnames : NContext n) ->
             Term q d n -> Doc HL
prettyTerm unicode dnames tnames term =
  let env = MakePrettyEnv {
    dnames = toSnocList' dnames,
    tnames = toSnocList' tnames,
    unicode, prec = Outer
  } in
  runReader env $ prettyM term