module FP.Pretty.Pretty where

import FP.Prelude
import FP.Pretty.Color

-- # Format

data Format = 
    FG Color
  | BG Color
  | UL
  | BD
  deriving (Eq, Ord)

-- # PrettyEnv

data Layout = Flat | Break
  deriving (Eq,Ord)

data FailureMode = CanFail | CantFail
  deriving (Eq,Ord)

data PrettyParams = PrettyParams
  { punctuationFormat        ∷ [Format]
  , keywordPunctuationFormat ∷ [Format]
  , keywordFormat            ∷ [Format]
  , constructorFormat        ∷ [Format]
  , operatorFormat           ∷ [Format]
  , binderFormat             ∷ [Format]
  , literalFormat            ∷ [Format]
  , highlightFormat          ∷ [Format]
  , headerFormat             ∷ [Format]
  , errorFormat              ∷ [Format]
  , appLevel                 ∷ ℕ
  }
makeLenses ''PrettyParams

prettyParams₀ ∷ PrettyParams
prettyParams₀ = PrettyParams
  { punctuationFormat        = [FG darkGray]
  , keywordPunctuationFormat = [FG darkYellow,BD]
  , keywordFormat            = [FG darkYellow,BD,UL]
  , constructorFormat        = [FG darkGreen,BD]
  , operatorFormat           = [FG darkBlue]
  , binderFormat             = [FG darkTeal]
  , literalFormat            = [FG darkRed]
  , highlightFormat          = [BG highlight]
  , headerFormat             = [FG darkPink,BD,UL]
  , errorFormat              = [FG white,BG darkRed]
  , appLevel                 = 𝕟 100
  }

data PrettyEnv = PrettyEnv
  { prettyParams ∷ PrettyParams
  , maxColumnWidth ∷ ℕ
  , maxRibbonWidth ∷ ℕ
  , layout ∷ Layout
  , failureMode ∷ FailureMode
  , nesting ∷ ℕ
  , level ∷ ℕ
  , bumped ∷ 𝔹
  , undertagMode ∷ Maybe (ℂ,Color)
  , doOutput ∷ 𝔹
  , doFormat ∷ 𝔹
  , doLineNumbers ∷ 𝔹
  , lineNumberDisplayWidth ∷ ℕ
  , formats ∷ [Format]
  , blinders ∷ Maybe (ℕ,ℕ)
  }
makeLenses ''PrettyEnv

prettyEnv₀ ∷ PrettyEnv
prettyEnv₀ = PrettyEnv
  { prettyParams = prettyParams₀
  , maxColumnWidth = 𝕟 100
  , maxRibbonWidth = 𝕟 60
  , layout = Break
  , failureMode = CantFail
  , nesting = 𝕟 0
  , level = 𝕟 0
  , bumped = False
  , undertagMode = Nothing
  , doOutput = True
  , doFormat = True
  , doLineNumbers = False
  , lineNumberDisplayWidth = 𝕟 3
  , formats = []
  , blinders = Nothing
  }

-- # PrettyOut

data Chunk = Text 𝕊 | Newline
  deriving (Eq, Ord)
data PrettyOut = 
    ChunkOut Chunk 
  | FormatOut [Format] PrettyOut
  | NullOut 
  | AppendOut PrettyOut PrettyOut
  deriving (Eq, Ord)
instance Monoid PrettyOut where
  null = NullOut
  (⧺) = AppendOut

-- # PrettyState

data PrettyState = PrettyState
  { column ∷ ℕ
  , ribbon ∷ ℕ
  , beginningOfLine ∷ 𝔹
  , lineNumber ∷ ℕ
  , undertags ∷ [(ℕ,ℕ,ℂ,Color)]
  }
makeLenses ''PrettyState

prettyState₀ ∷ PrettyState
prettyState₀ = PrettyState
  { column = 𝕟 0
  , ribbon = 𝕟 0
  , beginningOfLine = True
  , lineNumber = 𝕟 0
  , undertags = []
  }

