{-# LANGUAGE BlockArguments             #-}
{-# LANGUAGE DeriveGeneric              #-}
{-# LANGUAGE DerivingStrategies         #-}
{-# LANGUAGE DoAndIfThenElse            #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase                 #-}
{-# LANGUAGE RecordWildCards            #-}
module Language.Haskell.Stylish.Printer
  ( Printer(..)
  , PrinterConfig(..)
  , PrinterState(..)

    -- * Alias
  , P

    -- * Functions to use the printer
  , runPrinter
  , runPrinter_

    -- ** Combinators
  , comma
  , dot
  , getCurrentLine
  , getCurrentLineLength
  , newline
  , parenthesize
  , prefix
  , putComment
  , putMaybeLineComment
  , putOutputable
  , putCond
  , putType
  , putRdrName
  , putText
  , sep
  , space
  , spaces
  , suffix
  , pad

    -- ** Advanced combinators
  , withColumns
  , modifyCurrentLine
  , wrapping
  ) where

--------------------------------------------------------------------------------
import           Prelude                         hiding (lines)

--------------------------------------------------------------------------------
import qualified GHC.Hs                          as GHC
import           GHC.Hs.Extension                (GhcPs)
import           GHC.Types.Name.Reader           (RdrName (..))
import           GHC.Types.SrcLoc                (GenLocated (..))
import qualified GHC.Types.SrcLoc                as GHC
import           GHC.Utils.Outputable            (Outputable)

--------------------------------------------------------------------------------
import           Control.Monad                   (forM_, replicateM_)
import           Control.Monad.Reader            (MonadReader, ReaderT (..),
                                                  asks, local)
import           Control.Monad.State             (MonadState, State, get, gets,
                                                  modify, put, runState)
import           Data.List                       (foldl')

--------------------------------------------------------------------------------
import           Language.Haskell.Stylish.GHC    (showOutputable)
import           Language.Haskell.Stylish.Module (Lines)

-- | Shorthand for 'Printer' monad
type P = Printer

-- | Printer that keeps state of file
newtype Printer a = Printer (ReaderT PrinterConfig (State PrinterState) a)
  deriving (Applicative, Functor, Monad, MonadReader PrinterConfig, MonadState PrinterState)

-- | Configuration for printer, currently empty
data PrinterConfig = PrinterConfig
    { columns :: !(Maybe Int)
    }

-- | State of printer
data PrinterState = PrinterState
  { lines       :: !Lines
  , linePos     :: !Int
  , currentLine :: !String
  }

-- | Run printer to get printed lines out of module as well as return value of monad
runPrinter :: PrinterConfig -> Printer a -> (a, Lines)
runPrinter cfg (Printer printer) =
  let
    (a, PrinterState parsedLines _ startedLine) = runReaderT printer cfg `runState` PrinterState [] 0 ""
  in
    (a, parsedLines <> if startedLine == [] then [] else [startedLine])

-- | Run printer to get printed lines only
runPrinter_ :: PrinterConfig -> Printer a -> Lines
runPrinter_ cfg printer = snd (runPrinter cfg printer)

-- | Print text
putText :: String -> P ()
putText txt = do
  l <- gets currentLine
  modify \s -> s { currentLine = l <> txt }

-- | Check condition post action, and use fallback if false
putCond :: (PrinterState -> Bool) -> P b -> P b -> P b
putCond p action fallback = do
  prevState <- get
  res <- action
  currState <- get
  if p currState then pure res
  else put prevState >> fallback

-- | Print an 'Outputable'
putOutputable :: Outputable a => a -> P ()
putOutputable = putText . showOutputable

-- | Put all comments that has positions within 'SrcSpan' and separate by
--   passed @P ()@
{-
putAllSpanComments :: P () -> SrcSpan -> P ()
putAllSpanComments suff = \case
  UnhelpfulSpan _ -> pure ()
  RealSrcSpan rspan -> do
    cmts <- removeComments \(L rloc _) ->
      srcSpanStartLine rloc >= srcSpanStartLine rspan &&
      srcSpanEndLine rloc <= srcSpanEndLine rspan

    forM_ cmts (\c -> putComment c >> suff)
-}

-- | Print any comment
putComment :: GHC.EpaComment -> P ()
putComment epaComment = case GHC.ac_tok epaComment of
  GHC.EpaDocComment hs  -> putText $ show hs
  GHC.EpaLineComment s  -> putText s
  GHC.EpaDocOptions s   -> putText s
  GHC.EpaBlockComment s -> putText s
  GHC.EpaEofComment     -> pure ()

putMaybeLineComment :: Maybe GHC.EpaComment -> P ()
putMaybeLineComment = \case
    Nothing  -> pure ()
    Just cmt -> space >> putComment cmt

-- | Print a 'RdrName'
putRdrName :: GenLocated GHC.SrcSpanAnnN RdrName -> P ()
putRdrName rdrName = case GHC.unLoc rdrName of
    Unqual name -> do
      let (pre, post) = nameAnnAdornments $
            GHC.epAnnAnnsL $ GHC.ann $ GHC.getLoc rdrName
      putText pre
      putText (showOutputable name)
      putText post
    Qual modulePrefix name ->
      putModuleName modulePrefix >> dot >> putText (showOutputable name)
    Orig _ name ->
      putText (showOutputable name)
    Exact name ->
      putText (showOutputable name)

nameAnnAdornments :: [GHC.NameAnn] -> (String, String)
nameAnnAdornments = foldl'
    (\(accl, accr) nameAnn ->
        let (l, r) = nameAnnAdornment nameAnn in (accl ++ l, r ++ accr))
    (mempty, mempty)

nameAnnAdornment :: GHC.NameAnn -> (String, String)
nameAnnAdornment = \case
    GHC.NameAnn {..}       -> fromAdornment nann_adornment
    GHC.NameAnnCommas {..} -> fromAdornment nann_adornment
    GHC.NameAnnBars {..}   -> fromAdornment nann_adornment
    GHC.NameAnnOnly {..}   -> fromAdornment nann_adornment
    GHC.NameAnnRArrow {}   -> (mempty, mempty)
    GHC.NameAnnQuote {}    -> ("'", mempty)
    GHC.NameAnnTrailing {} -> (mempty, mempty)
  where
    fromAdornment GHC.NameParens     = ("(", ")")
    fromAdornment GHC.NameBackquotes = ("`", "`")
    fromAdornment GHC.NameParensHash = ("#(", "#)")
    fromAdornment GHC.NameSquare     = ("[", "]")

-- | Print module name
putModuleName :: GHC.ModuleName -> P ()
putModuleName = putText . GHC.moduleNameString

-- | Print type
putType :: GHC.LHsType GhcPs -> P ()
putType ltp = case GHC.unLoc ltp of
  GHC.HsFunTy _ arrowTp argTp funTp -> do
    putOutputable argTp
    space
    case arrowTp of
        GHC.HsUnrestrictedArrow {} -> putText "->"
        GHC.HsLinearArrow {}       -> putText "%1 ->"
        GHC.HsExplicitMult {}      -> putOutputable arrowTp
    space
    putType funTp
  GHC.HsAppTy _ t1 t2 ->
    putType t1 >> space >> putType t2
  GHC.HsExplicitListTy _ _ xs -> do
    putText "'["
    sep
      (comma >> space)
      (fmap putType xs)
    putText "]"
  GHC.HsExplicitTupleTy _ xs -> do
    putText "'("
    sep
      (comma >> space)
      (fmap putType xs)
    putText ")"
  GHC.HsOpTy _ _ lhs op rhs -> do
    putType lhs
    space
    putRdrName op
    space
    putType rhs
  GHC.HsTyVar _ flag rdrName -> do
    case flag of
      GHC.IsPromoted  -> putText "'"
      GHC.NotPromoted -> pure ()
    putRdrName rdrName
  GHC.HsTyLit _ tp ->
    putOutputable tp
  GHC.HsParTy _ tp -> do
    putText "("
    putType tp
    putText ")"
  GHC.HsTupleTy _ _ xs -> do
    putText "("
    sep
      (comma >> space)
      (fmap putType xs)
    putText ")"
  GHC.HsForAllTy {} ->
    putOutputable ltp
  GHC.HsQualTy {} ->
    putOutputable ltp
  GHC.HsAppKindTy _ _ _ ->
    putOutputable ltp
  GHC.HsListTy _ _ ->
    putOutputable ltp
  GHC.HsSumTy _ _ ->
    putOutputable ltp
  GHC.HsIParamTy _ _ _ ->
    putOutputable ltp
  GHC.HsKindSig _ _ _ ->
    putOutputable ltp
  GHC.HsStarTy _ _ ->
    putOutputable ltp
  GHC.HsSpliceTy _ _ ->
    putOutputable ltp
  GHC.HsDocTy _ _ _ ->
    putOutputable ltp
  GHC.HsBangTy _ _ _ ->
    putOutputable ltp
  GHC.HsRecTy _ _ ->
    putOutputable ltp
  GHC.HsWildCardTy _ ->
    putOutputable ltp
  GHC.XHsType _ ->
    putOutputable ltp

-- | Print a newline
newline :: P ()
newline = do
  l <- gets currentLine
  modify \s -> s { currentLine = "", linePos = 0, lines = lines s <> [l] }

-- | Print a space
space :: P ()
space = putText " "

-- | Print a number of spaces
spaces :: Int -> P ()
spaces i = replicateM_ i space

-- | Print a dot
dot :: P ()
dot = putText "."

-- | Print a comma
comma :: P ()
comma = putText ","

-- | Add parens around a printed action
parenthesize :: P a -> P a
parenthesize action = putText "(" *> action <* putText ")"

-- | Add separator between each element of the given printers
sep :: P a -> [P a] -> P ()
sep _ []             = pure ()
sep s (first : rest) = first >> forM_ rest ((>>) s)

-- | Prefix a printer with another one
prefix :: P a -> P b -> P b
prefix pa pb = pa >> pb

-- | Suffix a printer with another one
suffix :: P a -> P b -> P a
suffix pa pb = pb >> pa

-- | Indent to a given number of spaces.  If the current line already exceeds
-- that number in length, nothing happens.
pad :: Int -> P ()
pad n = do
    len <- length <$> getCurrentLine
    spaces $ n - len

-- | Get current line
getCurrentLine :: P String
getCurrentLine = gets currentLine

-- | Get current line length
getCurrentLineLength :: P Int
getCurrentLineLength = fmap length getCurrentLine

modifyCurrentLine :: (String -> String) -> P ()
modifyCurrentLine f = do
    s0 <- get
    put s0 {currentLine = f $ currentLine s0}

wrapping
    :: P a  -- ^ First printer to run
    -> P a  -- ^ Printer to run if first printer violates max columns
    -> P a  -- ^ Result of either the first or the second printer
wrapping p1 p2 = do
    maxCols <- asks columns
    case maxCols of
        -- No wrapping
        Nothing -> p1
        Just c  -> do
            s0 <- get
            x <- p1
            s1 <- get
            if length (currentLine s1) <= c
                -- No need to wrap
                then pure x
                else do
                    put s0
                    y <- p2
                    s2 <- get
                    if length (currentLine s1) == length (currentLine s2)
                        -- Wrapping didn't help!
                        then put s1 >> pure x
                        -- Wrapped
                        else pure y

withColumns :: Maybe Int -> P a -> P a
withColumns c = local $ \pc -> pc {columns = c}