{-# LANGUAGE NoImplicitPrelude #-}

module Language.Haskell.Brittany.Internal.WriteBriDoc
  ( ppBriDoc
  )
where

import           Language.Haskell.Brittany.Internal.Prelude

import qualified Control.Monad.Trans.MultiRWS.Strict
                                               as MultiRWSS
import qualified Data.Text                     as Text
import qualified Data.Text.Lazy.Builder        as TextL.Builder
import           GHC                            ( Anchor(Anchor)
                                                , AnchorOperation
                                                  ( MovedAnchor
                                                  , UnchangedAnchor
                                                  )
                                                , EpaComment(EpaComment)
                                                , EpaCommentTok
                                                  ( EpaBlockComment
                                                  , EpaDocCommentNamed
                                                  , EpaDocCommentNext
                                                  , EpaDocCommentPrev
                                                  , EpaDocOptions
                                                  , EpaDocSection
                                                  , EpaEofComment
                                                  , EpaLineComment
                                                  )
                                                , GenLocated(L)
                                                , LEpaComment
                                                , RealSrcLoc
                                                , DeltaPos(SameLine, DifferentLine)
                                                , srcLocLine
                                                , srcLocCol
                                                , srcSpanStartCol
                                                , srcSpanStartLine
                                                )
import           GHC.Types.SrcLoc               ( realSrcSpanEnd
                                                , realSrcSpanStart
                                                )
import qualified GHC.OldList                   as List

import qualified Language.Haskell.GHC.ExactPrint.Utils as ExactPrint

import           Language.Haskell.Brittany.Internal.Config.Types
import           Language.Haskell.Brittany.Internal.Transformations.T1_Alt
import           Language.Haskell.Brittany.Internal.Transformations.T2_Floating
import           Language.Haskell.Brittany.Internal.Transformations.T3_Par
import           Language.Haskell.Brittany.Internal.Transformations.T4_Columns
import           Language.Haskell.Brittany.Internal.Transformations.T5_Indent
import           Language.Haskell.Brittany.Internal.Types
import           Language.Haskell.Brittany.Internal.Utils
import           Language.Haskell.Brittany.Internal.WriteBriDoc.AlignmentAlgo
import           Language.Haskell.Brittany.Internal.WriteBriDoc.Types
import           Language.Haskell.Brittany.Internal.WriteBriDoc.Operators
import Language.Haskell.Brittany.Internal.Components.BriDoc
-- import Language.Haskell.Brittany.Internal.S3_ToBriDocTools


ppBriDoc :: BriDocNumbered -> Bool -> PPMLocal ()
ppBriDoc briDoc flush = do
  -- first step: transform the briDoc.
  briDoc' <- MultiRWSS.withMultiStateS BDEmpty $ do
    -- Note that briDoc is BriDocNumbered, but state type is BriDoc.
    -- That's why the alt-transform looks a bit special here.
    traceIfDumpConf "bridoc raw" _dconf_dump_bridoc_raw
      $ briDocToDoc
      $ unwrapBriDocNumbered
      $ briDoc
    -- bridoc transformation: remove alts
    transformAlts briDoc >>= mSet
    mGet
      >>= briDocToDoc
      .>  traceIfDumpConf "bridoc post-alt" _dconf_dump_bridoc_simpl_alt
    -- bridoc transformation: float stuff in
    mGet >>= transformSimplifyFloating .> mSet
    mGet
      >>= briDocToDoc
      .>  traceIfDumpConf "bridoc post-floating"
                          _dconf_dump_bridoc_simpl_floating
    -- bridoc transformation: par removal
    mGet >>= transformSimplifyPar .> mSet
    mGet
      >>= briDocToDoc
      .>  traceIfDumpConf "bridoc post-par" _dconf_dump_bridoc_simpl_par
    -- bridoc transformation: float stuff in
    mGet >>= transformSimplifyColumns .> mSet
    mGet
      >>= briDocToDoc
      .>  traceIfDumpConf "bridoc post-columns" _dconf_dump_bridoc_simpl_columns
    -- bridoc transformation: indent
    mGet >>= transformSimplifyIndent .> mSet
    mGet
      >>= briDocToDoc
      .>  traceIfDumpConf "bridoc post-indent" _dconf_dump_bridoc_simpl_indent
    mGet
      >>= briDocToDoc
      .>  traceIfDumpConf "bridoc final" _dconf_dump_bridoc_final
    -- -- convert to Simple type
    -- simpl <- mGet <&> transformToSimple
    -- return simpl

  let state = LayoutState { _lstate_baseYs           = [0]
                          , _lstate_curY             = 0
                          , _lstate_indLevels        = [0]
                          , _lstate_indLevelLinger   = 0
                          , _lstate_plannedSpace = PlannedNone
                          , _lstate_commentNewlines  = 0
                          , _lstate_markerForDelta   = Nothing
                          }
  state' <-
    MultiRWSS.withMultiStateS state
    $ MultiRWSS.withMultiStateS ([] :: [LEpaComment])
    $ do
      layoutBriDocM briDoc'
      when flush layoutFlushLine
  case _lstate_plannedSpace state' of
    PlannedNone -> if _lstate_curY state' == 0
      then pure ()
      else mTell $ TextL.Builder.fromString "\n"
    PlannedSameline _ -> if _lstate_curY state' == 0
      then pure ()
      else mTell $ TextL.Builder.fromString "\n"
    PlannedNewline l -> mTell $ TextL.Builder.fromString (replicate l '\n')
    PlannedDelta l _ -> mTell $ TextL.Builder.fromString (replicate l '\n')


layoutBriDocM :: HasCallStack => forall m . LayoutConstraints m => BriDoc -> m ()
layoutBriDocM = \case
  BDEmpty -> do
    return () -- can it be that simple
  BDLit t -> do
    layoutSetMarker Nothing
    layoutRemoveIndentLevelLinger
    layoutWriteAppend t
  BDSeq list -> do
    list `forM_` layoutBriDocM
  -- in this situation, there is nothing to do about cols.
  -- i think this one does not happen anymore with the current simplifications.
  -- BDCols cSig list | BDPar sameLine lines <- List.last list ->
  --   alignColsPar $ BDCols cSig (List.init list ++ [sameLine]) : lines
  BDCols _ list -> do
    list `forM_` layoutBriDocM
  BDSeparator -> do
    layoutAddSepSpace
  BDAddBaseY indent bd -> do
    let indentF = case indent of
          BrIndentNone         -> id
          BrIndentRegular      -> layoutWithAddBaseCol
          BrIndentRegularForce -> layoutWithAddBaseCol
          BrIndentSpecial i    -> layoutWithAddBaseColN i
    indentF $ layoutBriDocM bd
  BDBaseYPushCur bd -> do
    layoutBaseYPushCur
    layoutBriDocM bd
    layoutBaseYPop
  BDIndentLevelPushCur bd -> do
    layoutIndentLevelPushCur
    layoutBriDocM bd
  BDIndentLevelPop bd -> do
    layoutBriDocM bd
    layoutIndentLevelPop
  BDEnsureIndent indent bd -> do
    let indentF = case indent of
          BrIndentNone         -> id
          BrIndentRegular      -> layoutWithAddBaseCol
          BrIndentRegularForce -> layoutWithAddBaseCol
          BrIndentSpecial i    -> layoutWithAddBaseColN i
    indentF $ do
      layoutWriteEnsureBlock
      layoutBriDocM bd
  BDPar indent sameLine indented -> do
    layoutBriDocM sameLine
    let indentF = case indent of
          BrIndentNone         -> id
          BrIndentRegular      -> layoutWithAddBaseCol
          BrIndentRegularForce -> layoutWithAddBaseCol
          BrIndentSpecial i    -> layoutWithAddBaseColN i
    indentF $ do
      layoutWriteNewline
      layoutBriDocM indented
  BDEntryDelta dp bd -> do
    case dp of
      GHC.SameLine _ -> pure ()
      GHC.DifferentLine l _i -> mModify $ \s -> s {
        _lstate_plannedSpace = PlannedNewline l
      }
    layoutBriDocM bd
  BDLines           lines       -> alignColsLines layoutBriDocM lines
  BDAlt             []          -> error "empty BDAlt"
  BDAlt             (alt : _)   -> layoutBriDocM alt
  BDForceAlt _ bd               -> layoutBriDocM bd
  BDForwardLineMode bd          -> layoutBriDocM bd
  BDExternal shouldAddComment t -> do
    let tlines     = Text.lines $ t <> Text.pack "\n"
        tlineCount = length tlines
    when shouldAddComment $ do
      layoutWriteAppend $ Text.pack $ "{- via external! -}"
    zip [1 ..] tlines `forM_` \(i, l) -> do
      layoutWriteAppend $ l
      unless (i == tlineCount) layoutWriteNewline
  BDPlain t -> do
    layoutWriteAppend t
  -- BDAnnotationPrior comms bd -> do
  --   -- state <- mGet
  --   -- let m = _lstate_comments state
  --   -- let
  --   --   moveToExactLocationAction = case _lstate_curYOrAddNewline state of
  --   --     Left{} -> pure ()
  --   --     Right{} -> moveToExactAnn annKey
  --   -- case mAnn of
  --   --   Nothing -> moveToExactLocationAction
  --   --   Just [] -> moveToExactLocationAction
  --   --   Just priors -> do
  --   --     -- layoutResetSepSpace
  --   --     priors
  --   --       `forM_` \(ExactPrint.Types.Comment comment _ _, DifferentLine (y, x)) ->
  --   --                 when (comment /= "(" && comment /= ")") $ do
  --   --                   let commentLines = Text.lines $ Text.pack $ comment
  --   --                   case comment of
  --   --                     ('#' : _) ->
  --   --                       layoutMoveToCommentPos y (-999) (length commentLines)
  --   --                                --  ^ evil hack for CPP
  --   --                     _ -> layoutMoveToCommentPos y x (length commentLines)
  --   --                   -- fixedX <- fixMoveToLineByIsNewline x
  --   --                   -- replicateM_ fixedX layoutWriteNewline
  --   --                   -- layoutMoveToIndentCol y
  --   --                   layoutWriteAppendMultiline commentLines
  --   --       -- mModify $ \s -> s { _lstate_curYOrAddNewline = Right 0 }
  --   --     moveToExactLocationAction
  --   printComments comms
  --   mModify (\s -> s + CommentCounter (length comms))
  --   layoutBriDocM bd
  -- BDAnnotationPost comms bd -> do
  --   layoutBriDocM bd
  --   printComments comms
  --   mModify (\s -> s + CommentCounter (length comms))
    -- annMay <- do
    --   state <- mGet
    --   let m = _lstate_comments state
    --   pure $ Map.lookup annKey m
    -- let mComments = nonEmpty . extractAllComments =<< annMay
    -- -- let
    -- --   semiCount = length
    -- --     [ ()
    -- --     | Just ann <- [annMay]
    -- --     , (ExactPrint.Types.AnnSemiSep, _) <- ExactPrint.Types.annsDP ann
    -- --     ]
    -- shouldAddSemicolonNewlines <-
    --   mAsk
    --   <&> _conf_layout
    --   .> _lconfig_experimentalSemicolonNewlines
    --   .> confUnpack
    -- case mComments of
    --   Nothing -> do
    --     when shouldAddSemicolonNewlines $ do
    --       [1 .. semiCount] `forM_` const layoutWriteNewline
    --   Just comments -> do
    --     comments
    --       `forM_` \(ExactPrint.Types.Comment comment _ _, DifferentLine (y, x)) ->
    --                 when (comment /= "(" && comment /= ")") $ do
    --                   let commentLines = Text.lines $ Text.pack comment
    --                   case comment of
    --                     ('#' : _) -> layoutMoveToCommentPos y (-999) 1
    --                                --  ^ evil hack for CPP
    --                     ")" -> pure ()
    --                                --  ^ fixes the formatting of parens
    --                                --    on the lhs of type alias defs
    --                     _ -> layoutMoveToCommentPos y x (length commentLines)
    --                   -- fixedX <- fixMoveToLineByIsNewline x
    --                   -- replicateM_ fixedX layoutWriteNewline
    --                   -- layoutMoveToIndentCol y
    --                   layoutWriteAppendMultiline commentLines
    --   -- mModify $ \s -> s { _lstate_curYOrAddNewline = Right 0 }
  BDQueueComments comms bd -> do
    existing :: [GHC.LEpaComment] <- mGet
    mSet $ mergeOn (\(L l _) -> l) existing comms
    layoutBriDocM bd
  BDFlushCommentsPrior loc bd -> do
    comms <- takeBefore loc
    startCurY <- mGet <&> _lstate_curY
    printComments comms
    mModify (\s -> s + CommentCounter (length comms))
    mModify $ \s -> s
      { _lstate_markerForDelta = Nothing
      , _lstate_plannedSpace =
        case _lstate_markerForDelta s of
          Nothing -> _lstate_plannedSpace s
          Just m ->
            let p1 = (srcLocLine m, srcLocCol m)
                p2 = (srcLocLine loc, srcLocCol loc)
            in -- trace ("plannedSpace = " ++ show (_lstate_plannedSpace s)
               --   ++ ", markerForDelta = " ++ show (_lstate_markerForDelta s)
               --   ++ ", curY = " ++ show (_lstate_curY s)
               --   ++ ", p1 = " ++ show p1
               --   ++ ", p2 = " ++ show p2
               --   ++ ", startCurY = " ++ show startCurY
               --   -- ++ ", delta = " ++ show (ExactPrint.pos2delta p1 p2)
               -- ) $
              case ExactPrint.pos2delta p1 p2 of
              SameLine iDelta -> case _lstate_plannedSpace s of
                -- There is a delta after a marker. That should only happen
                -- after a block-comment. And the marker-delta is SameLine,
                -- so if layout doesn't force a newline we respect the marker
                -- delta (i.e. respect whitespace immediately after a block
                -- comment) or we just use the plan.
                PlannedNone        -> PlannedSameline iDelta
                PlannedSameline i  -> PlannedSameline (max (iDelta - 0) i)
                p@PlannedNewline{} -> p
                p@PlannedDelta{}   -> p
              DifferentLine n _ -> case _lstate_plannedSpace s of
                -- Marker with DifferentLine delta means that we want to
                -- reproduce the newlines that were present in the input.
                -- These will be emptylines in cases where it makes sense to
                -- reproduce them, e.g. between statements in a do-block.
                PlannedNone       -> PlannedDelta n startCurY
                                     --             ^^^^^^^^^
                                     --
                                     -- after a block-comment, _lstate_curY
                                     -- matches the end of the comment.
                                     -- But to not break layout we need to
                                     -- restore the original.
                                     -- E.g. consider > let {-comm-}
                                     --               >     x = 13
                                     --
                                     --             vvvvvvvvvvvvvvv
                PlannedSameline i -> PlannedDelta n (startCurY + i)
                PlannedNewline{}  -> PlannedNewline n
                PlannedDelta _ i  -> PlannedDelta n i
      }
    layoutBriDocM bd
  BDFlushCommentsPost loc shouldMark bd -> do
    layoutBriDocM bd
    if shouldMark then
      layoutSetMarker $ Just loc
    else
      layoutSetMarker Nothing
    comms <- takeBefore loc
    mModify (\s -> s + CommentCounter (length comms))
    printComments comms
  BDDebug s bd            -> do
    mTell $ TextL.Builder.fromText $ Text.pack $ "{-" ++ s ++ "-}"
    layoutBriDocM bd

takeBefore
  :: MonadMultiState [LEpaComment] m => RealSrcLoc -> m [LEpaComment]
takeBefore loc = do
  comms <- mGet
  let (before, after) = List.span
        (\(L (Anchor spanC _) _) -> realSrcSpanEnd spanC <= loc)
        comms
  mSet after
  pure before

printComments :: LayoutConstraints m => [GHC.LEpaComment] -> m ()
printComments comms = do
  let addComment isBlock s anchor prior = do
        curY <- _lstate_curY <$> mGet
        case anchor of
          Anchor span UnchangedAnchor -> do
            case ExactPrint.ss2deltaEnd prior span of
                  -- this is a stupid shitty hacky workaround
                  -- because exactprint _sometimes_ decides that
                  -- block comments don't get a proper prior span
                  -- assigned and it breaks everything.
                  SameLine 0
                    | isBlock
                    , srcSpanStartLine span > 1
                    , srcSpanStartCol span /= curY + 1
                    -> do
                    layoutSetMarker $ Just $ realSrcSpanStart span
                    layoutWriteComment
                      True
                      isBlock
                      (DifferentLine 1 (srcSpanStartCol span))
                      1
                      (Text.pack s)
                  dp -> do
                    layoutWriteComment True isBlock dp 1 (Text.pack s)
            if isBlock
              then layoutSetMarker $ Just $ realSrcSpanEnd span
              else layoutUpdateMarker $ realSrcSpanEnd span
          Anchor span (MovedAnchor dp) -> do
            layoutWriteComment False isBlock dp 1 (Text.pack s)
            if isBlock
              then layoutSetMarker $ Just $ realSrcSpanEnd span
              else layoutUpdateMarker $ realSrcSpanEnd span
  comms `forM_` \(L anch (EpaComment tok prior)) -> case tok of
    EpaDocCommentNext  s -> addComment False s anch prior
    EpaDocCommentPrev  s -> addComment False s anch prior
    EpaDocCommentNamed s -> addComment False s anch prior
    EpaDocSection _ s    -> addComment False s anch prior
    EpaDocOptions   s    -> addComment False s anch prior
    EpaLineComment  s    -> addComment False s anch prior
    EpaBlockComment s    -> addComment True s anch prior
    EpaEofComment        -> pure ()