-- # PrettyM

newtype PrettyM a = PrettyM { runPrettyM ∷ RWST PrettyEnv PrettyOut PrettyState Maybe a }
  deriving
  ( Functor,Monad
  , MonadReader PrettyEnv
  , MonadWriter PrettyOut
  , MonadState PrettyState
  , MonadFailure
  )

runPrettyMWith ∷ PrettyEnv → PrettyState → PrettyM a → Maybe (a,PrettyOut,PrettyState)
runPrettyMWith r s aM = runRWSTWith r s $ runPrettyM aM

execOutPrettyMWith ∷ PrettyEnv → PrettyState → PrettyM a → Maybe PrettyOut
execOutPrettyMWith r s aM = do
  (_,o,_) ← runPrettyMWith r s aM
  return o

-- # Doc

newtype Doc = Doc { runDoc ∷ PrettyM () }
instance Eq Doc where
  (==) = (==) `on` (renderDoc ∘ ppFinal)
instance Ord Doc where
  compare = compare `on` (renderDoc ∘ ppFinal)
instance Monoid Doc where
  null = Doc $ return ()
  x ⧺ y = Doc $ runDoc x ≫ runDoc y

renderDoc ∷ Doc → PrettyOut
renderDoc aM =
  let errOut = FormatOut (errorFormat prettyParams₀) $ ChunkOut $ Text "<internal pretty printing error>"
  in ifNothing errOut $ execOutPrettyMWith prettyEnv₀ prettyState₀ $ runDoc aM

-- # Class

class Pretty a where
  pretty ∷ a → Doc

-- # Low-Level Interface

shouldOutputM ∷ PrettyM 𝔹
shouldOutputM = do
  ln ← getL lineNumberL
  bldrs ← askL blindersL
  outP ← askL doOutputL
  return $
    let inBlds = case bldrs of
          Nothing → True
          Just (low,high) → low ≤ ln ∧ ln ≤ high
    in outP ∧ inBlds

shouldOutputNewlineM ∷ PrettyM 𝔹
shouldOutputNewlineM = do
  so ← shouldOutputM
  ln ← getL lineNumberL
  bldrs ← askL blindersL
  uts ← getL undertagsL
  return $ so ∧ (case bldrs of {Nothing → True;Just (_,high) → ln < high} ∨ not (isEmpty uts))

  
-- ⟬s⟭ should contain no newlines
ppllSpit ∷ 𝕊 → PrettyM ()
ppllSpit s 
  | isEmpty s = return ()
  | otherwise = do
      fmtB ← askL doFormatL
      fmts ← askL formatsL
      let fmtF = if fmtB ∧ not (isEmpty fmts) then FormatOut fmts else id
      whenM shouldOutputM $ tell $ fmtF $ ChunkOut $ Text s
      modifyL columnL $ (+) $ length s
      modifyL ribbonL $ (+) $ countNonSpace s
      f ← askL $ failureModeL
      when (f == CanFail) $ do
        cmax ← askL $ maxColumnWidthL
        rmax ← askL $ maxRibbonWidthL
        c ← getL columnL
        r ← getL ribbonL
        when (c > cmax) abort
        when (r > rmax) abort
  where
    countNonSpace ∷ 𝕊 → ℕ
    countNonSpace = iter (\ c → if isSpace c then id else suc) (𝕟 0) ∘ stream

ppllFormat ∷ [Format] → PrettyM () → PrettyM ()
ppllFormat f = local (alter formatsL (f ⧺))

ppllNoFormat ∷ PrettyM () → PrettyM ()
ppllNoFormat = local (update doFormatL False)

ppllClearFormat ∷ PrettyM () → PrettyM ()
ppllClearFormat = local (update formatsL [])

ppllNewline ∷ PrettyM ()
ppllNewline = ppllNoFormat $ do
  whenM shouldOutputNewlineM $ tell $ ChunkOut Newline
  putL beginningOfLineL True
  putL columnL $ 𝕟 0
  putL ribbonL $ 𝕟 0

