{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE DataKinds #-}

module Language.Haskell.Brittany.Internal.PerModule
  ( processModule
  ) where

import           Language.Haskell.Brittany.Internal.Prelude

import qualified Control.Monad.Trans.MultiRWS.Strict
                                               as MultiRWSS
import           Data.CZipWith
import qualified Data.Map.Strict               as Map
import qualified Data.Text.Lazy                as TextL
import qualified Data.Text.Lazy.Builder        as TextL.Builder
import qualified Data.Text.Lazy.Builder        as Text.Builder
import qualified GHC
import           GHC                            ( EpaComment(EpaComment)
                                                , EpaCommentTok
                                                  ( EpaBlockComment
                                                  , EpaEofComment
                                                  , EpaLineComment
                                                  )
                                                , GenLocated(L)
                                                , HsModule(HsModule)
                                                , LHsDecl
                                                )
import qualified GHC.Types.SrcLoc              as GHC
import qualified GHC.OldList                   as List
import           GHC.Types.SrcLoc               ( srcSpanFileName_maybe )
import qualified Language.Haskell.GHC.ExactPrint
                                               as ExactPrint

import           Language.Haskell.Brittany.Internal.Components.BriDoc
import           Language.Haskell.Brittany.Internal.Config.Types
import           Language.Haskell.Brittany.Internal.Config.Types.Instances2
                                                ( )
import           Language.Haskell.Brittany.Internal.SplitExactModule
                                                ( getDeclBindingNames
                                                , splitModuleStart
                                                )
import           Language.Haskell.Brittany.Internal.ToBriDocTools
import           Language.Haskell.Brittany.Internal.WriteBriDoc
                                                ( ppBriDoc )
import           Language.Haskell.Brittany.Internal.ToBriDoc.Comment
                                                ( commentToDoc )
import           Language.Haskell.Brittany.Internal.ToBriDoc.Import
import           Language.Haskell.Brittany.Internal.ToBriDoc.Module
import           Language.Haskell.Brittany.Internal.Types
import           Language.Haskell.Brittany.Internal.Utils
import           Language.Haskell.Brittany.Internal.ToBriDoc
                                                ( layouters )
import           Language.Haskell.Brittany.Internal.PerDecl
                                                ( ppToplevelDecl )



-- BrittanyErrors can be non-fatal warnings, thus both are returned instead
-- of an Either.
-- This should be cleaned up once it is clear what kinds of errors really
-- can occur.
processModule
  :: TraceFunc
  -> Config
  -> PerItemConfig
  -> FinalList ModuleElement p
  -> IO ([BrittanyError], TextL.Text)
processModule traceFunc conf inlineConf moduleElems = do
  let FinalList moduleElementsStream = moduleElems
      ((out, errs), debugStrings) =
        runIdentity
          $ MultiRWSS.runMultiRWSTNil
          $ MultiRWSS.withMultiWriterAW
          $ MultiRWSS.withMultiWriterAW
          $ MultiRWSS.withMultiWriterW
          $ MultiRWSS.withMultiReader traceFunc
          $ moduleElementsStream
              (\modElem cont -> do
                processModuleElement modElem
                cont
              )
              (\x -> do
              -- mTell $ TextL.Builder.fromString "\n"
                pure x
              )
    -- _tracer = 
    --   -- if Seq.null debugStrings
    --   -- then id
    --   -- else
    --     trace ("---- DEBUGMESSAGES ---- ")
    --       . foldr (seq . join trace) id debugStrings
  debugStrings `forM_` \s -> useTraceFunc traceFunc s
  -- moduleElementsStream
  --   (\el rest -> do
  --     case el of
  --       MEExactModuleHead{}  -> useTraceFunc traceFunc "MEExactModuleHead"
  --       MEPrettyModuleHead{} -> useTraceFunc traceFunc "MEPrettyModuleHead"
  --       MEImportDecl{}       -> useTraceFunc traceFunc "MEImportDecl"
  --       MEDecl decl _ ->
  --         useTraceFunc
  --           traceFunc
  --           ("MEDecl " ++ intercalate "," (getDeclBindingNames decl))
  --       MEComment (y, L _ (EpaComment (EpaLineComment str) _)) ->
  --         useTraceFunc traceFunc ("MEComment " ++ show y ++ " " ++ str)
  --       MEComment (y, L _ (EpaComment (EpaBlockComment str) _)) ->
  --         useTraceFunc traceFunc ("MEComment " ++ show y ++ " " ++ take 5 str)
  --       MEComment (y, _) ->
  --         useTraceFunc traceFunc ("MEComment " ++ show y ++ " _")
  --       MEWhitespace dp -> useTraceFunc traceFunc ("MEWhitespace " ++ show dp)
  --     rest
  --   )
  --   (\_ -> pure ())
  pure (errs, TextL.Builder.toLazyText out)
 where
  shouldReformatHead =
    conf & _conf_layout & _lconfig_reformatModulePreamble & confUnpack
  wrapNonDeclToBriDoc =
    MultiRWSS.withMultiReader conf . MultiRWSS.withMultiState_
      (CommentCounter 0)
  processModuleElement
    :: ModuleElement
    -> MultiRWSS.MultiRWST
         '[TraceFunc]
         '[Text.Builder.Builder , [BrittanyError] , Seq String]
         '[]
         Identity
         ()
  processModuleElement = \case
    MEExactModuleHead modHead -> if shouldReformatHead
      then do
        let FinalList startElems =
              splitModuleStart
                modHead
                ( fmap GHC.realSrcSpanStart
                $ obtainAnnPos (GHC.hsmodAnn $ GHC.unLoc modHead) GHC.AnnWhere
                )
        startElems
          (\modElem cont -> do
            processModuleElement modElem
            cont
          )
          (\_ -> pure ())
      else wrapNonDeclToBriDoc $ do
        bdMay <- ppModuleHead modHead
        case bdMay of
          Nothing -> pure ()
          Just bd -> do
            ppBriDoc bd True
        mTell $ Text.Builder.fromString "\n"
    MEPrettyModuleHead (L loc modHead) -> wrapNonDeclToBriDoc $ do
      case modHead of
        HsModule epAnn _layoutInfo Nothing Nothing _ _ _ _ -> do
          (bd, _) <-
            briDocMToPPM layouters
            $ maybe id docFlushRemaining (srcSpanFileName_maybe loc)
            $ docHandleComms epAnn docSeparator
          ppBriDoc bd True
        HsModule _ _layoutInfo Nothing Just{} _ _ _ _ ->
          error "brittany internal error: exports without module name"
        HsModule epAnn _layoutInfo (Just n) les _ _ _ _ -> do
          (bd, _) <-
            briDocMToPPM layouters
            $ maybe id docFlushRemaining (srcSpanFileName_maybe loc)
            $ moduleNameExportBridoc epAnn n les
          ppBriDoc bd True
    MEImportDecl importDecl immediateAfterComms -> wrapNonDeclToBriDoc $ do
      (bd, _) <- briDocMToPPM layouters $ docSeq
        (layoutImport importDecl : map commentToDoc immediateAfterComms)
      ppBriDoc bd False
    MEDecl decl immediateAfterComms -> do
      let declConfig = getDeclConfig conf inlineConf decl
      MultiRWSS.withMultiReader declConfig
        $ MultiRWSS.withMultiState_ (CommentCounter 0)
        $ ppToplevelDecl decl immediateAfterComms
    MEComment (ind, L _ (EpaComment (EpaLineComment str) _)) -> do
      mTell $ TextL.Builder.fromString (replicate ind ' ' ++ str)
      mTell $ TextL.Builder.fromString "\n"
    MEComment (ind, L _ (EpaComment (EpaBlockComment str) _)) -> do
      mTell $ TextL.Builder.fromString (replicate ind ' ' ++ str)
      mTell $ TextL.Builder.fromString "\n"
    MEComment (_, L _ (EpaComment EpaEofComment _)) -> pure ()
    MEComment    _  -> mTell $ TextL.Builder.fromString "some other comment"
    MEWhitespace dp -> do
      -- mTell $ TextL.Builder.fromString "B"
      -- mTell $ TextL.Builder.fromString ("whitespace " ++ show dp)
      ppmMoveToExactLoc dp

-- Prints the information associated with the module annotation
-- This includes the imports
-- This returns a `Maybe` because it only produces a BriDocNumbered if
-- re-formatting the module head is enabled. We maybe should change that
-- for consistency.
ppModuleHead :: GHC.ParsedSource -> PPMLocal (Maybe BriDocNumbered)
ppModuleHead lmod = do
  processDefault lmod $> Nothing

processDefault
  :: (ExactPrint.ExactPrint ast, MonadMultiWriter Text.Builder.Builder m)
     -- , MonadMultiReader ExactPrint.Types.Anns m
  => GHC.Located ast
  -> m ()
processDefault x = do
  let str = ExactPrint.exactPrint x
  -- this hack is here so our print-empty-module trick does not add
  -- a newline at the start if there actually is no module header / imports
  -- / anything.
  -- TODO: instead the appropriate annotation could be removed when "cleaning"
  --       the module (header). This would remove the need for this hack!
  case str of
    "\n" -> return ()
    _    -> mTell $ Text.Builder.fromString $ List.dropWhileEnd ((==) '\n') str


getDeclConfig :: Config -> PerItemConfig -> GHC.LHsDecl GhcPs -> Config
getDeclConfig config inlineConf decl = cZipWith fromOptionIdentity config
  $ mconcat (catMaybes (mBindingConfs ++ [mDeclConf]))
 where
  declBindingNames = getDeclBindingNames decl
  mBindingConfs    = declBindingNames <&> \n ->
    Map.lookup n $ _icd_perBinding inlineConf
  mDeclConf = case GHC.locA $ GHC.getLoc decl of
    GHC.RealSrcSpan x _ -> Map.lookup x $ _icd_perAnchor inlineConf
    GHC.UnhelpfulSpan{} -> Nothing