{-# LANGUAGE NoImplicitPrelude #-}
{-# OPTIONS_GHC -fno-warn-unused-imports #-}
-- TODO92

module Language.Haskell.Brittany.Internal.S2_SplitModule
  ( splitModule
  )
where



import           Language.Haskell.Brittany.Internal.Prelude

import qualified Data.Generics                 as SYB
import qualified GHC
import           GHC                            ( AddEpAnn(AddEpAnn)
                                                , Anchor(Anchor)
                                                , EpAnn(EpAnn, EpAnnNotUsed)
                                                , EpAnnComments
                                                  ( EpaComments
                                                  , EpaCommentsBalanced
                                                  )
                                                , EpaComment(EpaComment)
                                                , EpaCommentTok
                                                  ( EpaBlockComment
                                                  , EpaDocCommentNamed
                                                  , EpaDocCommentNext
                                                  , EpaDocCommentPrev
                                                  , EpaDocOptions
                                                  , EpaDocSection
                                                  , EpaEofComment
                                                  , EpaLineComment
                                                  )
                                                , EpaLocation(EpaSpan)
                                                , GenLocated(L)
                                                , HsModule(HsModule)
                                                , LEpaComment
                                                , LHsDecl
                                                , LImportDecl
                                                , SrcSpan
                                                  ( RealSrcSpan
                                                  , UnhelpfulSpan
                                                  )
                                                , SrcSpanAnn'(SrcSpanAnn)
                                                , anchor
                                                , ideclName
                                                , moduleNameString
                                                , srcLocCol
                                                , srcLocLine
                                                , unLoc
                                                )
import qualified GHC.OldList                   as List
import           GHC.Parser.Annotation          ( DeltaPos
                                                  ( DifferentLine
                                                  , SameLine
                                                  )
                                                , EpaCommentTok(EpaEofComment)
                                                )
import           GHC.Types.SrcLoc               ( realSrcSpanEnd )
import qualified Language.Haskell.GHC.ExactPrint.Types
                                               as ExactPrint
import qualified Language.Haskell.GHC.ExactPrint.Utils
                                               as ExactPrint
import           Safe                           ( maximumMay )

import           Language.Haskell.Brittany.Internal.Components.BriDoc
import           Language.Haskell.Brittany.Internal.ToBriDoc.Module
import           Language.Haskell.Brittany.Internal.Types



splitModule
  :: Bool
  -> GHC.ParsedSource
  -> Maybe GHC.RealSrcLoc
  -> FinalList ModuleElement ExactPrint.Pos
splitModule shouldReformatHead lmod posWhere = do
  let L moduleSpan modl          = lmod
      HsModule _ _layout _name _exports imports decls _ _ = modl
      (hsModAnn', finalComments) = case GHC.hsmodAnn modl of
        EpAnn a modAnns (EpaCommentsBalanced prior post) ->
          (EpAnn a modAnns (EpaCommentsBalanced prior []), post)
        _ -> (GHC.hsmodAnn modl, [])
      moduleWithoutComments =
        L moduleSpan modl { GHC.hsmodAnn = hsModAnn', GHC.hsmodDecls = [] }
  lastSpan <- if shouldReformatHead
    then do
      finalYield $ MEPrettyModuleHead moduleWithoutComments
      let locBeforeImports =
            maximumMay 
              $  [ realSrcSpanEnd $ anchor a
                 | L a _ <- case hsModAnn' of
                   EpAnn _ _ (EpaComments cs             ) -> cs
                   EpAnn _ _ (EpaCommentsBalanced cs1 cs2) -> cs1 ++ cs2
                   EpAnnNotUsed -> error "unexpected EpAnnNotUsed"
                 ]
              ++ [ pos | Just pos <- [posWhere] ]
      let (importLines, lastSpan) = finalToList $ transformToImportLine
            ( maybe 0 srcLocLine locBeforeImports
            , maybe 1 srcLocCol  locBeforeImports
            )
            imports
      let commentedImports = groupifyImportLines importLines
      sortCommentedImports commentedImports `forM_` \case
        EmptyLines n ->
          finalYield $ MEWhitespace $ DifferentLine n 1
        SamelineComment{} ->
          error "brittany internal error: splitModule SamelineComment"
        NewlineComment  comm   -> finalYield $ MEComment comm
        ImportStatement record -> do
          forM_ (commentsBefore record) $ finalYield . MEComment
          finalYield
            $ MEImportDecl (importStatement record) (commentsSameline record)
          forM_ (commentsAfter record) $ finalYield . MEComment
      pure $ lastSpan
    else do
      finalYield $ MEExactModuleHead moduleWithoutComments
      pure
        $  maybe (1, 1) (ExactPrint.ss2posEnd)
        $  maximumMay
        $  [ GHC.anchor a
           | L a _ <- GHC.priorComments $ case hsModAnn' of
             EpAnn _ _ cs -> cs
             EpAnnNotUsed -> error "unexpected EpAnnNotUsed"
           ]
        ++ [ s | AddEpAnn _ (EpaSpan s) <- GHC.am_main $ GHC.anns hsModAnn' ]
        ++ [ GHC.anchor a
           | L da _ <- GHC.hsmodImports modl
           , L a  _ <- case GHC.ann da of
             EpAnn _ _ (EpaComments l          ) -> l
             EpAnn _ _ (EpaCommentsBalanced _ l) -> l
             EpAnnNotUsed                        -> []
           ]
        ++ [ span
           | L (SrcSpanAnn _ (GHC.RealSrcSpan span _)) _ <- GHC.hsmodImports
             modl
           ]
  spanAfterDecls <- enrichDecls lastSpan decls
  enrichComms spanAfterDecls finalComments


enrichComms
  :: ExactPrint.Pos -> [LEpaComment] -> FinalList ModuleElement ExactPrint.Pos
enrichComms lastSpanEnd = \case
  [] -> pure lastSpanEnd
  (L (Anchor span _) (EpaComment EpaEofComment _) : commRest) -> do
    finalYield $ MEWhitespace $ case ExactPrint.ss2delta lastSpanEnd span of -- TODO92 move this (l-1) bit into utility function
      SameLine i        -> SameLine i
      DifferentLine l c -> DifferentLine (l - 1) c
    enrichComms (ExactPrint.ss2posEnd span) commRest
  (L (Anchor span _) (EpaComment comm _) : commRest) -> do
    case ExactPrint.ss2delta lastSpanEnd span of
      SameLine i -> do
        finalYield $ MEComment (i, comm)
      DifferentLine l c -> do
        finalYield $ MEWhitespace $ DifferentLine (l - 1) c
        finalYield $ MEComment (0, comm)
    enrichComms (ExactPrint.ss2posEnd span) commRest

enrichDecls
  :: ExactPrint.Pos -> [LHsDecl GhcPs] -> FinalList ModuleElement ExactPrint.Pos
enrichDecls lastSpanEnd = \case
  [] -> finalPure $ lastSpanEnd
  (L (SrcSpanAnn dAnn rlspan@(GHC.RealSrcSpan span _)) decl : declRest) ->
    case dAnn of
      EpAnn dAnchor items (EpaComments dComments) -> do
        let
          withoutComments =
            (L (SrcSpanAnn (EpAnn dAnchor items (EpaComments [])) rlspan) decl)
          commentExtract = \case
            L (GHC.Anchor anch _) EpaComment{} -> ExactPrint.ss2posEnd anch
          -- It would be really nice if `ExactPrint.ss2posEnd span` was
          -- sufficient. But for some reason the comments are not
          -- (consistently) included in the length of the anchor. I.e.
          -- there are cases where a syntax tree node has an anchor from
          -- pos A -> pos B. But then somewhere _below_ that node is a
          -- comment that has an anchor pos B -> pos C.
          -- We simply detect this here.
          -- We probably do some redundant `SYB.everything` lookups
          -- throughout the code now. But optimizing it is not easy, and
          -- at worst it is larger constant factor on the size of the
          -- input, so it isn't _that_ bad.
          fixedSpanEnd = SYB.everything
            max
            (SYB.mkQ (ExactPrint.ss2posEnd span) commentExtract)
            decl
        case ExactPrint.ss2delta lastSpanEnd span of
          SameLine{} -> pure ()
          DifferentLine n _ ->
            finalYield $ MEWhitespace $ DifferentLine (n - 1) 1
        let (afterComms, span2) = finalToList $ enrichComms fixedSpanEnd (reverse dComments)
        let (immediate, later) = List.span (\case
              MEComment{} -> True
              _ -> False
              ) afterComms
        finalYield $ MEDecl withoutComments [ comm | MEComment comm <- immediate ]
        later `forM_` finalYield
        enrichDecls span2 declRest
      EpAnn _anchor _items (EpaCommentsBalanced{}) ->
        error "EpaCommentsBalanced"
      EpAnnNotUsed -> error "EpAnnNotUsed"
  (L (SrcSpanAnn _ann (GHC.UnhelpfulSpan{})) _decl : _declRest) ->
    error "UnhelpfulSpan"


-- module head pretty-printing

data ImportLine
  = EmptyLines Int
  | SamelineComment (Int, EpaCommentTok)
  | NewlineComment  (Int, EpaCommentTok) -- indentation and comment
  | ImportStatement ImportStatementRecord

instance Show ImportLine where
  show = \case
    EmptyLines n      -> "EmptyLines " ++ show n
    SamelineComment{} -> "SamelineComment"
    NewlineComment{}  -> "NewlineComment"
    ImportStatement r ->
      "ImportStatement " ++ show (length $ commentsBefore r) ++ " " ++ show
        (length $ commentsAfter r)

data ImportStatementRecord = ImportStatementRecord
  { commentsBefore   :: [(Int, EpaCommentTok)]
  , importStatement  :: LImportDecl GhcPs
  , commentsSameline :: [(Int, EpaCommentTok)]
  , commentsAfter    :: [(Int, EpaCommentTok)]
  }

instance Show ImportStatementRecord where
  show r =
    "ImportStatement " ++ show (length $ commentsBefore r) ++ " " ++ show
      (length $ commentsAfter r)


transformToImportLine
  :: ExactPrint.Pos
  -> [LImportDecl GhcPs]
  -> FinalList ImportLine ExactPrint.Pos
transformToImportLine startPos is =
  let
    flattenComms
      :: [LEpaComment] -> ExactPrint.Pos -> FinalList ImportLine ExactPrint.Pos
    flattenComms = \case
      [] -> finalPure
      (L (Anchor span _) (EpaComment comm _) : commRest) -> \lastSpanEnd -> do
        case ExactPrint.ss2delta lastSpanEnd span of
          SameLine i -> do
            finalYield $ SamelineComment (i, comm)
          DifferentLine l c -> do
            finalYield $ EmptyLines (l - 1)
            finalYield $ NewlineComment (c - 1, comm)
        flattenComms commRest (ExactPrint.ss2posEnd span)
    flattenDecls
      :: [LImportDecl GhcPs]
      -> ExactPrint.Pos
      -> FinalList ImportLine ExactPrint.Pos
    flattenDecls = \case
      [] -> finalPure
      (L (SrcSpanAnn epAnn srcSpan@(RealSrcSpan declSpan _)) decl : declRest)
        -> \lastSpanEnd ->
          let (commsBefore, commsAfter, cleanEpAnn) = case epAnn of
                EpAnn anch s (EpaComments cs) ->
                  ([], reverse cs, EpAnn anch s (EpaComments []))
                EpAnn anch s (EpaCommentsBalanced cs1 cs2) ->
                  (reverse cs1, reverse cs2, EpAnn anch s (EpaComments []))
                EpAnnNotUsed -> ([], [], EpAnnNotUsed)
          in
            do
              span1 <- flattenComms commsBefore lastSpanEnd
              let newlines = case ExactPrint.ss2delta span1 declSpan of
                    SameLine _        -> 0
                    DifferentLine i _ -> i - 1
              finalYield
                $ EmptyLines newlines
              finalYield $ ImportStatement ImportStatementRecord
                { commentsBefore   = []
                , importStatement  = L (SrcSpanAnn cleanEpAnn srcSpan) decl
                , commentsSameline = []
                , commentsAfter    = []
                }
              span2 <- flattenComms commsAfter (ExactPrint.ss2posEnd declSpan)
              flattenDecls declRest span2
      (L (SrcSpanAnn _epAnn UnhelpfulSpan{}) _decl : _declRest) ->
        error "UnhelpfulSpan"
  in
    flattenDecls is startPos

data Partial = PartialCommsOnly [(Int, EpaCommentTok)]
             | PartialImport ImportStatementRecord

groupifyImportLines :: [ImportLine] -> [ImportLine]
groupifyImportLines ls = finalToList_ $ go (PartialCommsOnly []) ls
 where
  go acc [] = case acc of
    PartialCommsOnly comms ->
      reverse comms `forM_` \comm -> finalYield $ NewlineComment comm
    PartialImport partialRecord ->
      finalYield $ ImportStatement $ unpartial partialRecord
  go acc (line1 : lineR) = do
    newAcc <- case acc of
      PartialCommsOnly comms -> case line1 of
        e@EmptyLines{} -> do
          reverse comms `forM_` \comm -> finalYield $ NewlineComment comm
          finalYield e
          pure $ PartialCommsOnly []
        SamelineComment comm -> do
          pure $ PartialCommsOnly (comm : comms)
        NewlineComment comm -> pure $ PartialCommsOnly (comm : comms)
        ImportStatement record ->
          pure $ PartialImport $ record { commentsBefore = comms }
      PartialImport partialRecord -> case line1 of
        e@EmptyLines{} -> do
          finalYield $ ImportStatement $ unpartial partialRecord
          finalYield e
          pure $ PartialCommsOnly []
        SamelineComment comm -> do
          if (null $ commentsAfter partialRecord)
            then pure $ PartialImport partialRecord
              { commentsSameline = comm : commentsSameline partialRecord
              }
            else pure $ PartialImport partialRecord
              { commentsAfter = comm : commentsAfter partialRecord
              }
        NewlineComment comm -> pure $ PartialImport $ partialRecord
          { commentsAfter = comm : commentsAfter partialRecord
          }
        ImportStatement record -> do
          let contestedComments = commentsAfter partialRecord
          finalYield $ ImportStatement $ unpartial $ partialRecord
            { commentsAfter = []
            }
          pure $ PartialImport $ record { commentsBefore = contestedComments }
            -- comments in between will stay connected to the following decl
    go newAcc lineR
  unpartial :: ImportStatementRecord -> ImportStatementRecord
  unpartial partialRecord = ImportStatementRecord
    { commentsBefore   = reverse (commentsBefore partialRecord)
    , importStatement  = importStatement partialRecord
    , commentsSameline = reverse (commentsSameline partialRecord)
    , commentsAfter    = reverse (commentsAfter partialRecord)
    }


sortCommentedImports :: [ImportLine] -> [ImportLine]
sortCommentedImports =
  -- TODO92 we don't need this unpackImports, it is implied later in the process
                       mergeGroups . map (fmap (sortGroups)) . groupify
 where
  -- unpackImports :: [ImportLine] -> [ImportLine]
  -- unpackImports xs = xs >>= \case
  --   l@EmptyLines{}      -> [l]
  --   l@NewlineComment{}  -> [l]
  --   l@SamelineComment{} -> [l]
  --   ImportStatement r ->
  --     map NewlineComment (commentsBefore r) ++ [ImportStatement r] ++ map
  --       NewlineComment
  --       (commentsAfter r)
  mergeGroups :: [Either ImportLine [ImportStatementRecord]] -> [ImportLine]
  mergeGroups xs = xs >>= \case
    Left  x -> [x]
    Right y -> ImportStatement <$> y
  sortGroups :: [ImportStatementRecord] -> [ImportStatementRecord]
  sortGroups =
    List.sortOn (moduleNameString . unLoc . ideclName . unLoc . importStatement)
  groupify :: [ImportLine] -> [Either ImportLine [ImportStatementRecord]]
  groupify cs = go [] cs
   where
    go [] = \case
      (l@EmptyLines{}      : rest) -> Left l : go [] rest
      (l@NewlineComment{}  : rest) -> Left l : go [] rest
      (l@SamelineComment{} : rest) -> Left l : go [] rest
      (ImportStatement r   : rest) -> go [r] rest
      []                           -> []
    go acc = \case
      (l@EmptyLines{}      : rest) -> Right (reverse acc) : Left l : go [] rest
      (l@NewlineComment{}  : rest) -> Left l : Right (reverse acc) : go [] rest
      (l@SamelineComment{} : rest) -> Left l : Right (reverse acc) : go [] rest
      (ImportStatement r   : rest) -> go (r : acc) rest
      []                           -> [Right (reverse acc)]