-- ⟬s⟭ should be non-empty and contain no newlines
ppllString ∷ 𝕊 → PrettyM ()
ppllString s = do
  ppllClearFormat $ whenM (getL beginningOfLineL) $ do
    whenM (askL doLineNumbersL) $ do
      ln ← getL lineNumberL
      w ← askL lineNumberDisplayWidthL
      ppllFormat [FG darkGray] $ ppllSpit $ alignRight w (𝕤 $ show ln) ⧺ ": "
    n ← askL nestingL 
    ppllSpit $ appendN n " "
    putL beginningOfLineL False
  col ← getL columnL
  ppllSpit s
  col' ← getL columnL
  whenM shouldOutputM $
    whenMaybeM (askL undertagModeL) $ \ (c,o) → do
      modifyL undertagsL $ (:) (col,col' - (col ⊓ col'),c,o)

ppllUndertags ∷ PrettyM ()
ppllUndertags = ppllClearFormat $ do
  uts ← reverse ^$ getL undertagsL
  when (not $ isEmpty uts) $ do
    ppllNewline
    foreachOn uts $ \ (utcol,len,c,o) → do
      col ← getL columnL
      let diff = utcol - (col ⊓ utcol)
      ppllSpit $ 𝕤 $ replicate diff ' '
      ppllFormat [FG o] $ ppllSpit $ 𝕤 $ replicate len c
  putL undertagsL []

ppllLineBreak ∷ PrettyM ()
ppllLineBreak = do
  ppllUndertags
  ppllNewline
  modifyL lineNumberL $ (+ 𝕟 1)

