module Quox.Pretty

import Quox.Loc
import Quox.Name

import Control.Monad.ST.Extra
import public Text.PrettyPrint.Bernardy
import public Text.PrettyPrint.Bernardy.Core.Decorate
import public Quox.EffExtra
import public Data.String
import Control.ANSI.SGR

import Data.DPair
import Data.SnocList
import Derive.Prelude

%default total
%language ElabReflection

%hide TT.Name


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


public export
data HL
= Delim
| Free | TVar
| Dim  | DVar
| Qty  | QVar
| Universe
| Syntax
| Constant
%runElab derive "HL" [Eq, Ord, Show]


public export
data Flavor = Unicode | Ascii
%runElab derive "Flavor" [Eq, Ord, Show]

export %inline
noHighlight : HL -> Highlight
noHighlight _ = MkHighlight "" ""


public export
data EffTag = PREC | FLAVOR | HIGHLIGHT | INDENT

public export
Pretty : List (Type -> Type)
Pretty = [StateL PREC PPrec, ReaderL FLAVOR Flavor,
          ReaderL HIGHLIGHT (HL -> Highlight), ReaderL INDENT Nat]

export %inline
runPrettyWith : PPrec -> Flavor -> (HL -> Highlight) -> Nat ->
                Eff Pretty a -> a
runPrettyWith prec flavor highlight indent act =
  runST $ do
    runEff act $ with Union.(::)
      [handleStateSTRef !(newSTRef prec),
       handleReaderConst flavor,
       handleReaderConst highlight,
       handleReaderConst indent]


export %inline
toSGR : HL -> List SGR
toSGR Delim    = []
toSGR Free     = [SetForeground BrightBlue]
toSGR TVar     = [SetForeground BrightYellow]
toSGR Dim      = [SetForeground BrightGreen]
toSGR DVar     = [SetForeground BrightGreen]
toSGR Qty      = [SetForeground BrightMagenta]
toSGR QVar     = [SetForeground BrightMagenta]
toSGR Universe = [SetForeground BrightRed]
toSGR Syntax   = [SetForeground BrightCyan]
toSGR Constant = [SetForeground BrightRed]

export %inline
highlightSGR : HL -> Highlight
highlightSGR h = MkHighlight (escapeSGR $ toSGR h) (escapeSGR [Reset])

export %inline
toClass : HL -> String
toClass Delim    = "dl"
toClass Free     = "fr"
toClass TVar     = "tv"
toClass Dim      = "dc"
toClass DVar     = "dv"
toClass Qty      = "qt"
toClass QVar     = "qv"
toClass Universe = "un"
toClass Syntax   = "sy"
toClass Constant = "co"

export %inline
highlightHtml : HL -> Highlight
highlightHtml h = MkHighlight #"<span class="\#{toClass h}">"# "</span>"


export %inline
runPrettyHL : (HL -> Highlight) -> Eff Pretty a -> a
runPrettyHL f = runPrettyWith Outer Unicode f 2

export %inline
runPretty : Eff Pretty a -> a
runPretty = runPrettyHL noHighlight


export %inline
hl : {opts : LayoutOpts} -> HL -> Doc opts -> Eff Pretty (Doc opts)
hl h doc = asksAt HIGHLIGHT $ \f => decorate (f h) doc


export %inline
indentD : {opts : LayoutOpts} -> Doc opts -> Eff Pretty (Doc opts)
indentD doc = pure $ indent !(askAt INDENT) doc

export %inline
hangD : {opts : LayoutOpts} -> Doc opts -> Doc opts -> Eff Pretty (Doc opts)
hangD d1 d2 = pure $ hangSep !(askAt INDENT) d1 d2

export %inline
hangSingle : {opts : LayoutOpts} -> Nat -> Doc opts -> Doc opts -> Doc opts
hangSingle n d1 d2 = ifMultiline (d1 <++> d2) (vappend d1 (indent n d2))

export %inline
hangDSingle : {opts : LayoutOpts} -> Doc opts -> Doc opts ->
              Eff Pretty (Doc opts)
hangDSingle d1 d2 = pure $ hangSingle !(askAt INDENT) d1 d2


export
tightDelims : {opts : LayoutOpts} -> (l, r : String) -> (inner : Doc opts) ->
              Eff Pretty (Doc opts)
tightDelims l r inner = do
  l <- hl Delim $ text l
  r <- hl Delim $ text r
  pure $ hcat [l, inner, r]

