{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings #-}

-- | This module provides the 'FormattedWriter' monad, which extends a basic
-- "Writer" over 'Text' by allowing formatting features like indentation, line
-- breaks, and ensuring the output is well-structured for readability.
--
-- The 'FormattedWriter' is useful for generating pretty-printed output, such as
-- source code generation, where indentation and newline management are essential.
module Language.Fiddle.Compiler.Backend.Internal.FormattedWriter
  ( -- * Types
    FormattedWriter,
    -- | The main writer monad with formatting features.
    -- * Core Operations
    ensureNL,
    -- | Ensures that there is a newline at the end of the current line.
    flush,
    -- | Flushes any pending text (writes the pending line).
    incIndent,
    -- | Increases the indentation level.
    decIndent,
    -- | Decreases the indentation level.
    indented,
    -- | Performs an action with increased indentation.
    execFormattedWriter,
    -- | Runs the 'FormattedWriter' and produces the final formatted 'Text'.
    cMacroModeStart,
    cMacroModeStop,
  )
where

import Control.Monad (forM_)
import Control.Monad.RWS.Strict
import Data.Text (Text)
import qualified Data.Text as Text

-- | Internal state for the 'FormattedWriter' monad.
--
-- * 'indentLevel': The current level of indentation.
-- * 'pendingLine': The current line that is being built but not yet written.
data Fmt = Fmt
  { -- | Current indentation level.
    indentLevel :: Int,
    linePost :: Text,
    -- | Text that has been added but not yet written to the output.
    pendingLine :: Text
  }

-- | The 'FormattedWriter' is a monad that provides functionality for writing
-- formatted text, with control over indentation and newlines.
newtype FormattedWriter a = FormattedWriter (RWS () Text Fmt a)
  deriving newtype (Functor, Applicative, Monad, MonadState Fmt)

-- | 'MonadWriter' instance for 'FormattedWriter', which allows writing text
-- while maintaining proper indentation and ensuring that newlines are handled
-- correctly. Each line is prefixed with the correct indentation level.
instance MonadWriter Text FormattedWriter where
  tell txt = FormattedWriter $ do
    -- Get the current indentation level and prepare the indentation prefix.
    indent <- (`Text.replicate` "  ") <$> gets indentLevel
    -- Split the text into lines.
    let lines = Text.splitOn "\n" txt
    -- For each line, write the indent, the pending line, and then the line itself.
    forM_ (Prelude.init lines) $ \line -> do
      pending <- gets pendingLine
      lp <- gets linePost
      modify $ \s -> s {pendingLine = ""} -- Reset pending line
      tell indent -- Add indentation
      tell pending -- Write any pending text
      tell line -- Write the actual line
      tell lp
      tell "\n" -- Add a newline after the line
      -- The last fragment is kept as the new pending line, to be written later.
    modify $ \s -> s {pendingLine = pendingLine s <> Prelude.last lines}

  -- Allow listening to the written text within the 'FormattedWriter' context.
  listen (FormattedWriter fn) = FormattedWriter $ listen fn

  -- Allow transformation of the written text via a pass operation.
  pass (FormattedWriter fn) = FormattedWriter $ pass fn

-- | Ensures that a newline is written if there is any pending text. This is
-- useful for ensuring that the output ends on a clean line.
ensureNL :: FormattedWriter ()
ensureNL = do
  p <- gets pendingLine
  if Text.null p
    then return ()
    else do
      modify $ \s -> s {pendingLine = ""}
      tell p
      tell "\n"

-- | Flushes the pending line by writing it to the output and clearing the
-- pending state. This is useful when you want to force the current line to
-- be written immediately.
flush :: FormattedWriter ()
flush = do
  p <- gets pendingLine
  modify $ \s -> s {pendingLine = ""}
  lp <- gets linePost
  tell (p <> lp)

-- | Increases the indentation level by one. This will affect all subsequent
-- lines written within the 'FormattedWriter' monad.
incIndent :: FormattedWriter ()
incIndent = modify (\f -> f {indentLevel = indentLevel f + 1})

-- | Decreases the indentation level by one. This will affect all subsequent
-- lines written within the 'FormattedWriter' monad.
decIndent :: FormattedWriter ()
decIndent = modify (\f -> f {indentLevel = indentLevel f - 1})

cMacroModeStart :: FormattedWriter ()
cMacroModeStart = modify (\f -> f { linePost = " \\" })

cMacroModeStop :: FormattedWriter ()
cMacroModeStop = modify (\f -> f { linePost = "" })

-- | Runs the given 'FormattedWriter' action with an additional indentation level.
-- Once the action completes, the indentation level is decreased.
indented :: FormattedWriter () -> FormattedWriter ()
indented fn = do
  incIndent -- Increase indentation level
  fn -- Run the action with the increased indentation
  decIndent -- Restore the original indentation level

-- | Runs a 'FormattedWriter' and returns the final formatted 'Text' output.
-- It ensures that any pending lines are flushed before returning the result.
execFormattedWriter :: FormattedWriter a -> Text
execFormattedWriter ((>> flush) -> FormattedWriter rws) =
  snd $ execRWS rws () (Fmt 0 "" "") -- Execute the RWS monad, starting with no indentation and no pending line.