module Quox.Pretty

import Quox.Loc
import Quox.Name

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 | TVarErr
| Dim  | DVar | DVarErr
| Qty
| Syntax
| Tag
%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 =
  extract $
  evalStateAt PREC prec $
  runReaderAt FLAVOR flavor $
  runReaderAt HIGHLIGHT highlight $
  runReaderAt INDENT indent act


export %inline
toSGR : HL -> List SGR
toSGR Delim   = []
toSGR TVar    = [SetForeground BrightYellow]
toSGR TVarErr = [SetForeground BrightYellow, SetStyle SingleUnderline]
toSGR Dim     = [SetForeground BrightGreen]
toSGR DVar    = [SetForeground BrightGreen]
toSGR DVarErr = [SetForeground BrightGreen, SetStyle SingleUnderline]
toSGR Qty     = [SetForeground BrightMagenta]
toSGR Free    = [SetForeground BrightBlue]
toSGR Syntax  = [SetForeground BrightCyan]
toSGR Tag     = [SetForeground BrightRed]

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


export %inline
runPretty : Eff Pretty a -> a
runPretty = runPrettyWith Outer Unicode noHighlight 2

export %inline
runPrettyColor : Eff Pretty a -> a
runPrettyColor = runPrettyWith Outer Unicode highlightSGR 2


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


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

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

export %inline
hangDSingle : {opts : _} -> Doc opts -> Doc opts -> Eff Pretty (Doc opts)
hangDSingle d1 d2 =
  pure $ ifMultiline (d1 <++> d2) (vappend d1 !(indentD d2))


export
tightDelims : {opts : _} -> (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 : _} -> (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 : _} -> Doc opts -> Eff Pretty (Doc opts)
parens = tightDelims "(" ")"

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

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

export %inline
parensIf : {opts : _} -> 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 : _} -> List (Doc opts) -> Doc opts
sepSingle xs = ifMultiline (hsep xs) (vsep xs)

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

export
exceptLast : {opts : _} -> (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
  fillSeparateTight : Doc opts -> t (Doc opts) -> Doc opts
  fillSeparateTight d = fillSep . exceptLast (<+> d) . toList


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

export %inline
parensIfM : {opts : _} -> 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
prettyFree : {opts : _} -> Name -> Eff Pretty (Doc opts)
prettyFree = hl Free . text . toDots

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

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

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



export %inline
typeD, arrowD, darrowD, timesD, lamD, eqndD, dlamD, annD, natD,
eqD, colonD, commaD, semiD, caseD, typecaseD, returnD,
ofD, dotD, zeroD, succD, coeD, compD, undD, cstD, pipeD :
  {opts : _} -> Eff Pretty (Doc opts)
typeD     = hl Syntax . text =<< ifUnicode "★" "Type"
arrowD    = hl Delim  . text =<< ifUnicode "→" "->"
darrowD   = hl Delim  . text =<< ifUnicode "⇒" "=>"
timesD    = hl Delim  . text =<< ifUnicode "×" "**"
lamD      = hl Syntax . text =<< ifUnicode "λ" "fun"
eqndD     = hl Delim  . text =<< ifUnicode "≡" "=="
dlamD     = hl Syntax . text =<< ifUnicode "δ" "dfun"
annD      = hl Delim  . text =<< ifUnicode "∷" "::"
natD      = hl Syntax . text =<< ifUnicode "ℕ" "Nat"
eqD       = hl Syntax $ text "Eq"
colonD    = hl Delim  $ text ":"
commaD    = hl Delim  $ text ","
semiD     = 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 Syntax $ text "zero"
succD     = hl Syntax $ text "succ"
coeD      = hl Syntax $ text "coe"
compD     = hl Syntax $ text "comp"
undD      = hl Syntax $ text "_"
cstD      = hl Syntax $ text "="
pipeD     = hl Syntax $ text "|"


export
prettyApp : {opts : _} -> Nat -> Doc opts -> List (Doc opts) -> Doc opts
prettyApp ind f args =
      hsep (f :: args)
  <|> hsep [f, vsep args]
  <|> vsep (f :: map (indent ind) args)

export
prettyAppD : {opts : _} -> Doc opts -> List (Doc opts) -> Eff Pretty (Doc opts)
prettyAppD f args = pure $ 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 : _} -> 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 : _} -> Loc -> Eff Pretty (Doc opts)
prettyLoc (L NoLoc) =
  hcat <$> sequence [hl TVarErr "no location", colonD]
prettyLoc (L (YesLoc file b)) =
  hcat <$> sequence [hl Free $ text file, colonD, prettyBounds b]