export
looseDelims : {opts : LayoutOpts} -> (l, r : String) -> (inner : Doc opts) ->
              Eff Pretty (Doc opts)
looseDelims l r inner = do
  l <- hl Delim $ text l
  r <- hl Delim $ text r
  let short = hsep [l, inner, r]
      long  = vsep [l, !(indentD inner), r]
  pure $ ifMultiline short long

export %inline
parens : {opts : LayoutOpts} -> Doc opts -> Eff Pretty (Doc opts)
parens = tightDelims "(" ")"

export %inline
bracks : {opts : LayoutOpts} -> Doc opts -> Eff Pretty (Doc opts)
bracks = tightDelims "[" "]"

export %inline
braces : {opts : LayoutOpts} -> Doc opts -> Eff Pretty (Doc opts)
braces = looseDelims "{" "}"

export %inline
tightBraces : {opts : LayoutOpts} -> Doc opts -> Eff Pretty (Doc opts)
tightBraces = tightDelims "{" "}"

export %inline
parensIf : {opts : LayoutOpts} -> Bool -> Doc opts -> Eff Pretty (Doc opts)
parensIf True  = parens
parensIf False = pure


||| uses hsep only if the whole list fits on one line
export
sepSingle : {opts : LayoutOpts} -> List (Doc opts) -> Doc opts
sepSingle xs = ifMultiline (hsep xs) (vsep xs)

export
fillSep : {opts : LayoutOpts} -> List (Doc opts) -> Doc opts
fillSep []        = empty
fillSep (x :: xs) = foldl (\x, y => sep [x, y]) x xs

export
exceptLast : {opts : LayoutOpts} -> (Doc opts -> Doc opts) ->
             List (Doc opts) -> List (Doc opts)
exceptLast f []        = []
exceptLast f [x]       = [x]
exceptLast f (x :: xs) = f x :: exceptLast f xs

parameters {opts : LayoutOpts} {auto _ : Foldable t}
  export
  separateLoose : Doc opts -> t (Doc opts) -> Doc opts
  separateLoose d = sep . exceptLast (<++> d) . toList

  export
  separateTight : Doc opts -> t (Doc opts) -> Doc opts
  separateTight d = sep . exceptLast (<+> d) . toList

  export
  sepSingleTight : Doc opts -> t (Doc opts) -> Doc opts
  sepSingleTight d = sepSingle . exceptLast (<+> d) . toList

  export
  hseparateTight : Doc opts -> t (Doc opts) -> Doc opts
  hseparateTight d = hsep . exceptLast (<+> d) . toList

  export
  vseparateTight : Doc opts -> t (Doc opts) -> Doc opts
  vseparateTight d = vsep . exceptLast (<+> d) . toList

  export
  fillSeparateTight : Doc opts -> t (Doc opts) -> Doc opts
  fillSeparateTight d = fillSep . exceptLast (<+> d) . toList


export %inline
pshow : {opts : LayoutOpts} -> Show a => a -> Doc opts
pshow = text . show


export %inline
ifUnicode : (uni, asc : Lazy a) -> Eff Pretty a
ifUnicode uni asc =
  asksAt FLAVOR $ \case
    Unicode => uni
    Ascii   => asc

export %inline
parensIfM : {opts : LayoutOpts} -> PPrec -> Doc opts -> Eff Pretty (Doc opts)
parensIfM d doc = parensIf (!(getAt PREC) > d) doc

export %inline
withPrec : PPrec -> Eff Pretty a -> Eff Pretty a
withPrec = localAt_ PREC

export %inline
qbrackets : {opts : LayoutOpts} -> Doc opts -> Eff Pretty (Doc opts)
qbrackets doc = tightDelims !(ifUnicode "‹" ".{") !(ifUnicode "›" "}") doc


export
prettyName : Name -> Doc opts
prettyName = text . toDots

export
prettyFree : {opts : LayoutOpts} -> Name -> Eff Pretty (Doc opts)
prettyFree = hl Free . prettyName

export
prettyBind' : BindName -> Doc opts
prettyBind' = text . baseStr . val

export
prettyTBind : {opts : LayoutOpts} -> BindName -> Eff Pretty (Doc opts)
prettyTBind = hl TVar . prettyBind'

export
prettyDBind : {opts : LayoutOpts} -> BindName -> Eff Pretty (Doc opts)
prettyDBind = hl DVar . prettyBind'

