{-# LANGUAGE BlockArguments  #-}
{-# LANGUAGE LambdaCase      #-}
{-# LANGUAGE RecordWildCards #-}
module Language.Haskell.Stylish.Step.ModuleHeader
  ( Config (..)
  , BreakWhere (..)
  , OpenBracket (..)
  , defaultConfig
  , step
  ) where


--------------------------------------------------------------------------------
import           Control.Applicative                   ((<|>))
import           Control.Monad                         (guard, unless, when)
import           Data.Foldable                         (forM_)
import           Data.Maybe                            (fromMaybe, isJust,
                                                        listToMaybe)
import qualified GHC.Hs                                as GHC
import qualified GHC.Types.SrcLoc                      as GHC


--------------------------------------------------------------------------------
import           Language.Haskell.Stylish.Comments
import qualified Language.Haskell.Stylish.Editor       as Editor
import           Language.Haskell.Stylish.GHC
import           Language.Haskell.Stylish.Module
import           Language.Haskell.Stylish.Ordering
import           Language.Haskell.Stylish.Printer
import           Language.Haskell.Stylish.Step
import qualified Language.Haskell.Stylish.Step.Imports as Imports
import           Language.Haskell.Stylish.Util         (flagEnds)
import qualified GHC.Unit.Module.Warnings as GHC


data Config = Config
    { indent        :: Int
    , sort          :: Bool
    , separateLists :: Bool
    , breakWhere    :: BreakWhere
    , openBracket   :: OpenBracket
    }

data OpenBracket
    = SameLine
    | NextLine
    deriving (Eq, Show)

data BreakWhere
    = Exports
    | Single
    | Inline
    | Always
    deriving (Eq, Show)

defaultConfig :: Config
defaultConfig = Config
    { indent        = 4
    , sort          = True
    , separateLists = True
    , breakWhere    = Exports
    , openBracket   = NextLine
    }

step :: Maybe Int -> Config -> Step
step maxCols = makeStep "Module header" . printModuleHeader maxCols

printModuleHeader :: Maybe Int -> Config -> Lines -> Module -> Lines
printModuleHeader maxCols conf ls lmodul =
    let modul = GHC.unLoc lmodul
        name = GHC.unLoc <$> GHC.hsmodName modul

        deprecMsg = GHC.hsmodDeprecMessage $ GHC.hsmodExt modul

        startLine = fromMaybe 1 $ moduleLine <|>
            (fmap GHC.srcSpanStartLine . GHC.srcSpanToRealSrcSpan $
                GHC.getLoc lmodul)

        endLine = fromMaybe 1 $ whereLine <|>
            (do
                loc <- GHC.getLocA <$> GHC.hsmodExports modul
                GHC.srcSpanEndLine <$> GHC.srcSpanToRealSrcSpan loc)

        keywordLine kw = listToMaybe $ do
            GHC.EpAnn {..} <- pure $ GHC.hsmodAnn $ GHC.hsmodExt modul
            GHC.AddEpAnn kw' (GHC.EpaSpan s _) <- GHC.am_main anns
            guard $ kw == kw'
            pure $ GHC.srcSpanEndLine s

        moduleLine = keywordLine GHC.AnnModule
        whereLine = keywordLine GHC.AnnWhere

        commentOnLine l = listToMaybe $ do
            comment <- epAnnComments $ GHC.hsmodAnn $ GHC.hsmodExt modul
            guard $ GHC.srcSpanStartLine (GHC.anchor $ GHC.getLoc comment) == l
            pure comment

        moduleComment = moduleLine >>= commentOnLine
        whereComment =
            guard (whereLine /= moduleLine) >> whereLine >>= commentOnLine

        exportGroups = case GHC.hsmodExports modul of
            Nothing -> Nothing
            Just lexports -> Just $ doSort $ commentGroups
                (GHC.srcSpanToRealSrcSpan . GHC.getLocA)
                (GHC.unLoc lexports)
                (epAnnComments . GHC.ann $ GHC.getLoc lexports)

        printedModuleHeader = runPrinter_
            (PrinterConfig maxCols)
            (printHeader
                conf name deprecMsg exportGroups moduleComment whereComment)

        changes = Editor.changeLines
            (Editor.Block startLine endLine)
            (const printedModuleHeader) in

    Editor.apply changes ls
  where
    doSort = if sort conf then fmap (commentGroupSort compareLIE) else id

printHeader
    :: Config
    -> Maybe GHC.ModuleName
    -> Maybe (GHC.LocatedP (GHC.WarningTxt GHC.GhcPs))
    -> Maybe [CommentGroup (GHC.LIE GHC.GhcPs)]
    -> Maybe GHC.LEpaComment  -- Comment attached to 'module'
    -> Maybe GHC.LEpaComment  -- Comment attached to 'where'
    -> P ()
printHeader conf mbName mbDeprec mbExps mbModuleComment mbWhereComment = do
    forM_ mbName $ \name -> do
        putText "module"
        space
        putText (showOutputable name)

    forM_ mbDeprec \deprec -> do
        putText " "
        putText (showOutputable deprec)

    case mbExps of
        Nothing -> do
            when (isJust mbName) $ case breakWhere conf of
                Always -> do
                    attachModuleComment
                    newline
                    spaces (indent conf)
                _      -> space
            putText "where"
        Just exports -> case breakWhere conf of
            Single  | [] <- exports -> do
                printSingleLineExportList conf []
                attachModuleComment
            Single  | [egroup] <- exports
                    , not (commentGroupHasComments egroup)
                    , [(export, _)] <- (cgItems egroup) -> do
                printSingleLineExportList conf [export]
                attachModuleComment
            Inline  | [] <- exports -> do
                printSingleLineExportList conf []
                attachModuleComment
            Inline  | [egroup] <- exports, not (commentGroupHasComments egroup) -> do
                wrapping
                    (printSingleLineExportList conf $ map fst $ cgItems egroup)
                    (do
                        attachOpenBracket
                        attachModuleComment
                        printMultiLineExportList conf exports)
            _ -> do
                attachOpenBracket
                attachModuleComment
                printMultiLineExportList conf exports

    putMaybeLineComment $ GHC.unLoc <$> mbWhereComment
  where
    attachModuleComment = putMaybeLineComment $ GHC.unLoc <$> mbModuleComment

    attachOpenBracket
        | openBracket conf == SameLine = putText " ("
        | otherwise                    = pure ()

printSingleLineExportList
    :: Config -> [GHC.LIE GHC.GhcPs] -> P ()
printSingleLineExportList conf exports = do
    space >> putText "("
    printExports exports
    putText ")" >> space >> putText "where"
  where
    printExports :: [GHC.LIE GHC.GhcPs] -> P ()
    printExports = \case
        []     -> pure ()
        [e]    -> putExport conf e
        (e:es) -> putExport conf e >> comma >> space >> printExports es

printMultiLineExportList
     :: Config
     -> [CommentGroup (GHC.LIE GHC.GhcPs)]
     -> P ()
printMultiLineExportList conf exports = do
    newline
    doIndent >> putText firstChar >> unless (null exports) space
    mapM_ printExport $ flagEnds exports
    when (null exports) $ newline >> doIndent
    putText ")" >> space >> putText "where"
  where
    printExport (CommentGroup {..}, firstGroup, _lastGroup) = do
        forM_ (flagEnds cgPrior) $ \(cmt, start, _end) -> do
            unless (firstGroup && start) $ space >> space
            putComment $ GHC.unLoc cmt
            newline >> doIndent

        forM_ (flagEnds cgItems) $ \((export, mbComment), start, _end) -> do
            if firstGroup && start then
                unless (null cgPrior) $ space >> space
            else
                comma >> space
            putExport conf export
            putMaybeLineComment $ GHC.unLoc <$> mbComment
            newline >> doIndent

    firstChar = case openBracket conf of
        SameLine -> " "
        NextLine -> "("

    doIndent = spaces (indent conf)

-- NOTE(jaspervdj): This code is almost the same as the import printing in
-- 'Imports' and should be merged.
putExport :: Config -> GHC.LIE GHC.GhcPs -> P ()
putExport conf = Imports.printImport (separateLists conf) . GHC.unLoc