{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE LambdaCase #-}
{-# OPTIONS_GHC -Wno-name-shadowing #-}

module Text.PrettyPrint.Avh4.Block
  ( -- * Line
    Line,
    space,
    string7,
    char7,
    stringUtf8,
    lineFromBuilder,
    commentByteString,

    -- * Block
    Block,

    -- ** convert to
    render,

    -- ** create
    blankLine,
    line,
    mustBreak,

    -- ** combine
    stack,
    indent,
    prefix,
    addSuffix,
    prefixOrIndent,
    rowOrStack,
    rowOrStackForce,
    rowOrIndent,
    rowOrIndentForce,
    spaceSeparatedOrStack,
    spaceSeparatedOrStackForce,
    spaceSeparatedOrIndent,
    spaceSeparatedOrIndentForce,

    -- *** deprecated
    stackForce,
    andThen,
    joinMustBreak,
  )
where

import Control.Applicative (Applicative (..))
import Data.ByteString (ByteString)
import qualified Data.ByteString as ByteString
import qualified Data.ByteString.Builder as B
import Data.List.NonEmpty (NonEmpty ((:|)))
import qualified Data.List.NonEmpty as NonEmpty
import Data.Semigroup (sconcat)
import Text.PrettyPrint.Avh4.Indent (Indent)
import qualified Text.PrettyPrint.Avh4.Indent as Indent

-- | A `Line` is ALWAYS just one single line of text,
-- and can always be combined horizontally with other `Line`s.
--
-- - `Space` is a single horizontal space,
-- - `Blank` is a line with no content.
-- - `Text` brings any text into the data structure. (Uses `ByteString.Builder` for the possibility of optimal performance)
-- - `Row` joins multiple elements onto one line.
data Line
  = Text B.Builder
  | Row Line Line
  | Space
  | Blank
  deriving (Show)

instance Semigroup Line where
  a <> b = Row a b

-- | Creates a @Line@ from the given @Char@.
-- You must guarantee that the given character is a valid 7-bit ASCII character,
-- and is not a space character (use `space` instead).
char7 :: Char -> Line
char7 = Text . B.char7

-- | Creates a @Line@ from the given @String@.
-- You must guarantee that all characters in the @String@ are valid 7-bit ASCII characters,
-- and that the string does not start or end with spaces (use `space` instead).
string7 :: String -> Line
string7 = Text . B.string7

-- | If you know the String only contains ASCII characters, then use `string7` instead for better performance.
stringUtf8 :: String -> Line
stringUtf8 = Text . B.stringUtf8

-- | You must guarantee that the content of the Builder does not contain newlines and does not start with whitespace.
lineFromBuilder :: B.Builder -> Line
lineFromBuilder = Text

{-# INLINE mkTextByteString #-}
mkTextByteString :: ByteString -> Line
mkTextByteString = Text . B.byteString

commentByteString :: ByteString -> Line
commentByteString bs =
  if ByteString.null bs
    then Blank
    else mkTextByteString bs

-- | A @Line@ containing a single space.  You **must** use this to create
-- space characters if the spaces will ever be at the start or end of a line that is joined in the context of indentation changes.
space :: Line
space =
  Space

data Indented a
  = Indented Indent a
  deriving (Functor, Show)

-- | `Block` contains Lines (at least one; it can't be empty).
--
-- Block either:
--
--  - can appear in the middle of a line
--      (Stack someLine [], thus can be joined without problems), or
--  - has to appear on its own
--      (Stack someLine moreLines OR MustBreak someLine).
--
-- Types of Blocks:
--
-- - `SingleLine` is a single line, and the indentation level for the line.
-- - `MustBreak` is a single line (and its indentation level)) that cannot have anything joined to its right side.
--  Notably, it is used for `--` comments.
-- - `Stack` contains two or more lines, and the indentation level for each.
--
-- Sometimes (see `prefix`) the first line of Stack
--  gets different treatment than the other lines.
data Block
  = SingleLine RequiredLineBreaks (Indented Line)
  | Stack (Indented Line) (NonEmpty (Indented Line))
  deriving (Show)

data RequiredLineBreaks
  = MustBreakAtEnd
  | NoRequiredBreaks
  deriving (Show, Eq)

-- | A blank line (taking up one vertical space), with no text content.
blankLine :: Block
blankLine =
  line Blank

-- | Promote a @Line@ into a @Block@.
line :: Line -> Block
line =
  SingleLine NoRequiredBreaks . mkIndentedLine

-- | Promote a @Line@ into a @Block@ that will always have a newline at the end of it,
-- meaning that this @Line@ will never have another @Line@ joined to its right side.
mustBreak :: Line -> Block
mustBreak =
  SingleLine MustBreakAtEnd . mkIndentedLine

mkIndentedLine :: Line -> Indented Line
mkIndentedLine Space = Indented (Indent.spaces 1) Blank
mkIndentedLine (Row Space next) =
  let (Indented i rest') = mkIndentedLine next
   in Indented (Indent.spaces 1 <> i) rest'
mkIndentedLine other = Indented mempty other

-- | A binary version of `stack`.
{-# DEPRECATED stackForce "Use `stack` instead." #-}
stackForce :: Block -> Block -> Block
stackForce b1 b2 =
  let (line1first :| line1rest) = toLines b1
   in Stack
        line1first
        (line1rest `NonEmpty.prependList` toLines b2)
  where
    toLines :: Block -> NonEmpty (Indented Line)
    toLines b =
      case b of
        SingleLine _ l1 ->
          -- We lose information about RequiredLineBreaks, but that's okay
          -- since the result will always be a multiline stack,
          -- which will never join with a single line
          pure l1
        Stack l1 rest ->
          NonEmpty.cons l1 rest

-- | An alternate style for `stack`.
--
-- @a & andThen [b, c]@ is the same as @stack [a, b, c]@
{-# DEPRECATED andThen "Use `stack` instead." #-}
andThen :: [Block] -> Block -> Block
andThen rest first =
  foldl stackForce first rest

-- | A vertical stack of @Block@s.  The left edges of all the @Block@s will be aligned.
stack :: NonEmpty Block -> Block
stack = foldr1 stackForce

{-# DEPRECATED joinMustBreak "The normal rowOr* functions should handle this automatically now.  Use `rowOrStack (Just space)` instead." #-}
joinMustBreak :: Block -> Block -> Block
joinMustBreak inner eol =
  case (inner, eol) of
    (SingleLine NoRequiredBreaks (Indented i1 inner'), SingleLine NoRequiredBreaks (Indented _ eol')) ->
      SingleLine NoRequiredBreaks $
        Indented i1 $
          inner' <> space <> eol'
    (SingleLine NoRequiredBreaks (Indented i1 inner'), SingleLine MustBreakAtEnd (Indented _ eol')) ->
      SingleLine MustBreakAtEnd $
        Indented i1 $
          inner' <> space <> eol'
    _ ->
      stackForce inner eol

{-# INLINE prefixOrIndent #-}
prefixOrIndent :: Maybe Line -> Line -> Block -> Block
prefixOrIndent joiner a b =
  let join a b =
        case joiner of
          Nothing -> a <> b
          Just j -> a <> j <> b
   in case b of
        SingleLine breaks (Indented _ b') ->
          SingleLine breaks $ mkIndentedLine $ join a b'
        _ ->
          stackForce (line a) (indent b)

mapLines :: (Indented Line -> Indented Line) -> Block -> Block
mapLines fn =
  mapFirstLine fn fn

mapFirstLine :: (Indented Line -> Indented Line) -> (Indented Line -> Indented Line) -> Block -> Block
mapFirstLine firstFn restFn b =
  case b of
    SingleLine breaks l1 ->
      SingleLine breaks (firstFn l1)
    Stack l1 ls ->
      Stack (firstFn l1) (fmap restFn ls)

mapLastLine :: (Indented Line -> Indented Line) -> Block -> Block
mapLastLine lastFn = \case
  SingleLine breaks l1 ->
    SingleLine breaks (lastFn l1)
  Stack l1 ls ->
    Stack l1 (NonEmpty.init ls `NonEmpty.prependList` pure (lastFn $ NonEmpty.last ls))

-- | Makes a new @Block@ with the contents of the input @Block@ indented by one additional level.
indent :: Block -> Block
indent =
  mapLines (\(Indented i l) -> Indented (Indent.tab <> i) l)

-- | This is the same as `rowOrStackForce` @False@.
{-# INLINE rowOrStack #-}
rowOrStack :: Maybe Line -> NonEmpty Block -> Block
rowOrStack = rowOrStackForce False

-- | If all given @Block@s are single-line and the @Bool@ is @False@,
-- then makes a new single-line @Block@, with the @Maybe Line@ interspersed.
-- Otherwise, makes a vertical `stack` of the given @Block@s.
{-# INLINE rowOrStackForce #-}
rowOrStackForce :: Bool -> Maybe Line -> NonEmpty Block -> Block
rowOrStackForce _ _ (single :| []) = single
rowOrStackForce forceMultiline joiner blocks =
  case allSingles blocks of
    Right (lines, isMustBreak)
      | not forceMultiline ->
          mkLine $ join lines
      where
        mkLine = if isMustBreak then mustBreak else line
        join = case joiner of
          Nothing -> sconcat
          Just j -> sconcat . NonEmpty.intersperse j
    _ ->
      stack blocks

-- | Same as `rowOrIndentForce` @False@.
{-# INLINE rowOrIndent #-}
rowOrIndent :: Maybe Line -> NonEmpty Block -> Block
rowOrIndent = rowOrIndentForce False

-- | This is the same as `rowOrStackForce`, but all non-first lines in
-- the resulting block are indented one additional level.
{-# INLINE rowOrIndentForce #-}
rowOrIndentForce :: Bool -> Maybe Line -> NonEmpty Block -> Block
rowOrIndentForce _ _ (single :| []) = single
rowOrIndentForce forceMultiline joiner blocks@(b1 :| rest) =
  case allSingles blocks of
    Right (lines, isMustBreak)
      | not forceMultiline ->
          mkLine $ join lines
      where
        mkLine = if isMustBreak then mustBreak else line
        join = case joiner of
          Nothing -> sconcat
          Just j -> sconcat . NonEmpty.intersperse j
    _ ->
      stack (b1 :| (indent <$> rest))

data AllSingles a
  = Failed
  | Single (Maybe RequiredLineBreaks) a
  deriving (Show, Functor, Eq)

instance Applicative AllSingles where
  pure = Single Nothing
  liftA2 f (Single m1 a) (Single m2 b) =
    case (m1, m2) of
      (Just MustBreakAtEnd, Nothing) -> Single (Just MustBreakAtEnd) (f a b)
      (Just MustBreakAtEnd, Just _) -> Failed
      (m, Nothing) -> Single m (f a b)
      (_, m) -> Single m (f a b)
  liftA2 _ Failed _ = Failed
  liftA2 _ _ Failed = Failed

allSingles :: Traversable t => t Block -> Either (t Block) (t Line, Bool)
allSingles blocks =
  case traverse allSingles' blocks of
    Single (Just MustBreakAtEnd) lines -> Right (lines, True)
    Single (Just NoRequiredBreaks) lines -> Right (lines, False)
    _ -> Left blocks
  where
    allSingles' :: Block -> AllSingles Line
    allSingles' = \case
      SingleLine breaks (Indented _ l) -> Single (Just breaks) l
      _ -> Failed

-- | A convenience alias for `rowOrStack (Just space)`.
spaceSeparatedOrStack :: NonEmpty Block -> Block
spaceSeparatedOrStack = rowOrStack (Just space)

-- | A convenience alias for `rowOrStackForce (Just space)`.
spaceSeparatedOrStackForce :: Bool -> NonEmpty Block -> Block
spaceSeparatedOrStackForce force = rowOrStackForce force (Just space)

-- | A convenience alias for `rowOrIndentForce (Just space)`.
spaceSeparatedOrIndent :: NonEmpty Block -> Block
spaceSeparatedOrIndent = rowOrIndent (Just space)

-- | A convenience alias for `rowOrIndentForce (Just space)`.
spaceSeparatedOrIndentForce :: Bool -> NonEmpty Block -> Block
spaceSeparatedOrIndentForce force = rowOrIndentForce force (Just space)

-- | Adds the prefix to the first line,
-- and pads the other lines with spaces of the given length.
-- You are responsible for making sure that the given length is the actual length of the content of the given @Line@.
--
-- NOTE: An exceptional case that we haven't really designed for is if the first line of the input Block is indented.
--
-- EXAMPLE:
--
-- @
-- abcde
-- xyz
-- ----->
-- myPrefix abcde
--         xyz
-- @
prefix :: Word -> Line -> Block -> Block
prefix prefixLength pref =
  let padLineWithSpaces (Indented i l) = Indented (Indent.spaces prefixLength <> i) l

      addPrefixToLine Blank = stripEnd pref
      addPrefixToLine l = pref <> l
   in mapFirstLine (fmap addPrefixToLine) padLineWithSpaces

stripEnd :: Line -> Line
stripEnd = \case
  Space -> Blank
  Row r1 r2 ->
    case stripEnd r2 of
      Blank -> stripEnd r1
      r2' -> Row r1 r2'
  Text t -> Text t
  Blank -> Blank

-- | Adds the given suffix to then end of the last line of the @Block@.
addSuffix :: Line -> Block -> Block
addSuffix suffix =
  mapLastLine $ fmap (<> suffix)

renderIndentedLine :: Indented Line -> B.Builder
renderIndentedLine (Indented i line') =
  renderLine i (stripEnd line') <> B.char7 '\n'

spaces :: Int -> B.Builder
spaces i =
  B.byteString (ByteString.replicate i 0x20 {- space -})

renderLine :: Indent -> Line -> B.Builder
renderLine i = \case
  Text text ->
    spaces (Indent.width i) <> text
  Space ->
    spaces (1 + Indent.width i)
  Row left right ->
    renderLine i left <> renderLine mempty right
  Blank ->
    mempty

-- | Converts a @Block@ into a `Data.ByteString.Builder.Builder`.
--
-- You can then write to a file with `Data.ByteString.Builder.writeFile`,
-- or convert to @Text@ with @Data.Text.Encoding.decodeUtf8@ . `Data.ByteString.Builder.toLazyByteString`
render :: Block -> B.Builder
render = \case
  SingleLine _ line' ->
    renderIndentedLine line'
  Stack l1 rest ->
    foldMap renderIndentedLine (l1 : NonEmpty.toList rest)