export
prettyQBind : {opts : LayoutOpts} -> BindName -> Eff Pretty (Doc opts)
prettyQBind = hl QVar . prettyBind'


export %inline
typeD, ioStateD, arrowD, darrowD, timesD, lamD, eqndD, dlamD, annD, natD,
stringD, eqD, colonD, commaD, semiD, atD, caseD, typecaseD, returnD, ofD, dotD,
zeroD, succD, coeD, compD, undD, cstD, pipeD, fstD, sndD, letD, inD,
qplusD, qtimesD :
  {opts : LayoutOpts} -> Eff Pretty (Doc opts)
typeD     = hl Syntax . text =<< ifUnicode "★" "Type"
ioStateD  = hl Syntax $ text "IOState"
arrowD    = hl Syntax . text =<< ifUnicode "→" "->"
darrowD   = hl Syntax . text =<< ifUnicode "⇒" "=>"
timesD    = hl Syntax . text =<< ifUnicode "×" "**"
lamD      = hl Syntax . text =<< ifUnicode "λ" "fun"
eqndD     = hl Syntax . text =<< ifUnicode "≡" "=="
dlamD     = hl Syntax . text =<< ifUnicode "δ" "dfun"
annD      = hl Syntax . text =<< ifUnicode "∷" "::"
natD      = hl Syntax . text =<< ifUnicode "ℕ" "Nat"
stringD   = hl Syntax   $ text "String"
eqD       = hl Syntax   $ text "Eq"
colonD    = hl Syntax   $ text ":"
commaD    = hl Syntax   $ text ","
semiD     = hl Delim    $ text ";"
atD       = hl Delim    $ text "@"
caseD     = hl Syntax   $ text "case"
typecaseD = hl Syntax   $ text "type-case"
ofD       = hl Syntax   $ text "of"
returnD   = hl Syntax   $ text "return"
dotD      = hl Delim    $ text "."
zeroD     = hl Constant $ text "zero"
succD     = hl Constant $ text "succ"
coeD      = hl Syntax   $ text "coe"
compD     = hl Syntax   $ text "comp"
undD      = hl Syntax   $ text "_"
cstD      = hl Syntax   $ text "="
pipeD     = hl Delim    $ text "|"
fstD      = hl Syntax   $ text "fst"
sndD      = hl Syntax   $ text "snd"
letD      = hl Syntax   $ text "let"
inD       = hl Syntax   $ text "in"
qplusD    = hl Syntax   $ text "+"
qtimesD   = hl Syntax . text =<< ifUnicode "·" "*"


export
prettyApp : {opts : LayoutOpts} -> Nat -> Doc opts ->
            List (Doc opts) -> Doc opts
prettyApp ind f args =
  ifMultiline
    (hsep (f :: args))
    (f <++> vsep args <|> vsep (f :: map (indent ind) args))

export
prettyAppD : {opts : LayoutOpts} -> Doc opts -> List (Doc opts) ->
             Eff Pretty (Doc opts)
prettyAppD f args = parensIfM App $ prettyApp !(askAt INDENT) f args


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

export
quoteTag : String -> String
quoteTag tag =
  if isName tag then tag else
  "\"" ++ escapeString tag ++ "\""

export
prettyBounds : {opts : LayoutOpts} -> Bounds -> Eff Pretty (Doc opts)
prettyBounds (MkBounds l1 c1 l2 c2) =
  hcat <$> sequence
    [hl TVar $ text $ show l1, colonD,
     hl DVar $ text $ show c1, hl Delim "-",
     hl TVar $ text $ show l2, colonD,
     hl DVar $ text $ show c2, colonD]

export
prettyLoc : {opts : LayoutOpts} -> Loc -> Eff Pretty (Doc opts)
prettyLoc (L NoLoc) =
  hcat <$> sequence [hl TVar "no location", colonD]
prettyLoc (L (YesLoc file b)) =
  hcat <$> sequence [hl Free $ text file, colonD, prettyBounds b]

export
prettyTag : {opts : _} -> String -> Eff Pretty (Doc opts)
prettyTag tag = hl Constant $ text $ "'" ++ quoteTag tag

export
prettyStrLit : {opts : _} -> String -> Eff Pretty (Doc opts)
prettyStrLit s =
  let s = concatMap esc1 $ unpack s in
  hl Constant $ hcat ["\"", text s, "\""]
where
  esc1 : Char -> String
  esc1 '"' = "\""; esc1 '\\' = "\\"
  esc1 c   = singleton c