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


export %inline
typeD, arrowD, darrowD, timesD, lamD, eqndD, dlamD, annD, natD :
  Pretty.HasEnv m => m (Doc HL)
typeD   = hlF Syntax $ ifUnicode "★" "Type"
arrowD  = hlF Delim  $ ifUnicode "→" "->"
darrowD = hlF Delim  $ ifUnicode "⇒" "=>"
timesD  = hlF Delim  $ ifUnicode "×" "**"
lamD    = hlF Syntax $ ifUnicode "λ" "fun"
eqndD   = hlF Delim  $ ifUnicode "≡" "=="
dlamD   = hlF Syntax $ ifUnicode "δ" "dfun"
annD    = hlF Delim  $ ifUnicode "∷" "::"
natD    = hlF Syntax $ ifUnicode "ℕ" "Nat"

export %inline
eqD, colonD, commaD, semiD, caseD, typecaseD, returnD,
ofD, dotD, zeroD, succD, coeD, compD : Doc HL
eqD       = hl Syntax "Eq"
colonD    = hl Delim  ":"
commaD    = hl Delim  ","
semiD     = hl Delim  ";"
caseD     = hl Syntax "case"
typecaseD = hl Syntax "type-case"
ofD       = hl Syntax "of"
returnD   = hl Syntax "return"
dotD      = hl Delim "."
zeroD     = hl Syntax "zero"
succD     = hl Syntax "succ"
coeD      = hl Syntax "coe"
compD     = hl Syntax "compD"


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


public export
data WithQty a = MkWithQty Qty a

export
PrettyHL a => PrettyHL (WithQty a) where
  prettyM (MkWithQty q x) = do
    q <- pretty0M q
    x <- withPrec Arg $ prettyM x
    pure $ hcat [q, dotD, x]


public export
data Binder a = MkBinder BaseName a

export
PrettyHL a => PrettyHL (Binder a) where
  prettyM (MkBinder x ty) = do
    x  <- pretty0M $ TV x
    ty <- align <$> pretty0M ty
    pure $ parens $ sep [hsep [x, colonD], ty]


export
prettyBindType : PrettyHL a => PrettyHL b =>
                 Pretty.HasEnv m =>
                 Maybe Qty -> BindName -> a -> Doc HL -> b -> m (Doc HL)
prettyBindType q (BN x _) s arr t = do
  bind <- case q of
    Nothing => pretty0M $ MkBinder x s
    Just q  => pretty0M $ MkWithQty q $ MkBinder x s
  t <- withPrec AnnR $ under T x $ prettyM t
  parensIfM AnnR $ hang 2 $ bind <++> arr <%%> t

export
prettyArm : PrettyHL a => Pretty.HasEnv m =>
            BinderSort -> SnocList BindName -> Doc HL -> a -> m (Doc HL)
prettyArm sort xs pat body = do
  let xs = map name xs
  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 BindName -> 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


public export
data TypeLine a = MkTypeLine BindName a

export
PrettyHL a => PrettyHL (TypeLine a) where
  prettyM (MkTypeLine i ty) =
    if i.name == Unused then
      bracks <$> pretty0M ty
    else
      map bracks $ withPrec Outer $ prettyLams Nothing D [< i] ty


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

export
prettyApps : PrettyHL f => PrettyHL a => Pretty.HasEnv m =>
             Maybe (Doc HL) -> f -> List a -> m (Doc HL)
prettyApps pfx f args = prettyApps' f (map (pfx,) args)

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 =>
             BinderSort -> List (SnocList BindName, Doc HL, a) -> m (Doc HL)
prettyArms s =
  map (braces . aseparate semiD) .
  traverse (\(xs, l, r) => prettyArm s xs l r)

export
prettyCase' : (PrettyHL a, PrettyHL b, PrettyHL c, Pretty.HasEnv m) =>
              Doc HL -> a -> BindName -> b ->
              List (SnocList BindName, Doc HL, c) ->
              m (Doc HL)
prettyCase' intro elim r ret arms = do
  elim  <- pretty0M elim
  ret   <- case r.name of
    Unused => under T r.name $ pretty0M ret
    _      => prettyLams Nothing T [< r] ret
  arms  <- prettyArms T arms
  pure $ asep [intro <++> elim, returnD <++> ret, ofD <++> arms]

export
prettyCase : (PrettyHL a, PrettyHL b, PrettyHL c, Pretty.HasEnv m) =>
             Qty -> a -> BindName -> b ->
             List (SnocList BindName, Doc HL, c) ->
             m (Doc HL)
prettyCase pi elim r ret arms = do
  caseq <- (caseD <+>) <$> prettySuffix pi
  prettyCase' caseq elim r ret 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

export
prettyBoxVal : PrettyHL a => Pretty.HasEnv m => a -> m (Doc HL)
prettyBoxVal val = bracks <$> pretty0M val

export
prettyCompPat : Pretty.HasEnv m => DimConst -> BindName -> m (Doc HL)
prettyCompPat e j = hsep <$> sequence [pretty0M e, pretty0M $ DV j.name]

export
toNatLit : Term d n -> Maybe Nat
toNatLit (Zero _)   = Just 0
toNatLit (Succ n _) = [|S $ toNatLit n|]
toNatLit _          = Nothing

private
eterm : Term d n -> Exists (Term d)
eterm = Evidence n


