-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

-- | Utilities for rendering Michelson code in a format compatible with
-- Tezos software (e.g @tezos-client@)
module Morley.Michelson.Printer.Util
  ( RenderDoc(..)
  , Prettier(..)
  , printDoc
  , printDocB
  , printDocS
  , renderOps
  , renderOpsList
  , renderOpsListNoBraces
  , renderAnyBuildable
  , spaces
  , wrapInParens
  , buildRenderDoc
  , buildRenderDocExtended
  , renderDocList

    -- * Smart parentheses
  , RenderContext
  , needsParens
  , doesntNeedParens
  , addParens
  , addParensMultiline
  , assertParensNotNeeded
  ) where

import Prelude hiding (group)

import Control.Exception (assert)
import Data.Text.Lazy qualified as LT
import Data.Text.Lazy.Builder (Builder)
import Fmt (Buildable, pretty)
import Text.PrettyPrint.Leijen.Text
  (Doc, SimpleDoc, align, braces, displayB, displayT, enclose, encloseSep, hcat, hsep, isEmpty,
  lbracket, parens, punctuate, rbracket, renderOneLine, renderPretty, semi, sep, space, text, (<+>),
  (<//>))

-- | Environment carried during recursive rendering.
newtype RenderContext = RenderContext
  { _rcWillNeedParens :: Bool
    -- ^ Whether the current expression is going to be used as part of
    -- top-level expression or in a similar context.
    -- When set to 'True', you may need to wrap your rendered expression into
    -- parentheses.
  }

-- | Generalize converting a type into a
-- Text.PrettyPrint.Leijen.Text.Doc. Used to pretty print Michelson code
-- and define Fmt.Buildable instances.
class RenderDoc a where
  renderDoc :: RenderContext -> a -> Doc

  -- | Whether a value can be represented in Michelson code.
  -- Normally either all values of some type are renderable or not renderable.
  -- However, in case of instructions we have extra instructions which should
  -- not be rendered.
  -- Note: it's not sufficient to just return 'mempty' for such instructions,
  -- because sometimes we want to print lists of instructions and we need to
  -- ignore them complete (to avoid putting redundant separators).
  isRenderable :: a -> Bool
  isRenderable _ = True

-- | Renders a list of 'RenderDoc' elements surrounded with square brackets,
-- separated by a comma and a space.
renderDocList :: RenderDoc a => RenderContext -> [a] -> Doc
renderDocList context = encloseSep lbracket rbracket ", " . fmap (renderDoc context)

renderAnyBuildable :: Buildable a => a -> Doc
renderAnyBuildable = text . pretty

-- | A new type that can wrap values so that the RenderDoc
-- instances of the combined value can have a different
-- behavior for the pretty printer.
newtype Prettier a = Prettier a
  deriving stock Functor

-- | Convert 'Doc' to 'Text' with a line width of 80.
printDoc :: Bool -> Doc -> LT.Text
printDoc oneLine = displayT . doRender oneLine

-- | Convert 'Doc' to 'Builder' in the same manner as 'printDoc'.
printDocB :: Bool -> Doc -> Builder
printDocB oneLine = displayB . doRender oneLine

-- | Convert 'Doc' to 'String' in the same manner as 'printDoc'.
printDocS :: Bool -> Doc -> String
printDocS oneLine = toString . printDoc oneLine

-- | Generic way to render the different op types that get passed
-- to a contract.
renderOps :: (RenderDoc op) => Bool -> NonEmpty op -> Doc
renderOps oneLine = renderOpsList oneLine . toList

-- | Render a comma-separated list of items in braces
renderOpsList :: (RenderDoc op) => Bool -> [op] -> Doc
renderOpsList oneLine ops =
  braces $ enclose space space $ renderOpsListNoBraces oneLine ops

-- | Render a semi-colon-separated list of items without braces
renderOpsListNoBraces :: RenderDoc op => Bool -> [op] -> Doc
renderOpsListNoBraces oneLine =
  align . (if oneLine then hsep else sep) . punctuate semi .
    fmap (renderDoc doesntNeedParens) . filter isRenderable

-- | Create a specific number of spaces.
spaces :: Int -> Doc
spaces x = hcat $ replicate x space

-- | Wrap documents in parentheses if there are two or more in the list.
wrapInParens :: RenderContext -> NonEmpty Doc -> Doc
wrapInParens (RenderContext pn) ds =
  let moreThanOne = length (filter (not . isEmpty) (toList ds)) > 1
  in  addParens (RenderContext (pn && moreThanOne)) $
        foldr (<+>) mempty ds

-- | Turn something that is instance of `RenderDoc` into a `Builder`.
-- It's formatted the same way as `printDoc` formats docs.
buildRenderDoc :: RenderDoc a => a -> Builder
buildRenderDoc = printDocB True . renderDoc doesntNeedParens

-- | Works as 'buildRenderDoc' above, but doesn't force the doc to be printed in one line
buildRenderDocExtended :: RenderDoc a => a -> Builder
buildRenderDocExtended = printDocB False . renderDoc doesntNeedParens

-- | Here using a page width of 80 and a ribbon width of 1.0
-- https://hackage.haskell.org/package/wl-pprint-1.2.1/docs/Text-PrettyPrint-Leijen.html
doRender :: Bool -> Doc -> SimpleDoc
doRender oneLine = if oneLine then renderOneLine else renderPretty 1.0 80

-- Smart parentheses
----------------------------------------------------------------------------

{- Motivation:

Some expressions may need to be wrapped into parentheses, but only if they
are part of other expression, and are not already wrapped into braces or
brackets.
-}

-- | Constructors for 'RenderContext'
needsParens, doesntNeedParens :: RenderContext
needsParens = RenderContext True
doesntNeedParens = RenderContext False

-- | Add parentheses if needed.
addParens :: RenderContext -> Doc -> Doc
addParens = \case
  RenderContext True -> parens
  RenderContext False -> id

-- | Add parentheses if needed, multiline if necessary.
addParensMultiline :: RenderContext -> Doc -> Doc
addParensMultiline pn doc = case pn of
  RenderContext True -> "(" <> doc <//> ")"
  RenderContext False -> doc

-- | Ensure parentheses are not required, for case when you cannot
-- sensibly wrap your expression into them.
assertParensNotNeeded :: RenderContext -> a -> a
assertParensNotNeeded (RenderContext pn) = assert (not pn)