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

module Language.Haskell.Brittany.Internal.SplitExactModule
  ( extractDeclMap
  , splitModuleDecls
  , splitModuleStart
  , getDeclBindingNames
  ) where



import           Language.Haskell.Brittany.Internal.Prelude

import qualified Data.Generics                 as SYB
import qualified Data.List.Extra
import qualified Data.Map                      as Map
import qualified Data.Text                     as Text
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
                                                , moduleName
                                                , moduleNameString
                                                , srcLocCol
                                                , srcLocLine
                                                , unLoc
                                                )
import           GHC.Types.Name                 ( getOccString )
import           GHC.Types.Name.Occurrence      ( occNameString )
import           GHC.Types.Name.Reader          ( RdrName
                                                  ( Exact
                                                  , Orig
                                                  , Qual
                                                  , Unqual
                                                  )
                                                )
import qualified GHC.OldList                   as List
import           GHC.Parser.Annotation          ( DeltaPos
                                                  ( DifferentLine
                                                  , SameLine
                                                  )
                                                , EpaCommentTok(EpaEofComment)
                                                )
import           GHC.Types.SrcLoc               ( realSrcSpanEnd
                                                , realSrcSpanStart
                                                )
import qualified Language.Haskell.GHC.ExactPrint.Types
                                               as ExactPrint
import qualified Language.Haskell.GHC.ExactPrint.Utils
                                               as ExactPrint
import           Safe                           ( maximumMay )
import qualified Control.Monad.Trans.Writer.Strict
                                               as WriterS

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



extractDeclMap :: GHC.ParsedSource -> Map GHC.RealSrcSpan [String]
extractDeclMap modul =
  Map.fromList
    [ ( case span of
          GHC.RealSrcSpan s _ -> s
          GHC.UnhelpfulSpan _ -> error "unexpected UnhelpfulSpan"
      , getDeclBindingNames decl
      )
    | decl <- decls
    , let (L (GHC.SrcSpanAnn _ span) _) = decl
    ]
  where (L _ (GHC.HsModule _ _ _ _ _ decls _ _)) = modul

