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 public Control.Monad.Identity
import public Control.Monad.Reader

%default total


public export
data HL
  = Delim
  | TVar
  | TVarErr
  | Dim
  | DVar
  | DVarErr
  | Qty
  | Free
  | Syntax

private HLRepr : Type
HLRepr = Nat

private %inline hlRepr : HL -> Nat
hlRepr Delim   = 0
hlRepr TVar    = 1
hlRepr TVarErr = 2
hlRepr Dim     = 3
hlRepr DVar    = 4
hlRepr DVarErr = 5
hlRepr Qty     = 6
hlRepr Free    = 7
hlRepr Syntax  = 8

export Eq  HL where (==) = (==) `on` hlRepr
export Ord HL where compare = compare `on` hlRepr


public export
data PPrec
  = Outer
  | Ann  -- right of "::"
  | AnnL -- left  of "::"
  -- ...
  | App  -- term/dimension application
  | SApp -- substitution application
  | Arg

private PrecRepr : Type
PrecRepr = Nat

private %inline precRepr : PPrec -> PrecRepr
precRepr Outer = 0
precRepr Ann   = 1
precRepr AnnL  = 2
-- ...
precRepr App   = 98
precRepr SApp  = 99
precRepr Arg   = 100

export Eq  PPrec where (==) = (==) `on` precRepr
export Ord PPrec where compare = compare `on` precRepr


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
parens : Doc HL -> Doc HL
parens doc = hl Delim "(" <+> doc <+> hl Delim ")"

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


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


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

public export %inline 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
under : HasEnv m => BinderSort -> Name -> m a -> m a
under s x = local $
  {prec := Outer} .
  (case s of T => {tnames $= (x ::)}; D => {dnames $= (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
pretty0 : PrettyHL a => {default True unicode : Bool} -> a -> Doc HL
pretty0 x {unicode} =
  let env = MakePrettyEnv {dnames = [], tnames = [], unicode, prec = Outer} in
  runReader env $ prettyM x


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 %inline
prettyStr : PrettyHL a => {default True 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   = color BrightBlack
termHL TVar    = color BrightYellow
termHL TVarErr = color BrightYellow  <+> underline
termHL Dim     = color BrightGreen   <+> bold
termHL DVar    = color BrightGreen
termHL DVarErr = color BrightGreen   <+> underline
termHL Qty     = color BrightMagenta <+> bold
termHL Free    = color BrightWhite
termHL Syntax  = color BrightBlue

export %inline
prettyTerm : {default True color, unicode : Bool} -> PrettyHL a => a -> IO Unit
prettyTerm x {color, unicode} =
  let reann = if color then map termHL else unAnnotate in
  Terminal.putDoc $ reann $ pretty0 x {unicode}


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

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