ppllText ∷ 𝕊 → PrettyM ()
ppllText s =
  let (s',snl) = prefixUntil (== '\n') $ list s
  in if not $ isEmpty s'
    then ppllString (𝕤 s') ≫ ppllText (𝕤 snl)
    else case uncons snl of
      Nothing → return ()
      Just ('\n',snl') → ppllLineBreak ≫ ppllText (𝕤 snl')
      Just _ → error $ "<internal error> ppText"

-- # Mid-Level Interface
  
ppFinal ∷ Doc → Doc
ppFinal d = Doc $ do
  runDoc d
  ppllUndertags

ppText ∷ 𝕊 → Doc
ppText = Doc ∘ ppllText

ppFormat ∷ [Format] → Doc → Doc
ppFormat f = Doc ∘ ppllFormat f ∘ runDoc

ppSpace ∷ ℕ → Doc
ppSpace n = ppText $ 𝕤 $ replicate n ' '

ppNewline ∷ Doc
ppNewline = ppText "\n"

ppIfFlat ∷ Doc → Doc → Doc
ppIfFlat flatAction breakAction = Doc $ do
  l ← askL $ layoutL
  runDoc $ case l of
    Flat → flatAction
    Break → breakAction

ppFlat ∷ Doc → Doc
ppFlat = Doc ∘ local (update layoutL Flat) ∘ runDoc

ppCanFail ∷ Doc → Doc
ppCanFail = Doc ∘ local (update failureModeL CanFail) ∘ runDoc

ppGroup ∷ Doc → Doc
ppGroup x  = ppIfFlat x $ Doc $ tries 
  [ runDoc $ ppFlat $ ppCanFail x
  , runDoc x
  ]

ppNest ∷ ℕ → Doc → Doc
ppNest n = Doc ∘ local (alter nestingL (+ n)) ∘ runDoc

ppAlign ∷ Doc → Doc
ppAlign aM = Doc $ do
  i ← askL $ nestingL
  c ← getL columnL
  runDoc $ ppNest (c - (i ⊓ c)) aM

ppLength ∷ Doc → ℕ
ppLength d = case runPrettyMWith prettyEnv₀ prettyState₀ $ runDoc d of
  Nothing → 𝕟 0
  Just ((),_,s) → column s

-- # Formatting Helpers

paramFormat ∷ (Lens PrettyParams [Format]) → 𝕊 → Doc
paramFormat l s = Doc $ do
  fmt ← askL $ l ⌾ prettyParamsL
  runDoc $ ppFormat fmt $ ppText s

ppNoFormat ∷ Doc → Doc
ppNoFormat = Doc ∘ local (update doFormatL False) ∘ runDoc

ppLineNumbers ∷ Doc → Doc
ppLineNumbers = Doc ∘ local (update doLineNumbersL True) ∘ runDoc

ppBlinders ∷ ℕ → ℕ → Doc → Doc
ppBlinders low high = Doc ∘ local (update blindersL $ Just (low,high)) ∘ runDoc

ppSetLineNumber ∷ ℕ → Doc → Doc
ppSetLineNumber n d = Doc $ do
  l ← getL lineNumberL
  putL lineNumberL n 
  runDoc d
  putL lineNumberL l

ppFG ∷ Color → Doc → Doc
ppFG c = ppFormat [FG c]

ppBG ∷ Color → Doc → Doc
ppBG c = ppFormat [BG c]

ppUL ∷ Doc → Doc
ppUL = ppFormat [UL]

ppBD ∷ Doc → Doc
ppBD = ppFormat [BD]

ppPun ∷ 𝕊 → Doc
ppPun = paramFormat punctuationFormatL

ppKeyPun ∷ 𝕊 → Doc
ppKeyPun = paramFormat keywordPunctuationFormatL

ppKey ∷ 𝕊 → Doc
ppKey = paramFormat keywordFormatL

ppCon ∷ 𝕊 → Doc
ppCon = paramFormat constructorFormatL

ppOp ∷ 𝕊 → Doc
ppOp = paramFormat operatorFormatL

ppBdr ∷ 𝕊 → Doc
ppBdr = paramFormat binderFormatL

ppLit ∷ 𝕊 → Doc
ppLit = paramFormat literalFormatL

ppHl ∷ 𝕊 → Doc
ppHl = paramFormat highlightFormatL

ppHeader ∷ 𝕊 → Doc
ppHeader = paramFormat headerFormatL

ppErr ∷ 𝕊 → Doc
ppErr = paramFormat errorFormatL

ppUT ∷ ℂ → Color → Doc → Doc
ppUT c o = Doc ∘ local (update undertagModeL $ Just (c,o)) ∘ runDoc

ppAlignLeft ∷ ℕ → Doc → Doc
ppAlignLeft n d = 
  let len = ppLength d
  in case n ⋚ len of
    LT → d
    EQ → d
    GT → d ⧺ ppSpace (n - (len ⊓ n))

ppAlignRight ∷ ℕ → Doc → Doc
ppAlignRight n d =
  let len = ppLength d
  in case n ⋚ len of
    LT → d
    EQ → d
    GT → ppSpace (n - (len ⊓ n)) ⧺ d

-- # High Level Helpers

ppHorizontal ∷ [Doc] → Doc
ppHorizontal = concat ∘ intersperse (ppSpace $ 𝕟 1) ∘ map ppAlign

ppVertical ∷ [Doc] → Doc
ppVertical = concat ∘ intersperse ppNewline ∘ map ppAlign

ppBreak ∷ Doc
ppBreak = ppIfFlat (ppSpace $ 𝕟 1) ppNewline

ppSeparated ∷ [Doc] → Doc
ppSeparated = ppGroup ∘ concat ∘ intersperse ppBreak ∘ map ppAlign

ppBotLevel ∷ Doc → Doc
ppBotLevel = Doc ∘ local (update levelL (𝕟 0) ∘ update bumpedL False) ∘ runDoc

ppClosed ∷ Doc → Doc → Doc → Doc
ppClosed alM arM aM = concat $ map ppAlign
  [ alM
  , ppBotLevel aM
  , arM
  ]

ppParens ∷ Doc → Doc
ppParens = ppClosed (ppPun "(") (ppPun ")")

ppAtLevel ∷ ℕ → Doc → Doc
ppAtLevel i' aM = Doc $ do
  i ← askL $ levelL
  b ← askL $ bumpedL
  if (i < i') ∨ ((i == i') ∧ not b)
    then local (update levelL i' ∘ update bumpedL False) $ runDoc aM
    else runDoc $ ppParens aM

ppBump ∷ Doc → Doc
ppBump = Doc ∘ local (update bumpedL True) ∘ runDoc

ppInf ∷ ℕ → Doc → Doc → Doc → Doc
ppInf i oM x1M x2M = ppGroup $ ppAtLevel i $ ppSeparated [ppBump x1M,oM,ppBump x2M]

ppInfl ∷ ℕ → Doc → Doc → Doc → Doc
ppInfl i oM x1M x2M = ppGroup $ ppAtLevel i $ ppSeparated [x1M,oM,ppBump x2M]

ppInfr ∷ ℕ → Doc → Doc → Doc → Doc
ppInfr i oM x1M x2M = ppGroup $ ppAtLevel i $ ppSeparated [ppBump x1M,oM,x2M]

ppPre ∷ ℕ → Doc → Doc → Doc
ppPre i oM xM = ppGroup $ ppAtLevel i $ ppSeparated [oM,xM]

ppPost ∷ ℕ → Doc → Doc → Doc
ppPost i oM xM = ppGroup $ ppAtLevel i $ ppSeparated [xM,oM]

ppApp ∷ Doc → [Doc] → Doc
ppApp x [] = x
ppApp x xs = ppGroup $ Doc $ do
  l ← askL $ appLevelL ⌾ prettyParamsL
  runDoc $ ppAtLevel l $ ppSeparated $ ppAtLevel l x : map (ppAtLevel l ∘ ppBump) xs

ppCollectionAtLevel ∷ ℕ → 𝕊 → 𝕊 → 𝕊 → [Doc] → Doc
ppCollectionAtLevel i open close sep xs = ppGroup $ ppBotLevel $ ppAtLevel i $ ppIfFlat flatCollection breakCollection
  where
    flatCollection = concat [ppPun open,concat $ intersperse (ppPun sep) xs,ppPun close]
    breakCollection = ppVertical $ concat
      [ mapHead (\ x → ppHorizontal [ppPun open,x]) $ mapTail (\ x → ppHorizontal [ppPun sep,x]) xs
      , return $ ppPun close
      ]

ppCollection ∷ 𝕊 → 𝕊 → 𝕊 → [Doc] → Doc
ppCollection = ppCollectionAtLevel $ 𝕟 0

ppRecord ∷ 𝕊 → [(Doc,Doc)] → Doc
ppRecord rel kvs = ppCollection "{" "}" "," $ map mapping kvs
  where
    mapping (k,v) = concat
      [ ppAlign k
      , ppIfFlat null (ppSpace (𝕟 1))
      , ppPun rel
      , ppIfFlat null (ppSpace (𝕟 1))
      , ppNest (𝕟 2) $ ppGroup $ concat
          [ ppIfFlat null ppNewline
          , ppAlign v
          ]
      ]

-- # NoFormat

renderChunk ∷ Chunk → 𝕊
renderChunk (Text s) = s
renderChunk Newline = "\n"

renderNoFormat ∷ PrettyOut → 𝕊
renderNoFormat (ChunkOut c) = renderChunk c
renderNoFormat (FormatOut _ o) = renderNoFormat o
renderNoFormat NullOut = ""
renderNoFormat (AppendOut o₁ o₂) = renderNoFormat o₁ ⧺ renderNoFormat o₂

ppString ∷ (Pretty a) ⇒ a → 𝕊
ppString = renderNoFormat ∘ renderDoc ∘ ppFinal ∘ pretty