parameters (showSubsts : Bool)
 mutual
  export covering
  [TermSubst] PrettyHL (Term d n) using ElimSubst
  where
    prettyM (TYPE l _) =
      pure $ !typeD <+> hl Syntax !(prettyUnivSuffix l)

    prettyM (Pi qty s (S _ (N t)) _) = do
      dom <- pretty0M $ MkWithQty qty s
      cod <- withPrec AnnR $ prettyM t
      parensIfM AnnR $ asep [dom <++> !arrowD, cod]

    prettyM (Pi qty s (S [< x] (Y t)) _) =
      prettyBindType (Just qty) x s !arrowD t

    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 Nothing 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
      prettyApps Nothing (L eqD)
        [epretty $ MkTypeLine i ty, epretty l, epretty r]

    prettyM (DLam (S i t) _) =
      let GotDLams {names, body, _} = getDLams' i t.term Refl in
      prettyLams (Just !dlamD) D (toSnocList' names) body

    prettyM (Nat _) = natD

    prettyM (Zero _) = pure $ hl Syntax "0"

    prettyM (Succ n _) =
      case toNatLit n of
        Just n  => pure $ hl Syntax $ pretty $ S n
        Nothing => prettyApps Nothing (L succD) [n]

    prettyM (BOX pi ty _) = do
      pi <- pretty0M pi
      ty <- pretty0M ty
      pure $ bracks $ hcat [pi, dotD, align ty]

    prettyM (Box val _) = prettyBoxVal val

    prettyM (E e) = prettyM e

    prettyM (CloT (Sub s th)) =
      if showSubsts then
        parensIfM SApp . hang 2 =<<
          [|withPrec SApp (prettyM s) <%> prettyTSubst th|]
      else
        prettyM $ pushSubstsWith' id th s

    prettyM (DCloT (Sub 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 (Elim d n) using TermSubst
  where
    prettyM (F x _) =
      hl' Free <$> prettyM x

    prettyM (B i _) =
      prettyVar TVar TVarErr (!ask).tnames i

    prettyM (App 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 (CaseNat pi pi' nat (S [< r] ret) zer (S [< s, ih] suc) _) =
      prettyCase pi nat r ret.term
        [([<],       zeroD,    eterm zer),
         ([< s, ih], !succPat, eterm suc.term)]
    where
      succPat : m (Doc HL)
      succPat = case (ih, pi') of
        (BN Unused _, Zero) => pure $ succD <++> !(pretty0M s)
        _ => pure $ asep [succD <++> !(pretty0M s) <+> comma,
                          !(pretty0M $ MkWithQty pi' ih)]

    prettyM (CaseBox pi box (S [< r] ret) (S [< u] body) _) =
      prettyCase pi box r ret.term
        [([< u], !(prettyBoxVal $ TV u.name), body.term)]

    prettyM (DApp e d _) =
      let GotDArgs {fun, args, _} = getDArgs' e [d] in
      prettyApps (Just "@") fun args

    prettyM (Ann s a _) = do
      s <- withPrec AnnL $ prettyM s
      a <- withPrec AnnR $ prettyM a
      parensIfM AnnR $ hang 2 $ s <++> !annD <%%> a

    prettyM (Coe (S [< i] ty) p q val _) =
      let ty = case ty of
            Y ty => epretty $ MkTypeLine i ty
            N ty => epretty ty
      in
      prettyApps' (L coeD)
        [(Nothing,  ty),
         (Just "@", epretty p),
         (Just "@", epretty q),
         (Nothing,  epretty val)]

    prettyM (Comp ty p q val r (S [< z] zero) (S [< o] one) _) = do
      apps <- prettyApps' (L compD)
        [(Nothing,  epretty $ MkTypeLine (BN Unused noLoc) ty),
         (Just "@", epretty p),
         (Just "@", epretty q),
         (Nothing,  epretty val),
         (Just "@", epretty r)]
      arms <- prettyArms D
        [([< z], !(prettyCompPat Zero z), zero.term),
         ([< o], !(prettyCompPat One  o), one.term)]
      pure $ apps <++> arms

    prettyM (TypeCase ty ret arms def _) = do
      arms <- traverse fromArm (toList arms)
      prettyCase' typecaseD ty (BN Unused noLoc) ret $
        arms ++ [([<], hl Syntax "_", eterm def)]
    where
      v : BindName -> Doc HL
      v = pretty0 True . TV . name

      tyCasePat : (k : TyConKind) -> BContext (arity k) -> m (Doc HL)
      tyCasePat KTYPE [<] = typeD
      tyCasePat KPi   [< a, b] = pure $ parens $ hsep [v a, !arrowD, v b]
      tyCasePat KSig  [< a, b] = pure $ parens $ hsep [v a, !arrowD, v b]
      tyCasePat KEnum [<]   = pure $ hl Syntax "{}"
      tyCasePat KNat  [<]   = natD
      tyCasePat KBOX  [< a] = pure $ bracks $ v a
      tyCasePat KEq   vars  =
        prettyApps Nothing (L eqD) $ map (TV . name) $ toList' vars

      fromArm : TypeCaseArm d n ->
                m (SnocList BindName, Doc HL, Exists (Term d))
      fromArm (k ** S ns t) =
        pure (toSnocList' ns, !(tyCasePat k ns), eterm t.term)

    prettyM (CloE (Sub e th)) =
      if showSubsts then
        parensIfM SApp . hang 2 =<<
          [|withPrec SApp (prettyM e) <%> prettyTSubst th|]
      else
        prettyM $ pushSubstsWith' id th e

    prettyM (DCloE (Sub 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 => TSubst d from to -> m (Doc HL)
  prettyTSubst s =
    prettySubstM (prettyM @{ElimSubst}) (!ask).tnames TVar "[" "]" s

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

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


export covering
prettyTerm : (unicode : Bool) ->
             (dnames : BContext d) -> (tnames : BContext n) ->
             Term d n -> Doc HL
prettyTerm unicode dnames tnames term =
  pretty0With unicode (toNames dnames) (toNames tnames) term