module Quox.Pretty

import Quox.Name

import public Text.PrettyPrint.Prettyprinter.Doc
import Text.PrettyPrint.Prettyprinter.Render.String
import Text.PrettyPrint.Prettyprinter.Render.Terminal
import public Data.String
import Data.DPair
import Data.SnocList

import public Control.Monad.Identity
import public Control.Monad.Reader
import Derive.Prelude

%default total
%language ElabReflection

%hide TT.Name


public export
record PrettyOpts where
  constructor MakePrettyOpts
  unicode, color : Bool

public export
defPrettyOpts : PrettyOpts
defPrettyOpts = MakePrettyOpts {unicode = True, color = True}


public export
data HL
= Delim
| Free | TVar | TVarErr
| Dim  | DVar | DVarErr
| Qty
| Syntax
| Tag
%runElab derive "HL" [Eq, Ord, Show]

public export
data PPrec
= Outer
| AnnR    -- right of "∷"
| AnnL    -- left  of "∷"
| Eq      -- "_ ≡ _ : _"
| InEq    -- arguments of ≡
| Times   -- "_ × _"
| InTimes -- arguments of ×
  -- ...
| App  -- term/dimension application
| SApp -- substitution application
| Arg  -- argument to nonfix function
%runElab derive "PPrec" [Eq, Ord, Show]


export %inline
hl : HL -> Doc HL -> Doc HL
hl = annotate

export %inline
hl' : HL -> Doc HL -> Doc HL
hl' h = hl h . unAnnotate

export %inline
hlF : Functor f => HL -> f (Doc HL) -> f (Doc HL)
hlF = map . hl

export %inline
hlF' : Functor f => HL -> f (Doc HL) -> f (Doc HL)
hlF' = map . hl'


export %inline
delims : Doc HL -> Doc HL -> Doc HL -> Doc HL
delims l r doc = hl Delim l <+> doc <+> hl Delim r

export %inline
parens : Doc HL -> Doc HL
parens = delims "(" ")"

export %inline
bracks : Doc HL -> Doc HL
bracks = delims "[" "]"

||| includes spaces inside the braces
export %inline
braces : Doc HL -> Doc HL
braces doc = hl Delim "{" <++> doc <++> hl Delim "}"

export %inline
parensIf : Bool -> Doc HL -> Doc HL
parensIf True  = parens
parensIf False = id

export %inline
comma : Doc HL
comma = hl Delim ","


export %inline
asep : List (Doc a) -> Doc a
asep = align . sep

export
separate' : Doc a -> List (Doc a) -> List (Doc a)
separate' s []        = []
separate' s [x]       = [x]
separate' s (x :: xs) = x <+> s :: separate' s xs

export %inline
separate : Doc a -> List (Doc a) -> Doc a
separate s = sep . separate' s

export %inline
hseparate : Doc a -> List (Doc a) -> Doc a
hseparate s = hsep . separate' s

export %inline
vseparate : Doc a -> List (Doc a) -> Doc a
vseparate s = vsep . separate' s

export %inline
aseparate : Doc a -> List (Doc a) -> Doc a
aseparate s = align . separate s


public export
record PrettyEnv where
  constructor MakePrettyEnv
  ||| names of bound dimension variables
  dnames  : SnocList BaseName
  ||| names of bound term variables
  tnames  : SnocList BaseName
  ||| use non-ascii characters for syntax
  unicode : Bool
  ||| surrounding precedence level
  prec    : PPrec

public export
HasEnv : (Type -> Type) -> Type
HasEnv = MonadReader PrettyEnv

export %inline
ifUnicode : HasEnv m => (uni, asc : Lazy a) -> m a
ifUnicode uni asc = if (!ask).unicode then [|uni|] else [|asc|]

export %inline
parensIfM : HasEnv m => PPrec -> Doc HL -> m (Doc HL)
parensIfM d doc = pure $ parensIf ((!ask).prec > d) doc

export %inline
withPrec : HasEnv m => PPrec -> m a -> m a
withPrec d = local {prec := d}

public export data BinderSort = T | D

export %inline
unders : HasEnv m => BinderSort -> SnocList BaseName -> m a -> m a
unders T xs = local {prec := Outer, tnames $= (++ xs)}
unders D xs = local {prec := Outer, dnames $= (++ xs)}

export %inline
under : HasEnv m => BinderSort -> BaseName -> m a -> m a
under t x = unders t [< x]


public export
interface PrettyHL a where
  prettyM : HasEnv m => a -> m (Doc HL)

export %inline
pretty0M : (PrettyHL a, HasEnv m) => a -> m (Doc HL)
pretty0M = local {prec := Outer} . prettyM

export %inline
runPretty : (unicode : Bool) -> Reader PrettyEnv a -> a
runPretty unicode act =
  let env = MakePrettyEnv {dnames = [<], tnames = [<], unicode, prec = Outer} in
  runReader env act

export %inline
pretty0 : PrettyHL a => (unicode : Bool) -> a -> Doc HL
pretty0 unicode = runPretty unicode . prettyM


export
(forall a. PrettyHL (f a)) => PrettyHL (Exists f) where
  prettyM x = prettyM x.snd

export
PrettyHL a => PrettyHL (Subset a b) where
  prettyM x = prettyM x.fst


export PrettyHL BaseName where prettyM = pure . pretty . baseStr
export PrettyHL Name     where prettyM = pure . pretty . toDots


export
nameSeq : HL -> List Name -> Doc HL
nameSeq h = hl h . asep . map (pretty0 False)

export %inline
prettyStr : PrettyHL a => (unicode : Bool) -> a -> String
prettyStr unicode =
  let layout = layoutSmart (MkLayoutOptions (AvailablePerLine 80 0.8)) in
  renderString . layout . pretty0 unicode


export
termHL : HL -> AnsiStyle
termHL Delim   = neutral
termHL TVar    = color BrightYellow
termHL TVarErr = color BrightYellow  <+> underline
termHL Dim     = color BrightGreen
termHL DVar    = color BrightGreen
termHL DVarErr = color BrightGreen   <+> underline
termHL Qty     = color BrightMagenta
termHL Free    = color BrightBlue
termHL Syntax  = color BrightCyan
termHL Tag     = color BrightRed

export %inline
prettyIO : PrettyOpts -> PrettyHL a => a -> IO Unit
prettyIO opts x =
  let reann = if opts.color then map termHL else unAnnotate in
  Terminal.putDoc $ reann $ pretty0 opts.unicode x

export %inline
prettyIODef : PrettyHL a => a -> IO Unit
prettyIODef = prettyIO defPrettyOpts


infixr 6 <%%>, <%>
export %inline
(<%%>) : Doc a -> Doc a -> Doc a
a <%%> b = sep [a, b]

export %inline
(<%>) : Doc a -> Doc a -> Doc a
a <%> b = cat [a, b]



||| wrapper for names that pretty-prints highlighted as a `TVar`.
public export data TVarName = TV BaseName
export %inline PrettyHL TVarName where prettyM (TV x) = hlF TVar $ prettyM x

||| wrapper for names that pretty-prints highlighted as a `DVar`.
public export data DVarName = DV BaseName
export %inline PrettyHL DVarName where prettyM (DV x) = hlF DVar $ prettyM x