splitModuleDecls :: GHC.ParsedSource -> FinalList ModuleElement ExactPrint.Pos
splitModuleDecls lmod = 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, [])
    (newImports, commsAfterImports) = case Data.List.Extra.unsnoc imports of
      Just (allButLast, L (SrcSpanAnn epAnn s@(RealSrcSpan span _)) lastImp) ->
        case epAnn of
          EpAnnNotUsed -> (imports, [])
          EpAnn anch anns (EpaComments cs) ->
            let
              (keepImports, moveImports) =
                partition
                  (\(L cAnch _) ->
                    GHC.srcSpanEndLine (anchor cAnch) <= GHC.srcSpanEndLine span
                  )
                  cs
              newLastImport =
                L (SrcSpanAnn (EpAnn anch anns (EpaComments keepImports)) s)
                  lastImp
            in
              ( allButLast ++ [newLastImport]
              , List.sortOn (\(L l _) -> l) moveImports
              )
          EpAnn anch anns (EpaCommentsBalanced cs1 cs2) ->
            let newLastImport =
                  L (SrcSpanAnn (EpAnn anch anns (EpaComments cs1)) s) lastImp
            in  (allButLast ++ [newLastImport], List.sortOn (\(L l _) -> l) cs2)
      _ -> ([], [])
    moduleWithoutComments = L moduleSpan modl { GHC.hsmodAnn     = hsModAnn'
                                              , GHC.hsmodDecls   = []
                                              , GHC.hsmodImports = newImports
                                              }
  spanAfterImports <- do
    finalYield $ MEExactModuleHead moduleWithoutComments
    pure
      $  maybe (0, 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' ]
      ++ [ span
         | L (SrcSpanAnn _ (RealSrcSpan span _)) _ <- GHC.hsmodImports modl
         ]
      ++ [ span
         | L (SrcSpanAnn _ (GHC.RealSrcSpan span _)) _ <- GHC.hsmodImports modl
         ]
  spanBeforeDecls <- enrichComms spanAfterImports commsAfterImports
  spanAfterDecls  <- enrichDecls spanBeforeDecls decls
  enrichComms spanAfterDecls finalComments

splitModuleStart
  :: GHC.ParsedSource
  -> Maybe GHC.RealSrcLoc
  -> FinalList ModuleElement ExactPrint.Pos
splitModuleStart modul posWhere = do
  finalYield $ MEPrettyModuleHead modul
  let locBeforeImports =
        maximumMay
          $  [ realSrcSpanEnd $ anchor a
             | L a _ <- case GHC.hsmodAnn $ unLoc modul 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
        )
        (GHC.hsmodImports $ unLoc modul)
  let commentedImports = groupifyImportLines importLines
  sortCommentedImports commentedImports `forM_` \case
    EmptyLines n -> finalYield $ MEWhitespace $ DifferentLine n 1
    SamelineComment{} ->
      error "brittany internal error: splitModuleStart 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

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
  lcomm@(L (Anchor span _) _) : commRest -> do
    case ExactPrint.ss2delta lastSpanEnd span of
      SameLine i -> do
        finalYield $ MEComment (i, lcomm)
      DifferentLine l c -> do
        finalYield $ MEWhitespace $ DifferentLine (l - 1) c
        finalYield $ MEComment (0, lcomm)
    enrichComms (ExactPrint.ss2posEnd span) commRest

enrichDecls
  :: ExactPrint.Pos -> [LHsDecl GhcPs] -> FinalList ModuleElement ExactPrint.Pos
enrichDecls lastSpanEnd = \case
  [] -> finalPure $ lastSpanEnd
  ldecl@(L (SrcSpanAnn dAnn (GHC.RealSrcSpan span _)) _) : declRest ->
    case dAnn of
      EpAnn _dAnchor _items (EpaComments _dComments) -> do
        let
          commentExtract
            :: [LEpaComment] -> WriterS.Writer [LEpaComment] [LEpaComment]
          commentExtract comms = do
            let (innerComments, outerComments) =
                  partition
                    (\(L (Anchor anch _) _) ->
                      (  realSrcSpanStart anch < realSrcSpanEnd span
                      && realSrcSpanEnd anch > realSrcSpanStart span
                      )
                    )
                    comms
            WriterS.tell outerComments
            pure innerComments
          (ldecl', extractedComments) = WriterS.runWriter
            $ SYB.everywhereM (SYB.mkM commentExtract) ldecl
        case ExactPrint.ss2delta lastSpanEnd span of
          SameLine{} -> pure ()
          DifferentLine n _ ->
            finalYield $ MEWhitespace $ DifferentLine (n - 1) 1
        let fixedSpanEnd = ExactPrint.ss2posEnd span
        let (afterComms, span2) = finalToList
              $ enrichComms fixedSpanEnd
                            (List.sortOn (\(L l _) -> l) extractedComments)
        let (immediate, later) =
              List.span
                (\case
                  MEComment{} -> True
                  _           -> False
                )
                afterComms
        finalYield
          $ MEDecl
            ldecl'
            [ (ind, GHC.ac_tok comm) | MEComment (ind, L _ 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, LEpaComment)
  | NewlineComment  (Int, LEpaComment) -- 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, LEpaComment)]
  , importStatement  :: LImportDecl GhcPs
  , commentsSameline :: [(Int, EpaCommentTok)]
  , commentsAfter    :: [(Int, LEpaComment)]
  }

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
      lcomm@(L (Anchor span _) _) : commRest -> \lastSpanEnd -> do
        case ExactPrint.ss2delta lastSpanEnd span of
          SameLine i -> do
            finalYield $ SamelineComment (i, lcomm)
          DifferentLine l c -> do
            finalYield $ EmptyLines (l - 1)
            finalYield $ NewlineComment (c - 1, lcomm)
        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, LEpaComment)]
             | 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 (EmptyLines 0 : lineR) = go acc lineR
  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 = tokenOnly 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
  tokenOnly :: (Int, LEpaComment) -> (Int, EpaCommentTok)
  tokenOnly (ind, L _ (EpaComment tok _)) = (ind, tok)
  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)]

rdrNameToText :: RdrName -> Text
-- rdrNameToText = Text.pack . show . flip runSDoc unsafeGlobalDynFlags . ppr
rdrNameToText (Unqual occname) = Text.pack $ occNameString occname
rdrNameToText (Qual mname occname) =
  Text.pack $ moduleNameString mname ++ "." ++ occNameString occname
rdrNameToText (Orig modul occname) =
  Text.pack $ moduleNameString (moduleName modul) ++ occNameString occname
rdrNameToText (Exact name) = Text.pack $ getOccString name

getDeclBindingNames :: GHC.LHsDecl GhcPs -> [String]
getDeclBindingNames (L _ decl) = case decl of
  GHC.SigD _ (GHC.TypeSig _ ns _) ->
    ns <&> \(L _ n) -> Text.unpack (rdrNameToText n)
  GHC.ValD _ (GHC.FunBind _ (L _ n) _ _) -> [Text.unpack $ rdrNameToText n]
  _ -> []
