{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NoImplicitPrelude #-}

module Language.Haskell.Brittany.Internal.Layouters.Module where

import qualified Data.Maybe
import qualified Data.Semigroup as Semigroup
import qualified Data.Text as Text
import GHC (AnnKeywordId(..), GenLocated(L), moduleNameString, unLoc)
import GHC.Hs
import qualified GHC.OldList as List
import Language.Haskell.Brittany.Internal.Config.Types
import Language.Haskell.Brittany.Internal.LayouterBasics
import Language.Haskell.Brittany.Internal.Layouters.IE
import Language.Haskell.Brittany.Internal.Layouters.Import
import Language.Haskell.Brittany.Internal.Prelude
import Language.Haskell.Brittany.Internal.PreludeUtils
import Language.Haskell.Brittany.Internal.Types
import Language.Haskell.GHC.ExactPrint as ExactPrint
import Language.Haskell.GHC.ExactPrint.Types
  (DeltaPos(..), commentContents, deltaRow)



layoutModule :: ToBriDoc' HsModule
layoutModule lmod@(L _ mod') = case mod' of
    -- Implicit module Main
  HsModule _ Nothing  _   imports _ _ _ -> do
    commentedImports <- transformToCommentedImport imports
    -- groupify commentedImports `forM_` tellDebugMessShow
    docLines (commentedImportsToDoc <$> sortCommentedImports commentedImports)
    -- sortedImports <- sortImports imports
    -- docLines $ [layoutImport y i | (y, i) <- sortedImports]
  HsModule _ (Just n) les imports _ _ _ -> do
    commentedImports <- transformToCommentedImport imports
    -- groupify commentedImports `forM_` tellDebugMessShow
    -- sortedImports <- sortImports imports
    let tn = Text.pack $ moduleNameString $ unLoc n
    allowSingleLineExportList <- mAsk
      <&> _conf_layout
      .>  _lconfig_allowSingleLineExportList
      .>  confUnpack
    -- the config should not prevent single-line layout when there is no
    -- export list
    let allowSingleLine = allowSingleLineExportList || Data.Maybe.isNothing les
    docLines
      $ docSeq
          [ docNodeAnnKW lmod Nothing docEmpty
             -- A pseudo node that serves merely to force documentation
             -- before the node
          , docNodeMoveToKWDP lmod AnnModule True $ runFilteredAlternative $ do
            addAlternativeCond allowSingleLine $
              docForceSingleline
                $ docSeq
                [ appSep $ docLit $ Text.pack "module"
                , appSep $ docLit tn
                , docWrapNode lmod $ appSep $ case les of
                  Nothing -> docEmpty
                  Just x  -> layoutLLIEs True KeepItemsUnsorted x
                , docSeparator
                , docLit $ Text.pack "where"
                ]
            addAlternative
              $ docLines
              [ docAddBaseY BrIndentRegular $ docPar
                (docSeq [appSep $ docLit $ Text.pack "module", docLit tn]
                )
                (docSeq [
                          docWrapNode lmod $ case les of
                           Nothing -> docEmpty
                           Just x  -> layoutLLIEs False KeepItemsUnsorted x
                        , docSeparator
                        , docLit $ Text.pack "where"
                        ]
                )
              ]
          ]
      : (commentedImportsToDoc <$> sortCommentedImports commentedImports) -- [layoutImport y i | (y, i) <- sortedImports]

data CommentedImport
  = EmptyLine
  | IndependentComment (Comment, DeltaPos)
  | ImportStatement ImportStatementRecord

instance Show CommentedImport where
  show = \case
    EmptyLine            -> "EmptyLine"
    IndependentComment _ -> "IndependentComment"
    ImportStatement r ->
      "ImportStatement " ++ show (length $ commentsBefore r) ++ " " ++ show
        (length $ commentsAfter r)

data ImportStatementRecord = ImportStatementRecord
  { commentsBefore :: [(Comment, DeltaPos)]
  , commentsAfter :: [(Comment, DeltaPos)]
  , importStatement :: ImportDecl GhcPs
  }

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

transformToCommentedImport
  :: [LImportDecl GhcPs] -> ToBriDocM [CommentedImport]
transformToCommentedImport is = do
  nodeWithAnnotations <- is `forM` \i@(L _ rawImport) -> do
    annotionMay <- astAnn i
    pure (annotionMay, rawImport)
  let
    convertComment (c, DP (y, x)) =
      replicate (y - 1) EmptyLine ++ [IndependentComment (c, DP (1, x))]
    accumF
      :: [(Comment, DeltaPos)]
      -> (Maybe Annotation, ImportDecl GhcPs)
      -> ([(Comment, DeltaPos)], [CommentedImport])
    accumF accConnectedComm (annMay, decl) = case annMay of
      Nothing ->
        ( []
        , [ ImportStatement ImportStatementRecord { commentsBefore  = []
                                                  , commentsAfter   = []
                                                  , importStatement = decl
                                                  }
          ]
        )
      Just ann ->
        let
          blanksBeforeImportDecl = deltaRow (annEntryDelta ann) - 1
          (newAccumulator, priorComments') =
            List.span ((== 0) . deltaRow . snd) (annPriorComments ann)
          go
            :: [(Comment, DeltaPos)]
            -> [(Comment, DeltaPos)]
            -> ([CommentedImport], [(Comment, DeltaPos)], Int)
          go acc []                       = ([], acc, 0)
          go acc [c1@(_, DP (y, _))] = ([], c1 : acc, y - 1)
          go acc (c1@(_, DP (1, _)) : xs) = go (c1 : acc) xs
          go acc ((c1, DP (y, x)) : xs) =
            ( (convertComment =<< xs) ++ replicate (y - 1) EmptyLine
            , (c1, DP (1, x)) : acc
            , 0
            )
          (convertedIndependentComments, beforeComments, initialBlanks) =
            if blanksBeforeImportDecl /= 0
              then (convertComment =<< priorComments', [], 0)
              else go [] (reverse priorComments')
        in
          ( newAccumulator
          , convertedIndependentComments
          ++ replicate (blanksBeforeImportDecl + initialBlanks) EmptyLine
          ++ [ ImportStatement ImportStatementRecord
                 { commentsBefore  = beforeComments
                 , commentsAfter   = accConnectedComm
                 , importStatement = decl
                 }
             ]
          )
  let (finalAcc, finalList) = mapAccumR accumF [] nodeWithAnnotations
  pure $ join $ (convertComment =<< finalAcc) : finalList

sortCommentedImports :: [CommentedImport] -> [CommentedImport]
sortCommentedImports =
  unpackImports . mergeGroups . map (fmap (sortGroups)) . groupify
 where
  unpackImports :: [CommentedImport] -> [CommentedImport]
  unpackImports xs = xs >>= \case
    l@EmptyLine            -> [l]
    l@IndependentComment{} -> [l]
    ImportStatement r ->
      map IndependentComment (commentsBefore r) ++ [ImportStatement r]
  mergeGroups
    :: [Either CommentedImport [ImportStatementRecord]] -> [CommentedImport]
  mergeGroups xs = xs >>= \case
    Left  x -> [x]
    Right y -> ImportStatement <$> y
  sortGroups :: [ImportStatementRecord] -> [ImportStatementRecord]
  sortGroups =
    List.sortOn (moduleNameString . unLoc . ideclName . importStatement)
  groupify
    :: [CommentedImport] -> [Either CommentedImport [ImportStatementRecord]]
  groupify cs = go [] cs
   where
    go [] = \case
      (l@EmptyLine            : rest) -> Left l : go [] rest
      (l@IndependentComment{} : rest) -> Left l : go [] rest
      (ImportStatement r      : rest) -> go [r] rest
      []                              -> []
    go acc = \case
      (l@EmptyLine : rest) -> Right (reverse acc) : Left l : go [] rest
      (l@IndependentComment{} : rest) ->
        Left l : Right (reverse acc) : go [] rest
      (ImportStatement r : rest) -> go (r : acc) rest
      []                         -> [Right (reverse acc)]

commentedImportsToDoc :: CommentedImport -> ToBriDocM BriDocNumbered
commentedImportsToDoc = \case
  EmptyLine -> docLitS ""
  IndependentComment c -> commentToDoc c
  ImportStatement r ->
    docSeq
      ( layoutImport (importStatement r)
      : map commentToDoc (commentsAfter r)
      )
 where
  commentToDoc (c, DP (_y, x)) = docLitS (replicate x ' ' ++ commentContents c)