{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module Language.Haskell.Brittany.Internal.ToBriDocTools where

import qualified Control.Monad.Writer.Strict   as Writer
import qualified Data.Char                     as Char
import           Data.Data
import qualified Data.Generics                 as SYB
import qualified Data.Sequence                 as Seq
import qualified Data.Text                     as Text
import           DataTreePrint
import           GHC                            ( EpAnn(EpAnn, EpAnnNotUsed)
                                                , EpAnnComments
                                                  ( EpaComments
                                                  , EpaCommentsBalanced
                                                  )
                                                , GenLocated(L)
                                                , LEpaComment
                                                , Located
                                                , LocatedA
                                                , moduleName
                                                , moduleNameString
                                                )
import qualified GHC
import           GHC.Data.FastString            ( FastString )
import qualified GHC.OldList                   as List
import           GHC.Parser.Annotation          ( AnnKeywordId(..) )
import           GHC.Types.Name                 ( getOccString )
import           GHC.Types.Name.Occurrence      ( occNameString )
import           GHC.Types.Name.Reader          ( RdrName(..) )
import qualified GHC.Types.SrcLoc              as GHC
import           GHC.Utils.Outputable           ( Outputable )
import           Data.Coerce                    ( Coercible )
import           Language.Haskell.Brittany.Internal.Components.BriDoc
import           Language.Haskell.Brittany.Internal.Config.Types
import           Language.Haskell.Brittany.Internal.Prelude
import           Language.Haskell.Brittany.Internal.Types
import           Language.Haskell.Brittany.Internal.Utils
import qualified Language.Haskell.GHC.ExactPrint
                                               as ExactPrint
import qualified Language.Haskell.GHC.ExactPrint.Utils
                                               as ExactPrint



-- | Use ExactPrint's output for this node; add a newly generated inline comment
-- at insertion position (meant to point out to the user that this node is
-- not handled by brittany yet). Useful when starting implementing new
-- syntactic constructs when children are not handled yet.
briDocByExact
  :: (ExactPrint.ExactPrint ast, Data ast)
  => LocatedA ast
  -> ToBriDocM BriDocNumbered
briDocByExact ast = do
  traceIfDumpConf
    "ast"
    _dconf_dump_ast_unknown
    (printTreeWithCustom 160 customLayouterF ast)
  mModify (+ connectedCommentCount ast)
  docExt ast True

-- | Use ExactPrint's output for this node.
-- Consider that for multi-line input, the indentation of the code produced
-- by ExactPrint might be different, and even incompatible with the indentation
-- of its surroundings as layouted by brittany. But there are safe uses of
-- this, e.g. for any top-level declarations.
briDocByExactNoComment
  :: (ExactPrint.ExactPrint (GenLocated l ast), Data ast, Data l)
  => GenLocated l ast
  -> ToBriDocM BriDocNumbered
briDocByExactNoComment ast = do
  traceIfDumpConf
    "ast"
    _dconf_dump_ast_unknown
    (printTreeWithCustom 160 customLayouterF ast)
  mModify (+ connectedCommentCount ast)
  docExt ast False

-- | Use ExactPrint's output for this node, presuming that this output does
-- not contain any newlines. If this property is not met, the semantics
-- depend on the @econf_AllowRiskyExactPrintUse@ config flag.
briDocByExactInlineOnly
  :: ( ExactPrint.ExactPrint (GHC.XRec GhcPs a)
     , Data (GHC.XRec GhcPs a)
     , Data a
     , Data (GHC.Anno a)
     , Outputable (GHC.Anno a)
     )
  => String
  -> GHC.XRec GhcPs a
  -> ToBriDocM (BriDocNumbered)
briDocByExactInlineOnly infoStr ast = do
  traceIfDumpConf "ast"
                  _dconf_dump_ast_unknown
                  (printTreeWithCustom 160 customLayouterF ast)
  let exactPrinted =
        dropWhile Text.null $ Text.lines $ Text.pack $ ExactPrint.exactPrint ast
  fallbackMode <-
    mAsk <&> _conf_errorHandling .> _econf_ExactPrintFallback .> confUnpack
  let exactPrintNode t = allocateNode $ BDExternal False t
  let errorAction = do
        mTell [ErrorUnknownNode infoStr ast]
        docLit
          $ Text.pack "{- BRITTANY ERROR UNHANDLED SYNTACTICAL CONSTRUCT -}"
  mModify (+ connectedCommentCount ast)
  case (fallbackMode, exactPrinted) of
    (ExactPrintFallbackModeNever, _  ) -> errorAction
    (_                          , [t]) -> exactPrintNode
      (Text.dropWhile Char.isSpace . Text.dropWhileEnd Char.isSpace $ t)
    (ExactPrintFallbackModeRisky, _) -> exactPrintNode
      (Text.unlines exactPrinted)
    _ -> errorAction

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

lrdrNameToText :: GenLocated l RdrName -> Text
lrdrNameToText (L _ n) = rdrNameToText n

class PrintRdrNameWithAnns l where
  printRdrNameWithAnns :: GenLocated l RdrName -> Text

instance PrintRdrNameWithAnns GHC.SrcSpanAnnN where
  printRdrNameWithAnns (L (GHC.SrcSpanAnn epAnn _) name) =
    case epAnn of
      EpAnn _ (GHC.NameAnn GHC.NameParens _ _ _ _) _ -> f "(" name ")"
      EpAnn _ (GHC.NameAnn GHC.NameParensHash _ _ _ _) _ -> f "(#" name "#)"
      EpAnn _ (GHC.NameAnn GHC.NameBackquotes _ _ _ _) _ -> f "`" name "`"
      EpAnn _ (GHC.NameAnn GHC.NameSquare _ _ _ _) _ -> f "[" name "]"
      EpAnn _ (GHC.NameAnnQuote _ _ _) _ -> f "'" name ""
      -- TODO92 There are way more possible constructors here
      -- see https://hackage.haskell.org/package/ghc-9.2.5/docs/GHC-Parser-Annotation.html#t:NameAnn
      EpAnn _ _ _ -> rdrNameToText name
      EpAnnNotUsed -> rdrNameToText name
   where
    f a b c = Text.pack a <> rdrNameToText b <> Text.pack c

lrdrNameToTextAnnGen
  :: (MonadMultiReader Config m, PrintRdrNameWithAnns l)
  => (Text -> Text)
  -> GenLocated l RdrName
  -> m Text
  -- TODO this doesn't need to be monadic. I am pretty sure it started of as
  -- a pure function, then at some point annotations were inspected
  -- (from reader) but now it is pure again.
  -- Leaving it as pseudo-monadic is harmless though (I think? Maybe I should
  -- check I don't force some mapM/sequence/… garbage at common callsides
    -- for this).
lrdrNameToTextAnnGen f ast = pure $ f $ printRdrNameWithAnns ast

lrdrNameToTextAnn
  :: (MonadMultiReader Config m, PrintRdrNameWithAnns l)
  => GenLocated l RdrName
  -> m Text
lrdrNameToTextAnn = lrdrNameToTextAnnGen id

lrdrNameToTextAnnTypeEqualityIsSpecial
  :: (MonadMultiReader Config m, PrintRdrNameWithAnns l)
  => GenLocated l RdrName
  -> m Text
lrdrNameToTextAnnTypeEqualityIsSpecial ast = do
  let
    f x = if x == Text.pack "Data.Type.Equality~"
      then Text.pack "~" -- rraaaahhh special casing rraaahhhhhh
      else x
  lrdrNameToTextAnnGen f ast

-- | Same as lrdrNameToTextAnnTypeEqualityIsSpecial, but also inspects
-- the annotations for a (parent) node for a tick to be added to the
-- literal.
-- Excessively long name to reflect on us having to work around such
-- excessively obscure special cases in the exactprint API.
lrdrNameToTextAnnTypeEqualityIsSpecialAndRespectTick
  :: (MonadMultiReader Config m, PrintRdrNameWithAnns l)
  => Located ast
  -> GenLocated l RdrName
  -> m Text
lrdrNameToTextAnnTypeEqualityIsSpecialAndRespectTick _ast1 ast2 = do
  -- TODO92
  -- hasQuote <- hasAnnKeyword ast1 AnnSimpleQuote
  x <- lrdrNameToTextAnn ast2
  let
    lit = if x == Text.pack "Data.Type.Equality~"
      then Text.pack "~" -- rraaaahhh special casing rraaahhhhhh
      else x
  return lit -- $ if hasQuote then Text.cons '\'' lit else lit

askIndent :: (MonadMultiReader Config m) => m Int
askIndent = confUnpack . _lconfig_indentAmount . _conf_layout <$> mAsk

-- TODO92 this is not filtering enough yet, see old code below
hasAnyCommentsBelow :: Data ast => ast -> Bool
hasAnyCommentsBelow =
  getAny . SYB.everything (<>) (SYB.mkQ (Any False) (\(_ :: LEpaComment) -> Any True))
-- -- | True if there are any comments that are
-- -- a) connected to any node below (in AST sense) the given node AND
-- -- b) after (in source code order) the node.
-- hasAnyCommentsBelow :: Data ast => GHC.Located ast -> ToBriDocM Bool
-- hasAnyCommentsBelow ast@(L l _) =
--   List.any (\(c, _) -> ExactPrint.commentIdentifier c > ExactPrint.Utils.rs l)
--     <$> astConnectedComments ast

-- extractRestComments
--   :: ExactPrint.Annotation -> [(ExactPrint.Comment, ExactPrint.DeltaPos)]
-- extractRestComments ann =
--   ExactPrint.annFollowingComments ann
--     ++ (ExactPrint.annsDP ann >>= \case
--          (ExactPrint.AnnComment com, dp) -> [(com, dp)]
--          _ -> []
--        )

-- filterAnns :: Data.Data.Data ast => ast -> ExactPrint.Anns -> ExactPrint.Anns
-- filterAnns ast = Map.filterWithKey (\k _ -> k `Set.member` foldedAnnKeys ast)

hasCommentsBetween
  :: Data ast
  => ast
  -> Maybe GHC.RealSrcSpan
  -> Maybe GHC.RealSrcSpan
  -> Bool
hasCommentsBetween ast left right = do
  getAny
    $ SYB.everything
      (<>)
      (SYB.mkQ
        (Any False)
        (\(L (GHC.Anchor pos _) _ :: LEpaComment) -> Any
          (  ( maybe True
                     (\l -> GHC.realSrcSpanStart pos >= GHC.realSrcSpanEnd l)
                     left
             )
          && (maybe True
                    (\l -> GHC.realSrcSpanEnd pos <= GHC.realSrcSpanStart l)
                    right
             )
          )
        )
      )
      ast

startsWithComments :: EpAnn a -> Bool
startsWithComments = \case
  EpAnnNotUsed                         -> False
  EpAnn (GHC.Anchor srcSpan _) _ comms -> case comms of
    EpaComments cs                 -> anyCheck cs
    EpaCommentsBalanced comms1 comms2 -> anyCheck comms1 || anyCheck comms2
   where
    anyCheck cs =
      any
        (\(L _ (GHC.EpaComment _ commSpan)) ->
          GHC.realSrcSpanStart srcSpan == GHC.realSrcSpanStart commSpan
        )
        cs


  -- mAnn <- astAnn ast
  -- let
  --   go1 [] = False
  --   go1 ((ExactPrint.G kw, _dp) : rest) | kw == leftKey = go2 rest
  --   go1 (_ : rest) = go1 rest
  --   go2 [] = False
  --   go2 ((ExactPrint.AnnComment _, _dp) : _rest) = True
  --   go2 ((ExactPrint.G kw, _dp) : _rest) | kw == rightKey = False
  --   go2 (_ : rest) = go2 rest
  -- case mAnn of
  --   Nothing -> pure False
  --   Just ann -> pure $ go1 $ ExactPrint.annsDP ann

-- | True if there are any comments that are connected to any node below (in AST
--   sense) the given node
hasAnyCommentsConnected :: (Data ann, Data ast) => GHC.GenLocated ann ast -> Bool
hasAnyCommentsConnected =
  getAny . SYB.everything (<>) (SYB.mkQ (Any False) (\(_ :: LEpaComment) -> Any True))

connectedCommentCount :: (Data ann, Data ast) => GHC.GenLocated ann ast -> CommentCounter
connectedCommentCount =
  getSum . SYB.everything (<>) (SYB.mkQ (Sum 0) (\(_ :: LEpaComment) -> Sum 1))

-- | True if there are any regular comments connected to any node below (in AST
--   sense) the given node
-- hasAnyRegularCommentsConnected :: GenLocated ann ast -> Bool
-- hasAnyRegularCommentsConnected ast =
--   any isRegularComment $ astConnectedComments ast

-- | Regular comments are comments that are actually "source code comments",
-- i.e. things that start with "--" or "{-". In contrast to comment-annotations
-- used by ghc-exactprint for capturing symbols (and their exact positioning).
--
-- Only the type instance layouter makes use of this filter currently, but
-- it might make sense to apply it more aggressively or make it the default -
-- I believe that most of the time we branch on the existence of comments, we
-- only care about "regular" comments. We simply did not need the distinction
-- because "irregular" comments are not that common outside of type/data decls.
-- isRegularComment :: (ExactPrint.Comment, ExactPrint.DeltaPos) -> Bool
-- isRegularComment = (== Nothing) . ExactPrint.Types.commentOrigin . fst

astConnectedComments
  :: (Data ann, Data ast)
  => GHC.GenLocated ann ast
  -> [LEpaComment]
astConnectedComments =
  SYB.listify (\(_ :: LEpaComment) -> True)
  -- anns <- filterAnns ast <$> mAsk
  -- pure $ extractAllComments =<< Map.elems anns
-- 
-- hasAnyCommentsPrior :: Data ast => GHC.Located ast -> ToBriDocM Bool
-- hasAnyCommentsPrior ast = astAnn ast <&> \case
--   Nothing -> False
--   Just (ExactPrint.Types.Ann _ priors _ _ _ _) -> not $ null priors

-- hasAnyRegularCommentsRest :: Data ast => GHC.Located ast -> ToBriDocM Bool
-- hasAnyRegularCommentsRest ast = astAnn ast <&> \case
--   Nothing -> False
--   Just ann -> any isRegularComment (extractRestComments ann)

-- hasAnnKeywordComment
--   :: GHC.LocatedA ast -> AnnKeywordId -> Bool
-- hasAnnKeywordComment (L (GHC.SrcSpanAnn ann _) _) annKeyword = False -- _ ann
  -- Nothing -> False
  -- Just ann -> any hasK (extractAllComments ann)
  -- where hasK = (== Just annKeyword) . ExactPrint.Types.commentOrigin . fst

-- new BriDoc stuff

allocateNode
  :: MonadMultiState NodeAllocIndex m => BriDocWrapped -> m BriDocNumbered
allocateNode bd = do
  i <- allocNodeIndex
  return (i, bd)

allocNodeIndex :: MonadMultiState NodeAllocIndex m => m Int
allocNodeIndex = do
  NodeAllocIndex i <- mGet
  mSet $ NodeAllocIndex (i + 1)
  return i

docEmpty :: ToBriDocM BriDocNumbered
docEmpty = allocateNode BDEmpty

docLit :: Text -> ToBriDocM BriDocNumbered
docLit t = allocateNode $ BDLit t

docLitS :: String -> ToBriDocM BriDocNumbered
docLitS s = allocateNode $ BDLit $ Text.pack s

docExt
  :: (ExactPrint.ExactPrint (GenLocated l ast))
  => GenLocated l ast
  -> Bool
  -> ToBriDocM BriDocNumbered
docExt x shouldAddComment = allocateNode $ BDExternal
    -- (ExactPrint.Types.mkAnnKey x)
    -- (foldedAnnKeys x)
  shouldAddComment
  (Text.pack
  $ List.dropWhile ((==) '\n')
  $ ExactPrint.exactPrint
  $ ExactPrint.makeDeltaAst x
  )

docAlt :: [ToBriDocM BriDocNumbered] -> ToBriDocM BriDocNumbered
docAlt l = allocateNode . BDAlt =<< sequence l

newtype CollectAltM a = CollectAltM (Writer.Writer [ToBriDocM BriDocNumbered] a)
  deriving (Functor, Applicative, Monad)

addAlternativeCond :: Bool -> ToBriDocM BriDocNumbered -> CollectAltM ()
addAlternativeCond cond doc = when cond (addAlternative doc)

addAlternative :: ToBriDocM BriDocNumbered -> CollectAltM ()
addAlternative = CollectAltM . Writer.tell . (: [])

runFilteredAlternative :: CollectAltM () -> ToBriDocM BriDocNumbered
runFilteredAlternative (CollectAltM action) = docAlt $ Writer.execWriter action


docSeq :: [ToBriDocM BriDocNumbered] -> ToBriDocM BriDocNumbered
docSeq [] = docEmpty
docSeq l = allocateNode . BDSeq =<< sequence l

docSeqSep :: [ToBriDocM BriDocNumbered] -> ToBriDocM BriDocNumbered
docSeqSep = docSeq . List.intersperse docSeparator

docLines :: [ToBriDocM BriDocNumbered] -> ToBriDocM BriDocNumbered
docLines l = allocateNode . BDLines =<< sequence l

docCols :: ColSig -> [ToBriDocM BriDocNumbered] -> ToBriDocM BriDocNumbered
docCols sig l = allocateNode . BDCols sig =<< sequence l

docAddBaseY :: BrIndent -> ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered
docAddBaseY ind bdm = allocateNode . BDAddBaseY ind =<< bdm

docSetBaseY :: ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered
docSetBaseY bdm = do
  bd <- bdm
  -- the order here is important so that these two nodes can be treated
  -- properly over at `transformAlts`.
  allocateNode $ BDBaseYPushCur bd

docSetIndentLevel :: ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered
docSetIndentLevel bdm = do
  bd <- bdm
  n1 <- allocateNode $ BDIndentLevelPushCur bd
  n2 <- allocateNode $ BDIndentLevelPop n1
  return n2

docSetBaseAndIndent :: ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered
docSetBaseAndIndent = docSetBaseY . docSetIndentLevel

docSeparator :: ToBriDocM BriDocNumbered
docSeparator = allocateNode BDSeparator

docNonBottomSpacing :: ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered
docNonBottomSpacing bdm = allocateNode . BDForceAlt (NonBottomSpacing False) =<< bdm

docNonBottomSpacingS :: ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered
docNonBottomSpacingS bdm = allocateNode . BDForceAlt (NonBottomSpacing True) =<< bdm

docSetParSpacing :: ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered
docSetParSpacing bdm = allocateNode . BDForceAlt SetParSpacing =<< bdm

docForceParSpacing :: ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered
docForceParSpacing bdm = allocateNode . BDForceAlt ForceParSpacing =<< bdm

docForceZeroAdd :: ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered
docForceZeroAdd bdm = allocateNode . BDForceAlt ForceZeroAdd =<< bdm

docDebug :: String -> ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered
docDebug s bdm = allocateNode . BDDebug s =<< bdm

appSep :: ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered
appSep x = docSeq [x, docSeparator]

docCommaSep :: ToBriDocM BriDocNumbered
docCommaSep = appSep $ docLit $ Text.pack ","

docParenLSep :: ToBriDocM BriDocNumbered
docParenLSep = appSep docParenL

-- TODO: we don't make consistent use of these (yet). However, I think the
-- most readable approach overall might be something else: define
-- `lit = docLit . Text.pack` and `prepSep = docSeq [docSeparator, x]`.
-- I think those two would make the usage most readable.
-- lit "("  and  appSep (lit "(")  are understandable and short without
-- introducing a new top-level binding for all types of parentheses.
docParenL :: ToBriDocM BriDocNumbered
docParenL = docLit $ Text.pack "("

docParenR :: ToBriDocM BriDocNumbered
docParenR = docLit $ Text.pack ")"

docParenHashLSep :: ToBriDocM BriDocNumbered
docParenHashLSep = docSeq [docLit $ Text.pack "(#", docSeparator]

docParenHashRSep :: ToBriDocM BriDocNumbered
docParenHashRSep = docSeq [docSeparator, docLit $ Text.pack "#)"]

docBracketL :: ToBriDocM BriDocNumbered
docBracketL = docLit $ Text.pack "["

docBracketR :: ToBriDocM BriDocNumbered
docBracketR = docLit $ Text.pack "]"


docTick :: ToBriDocM BriDocNumbered
docTick = docLit $ Text.pack "'"


docPar
  :: ToBriDocM BriDocNumbered
  -> ToBriDocM BriDocNumbered
  -> ToBriDocM BriDocNumbered
docPar lineM indentedM = do
  line <- lineM
  indented <- indentedM
  allocateNode $ BDPar BrIndentNone line indented

docForceSingleline :: ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered
docForceSingleline bdm = allocateNode . BDForceAlt ForceSingleline =<< bdm

docForceMultiline :: ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered
docForceMultiline bdm = allocateNode . BDForceAlt ForceMultiline =<< bdm

docEnsureIndent
  :: BrIndent -> ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered
docEnsureIndent ind mbd = mbd >>= \bd -> allocateNode $ BDEnsureIndent ind bd

docAddEntryDelta :: GHC.DeltaPos -> ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered
docAddEntryDelta dp bdm = do
  bd <- bdm
  allocateNode (BDEntryDelta dp bd)

docFlushRemaining :: FastString -> ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered
docFlushRemaining fileThing = docFlushCommsPost False
  (Just $ GHC.mkRealSrcLoc fileThing 999999 999999)

-- CLASS DocHandleComms --------------------------------------------------------

class DocHandleComms ann a where
  docHandleComms :: HasCallStack => ann -> a -> a

instance DocHandleComms [LEpaComment] (ToBriDocM BriDocNumbered) where
  docHandleComms comms bdm = do
    bd <- bdm
    i1 <- allocNodeIndex
    pure (i1, BDQueueComments comms bd)

instance DocHandleComms (EpAnn a) (ToBriDocM BriDocNumbered) where
  docHandleComms epAnn bdm = case epAnn of
    EpAnn anch _ (EpaComments []) -> do
      bd <- bdm
      i1 <- allocNodeIndex
      pure
        (i1, BDFlushCommentsPrior (GHC.realSrcSpanStart $ GHC.anchor anch) bd)
    EpAnn anch _ (EpaComments comms) -> do
      bd <- bdm
      i1 <- allocNodeIndex
      i2 <- allocNodeIndex
      pure
        ( i1
        , BDFlushCommentsPrior (GHC.realSrcSpanStart $ GHC.anchor anch)
                                (i2, BDQueueComments (reverse comms) bd)
        )
    EpAnn anch _ (EpaCommentsBalanced commsB commsA) -> do
      bd <- bdm
      i1 <- allocNodeIndex
      i2 <- allocNodeIndex
      pure
        ( i1
        , BDQueueComments
          (reverse commsB ++ reverse commsA)
          ( i2
          , BDFlushCommentsPrior (GHC.realSrcSpanStart $ GHC.anchor anch) bd
          )
        )
    EpAnnNotUsed -> bdm

instance DocHandleComms (Maybe (EpAnn a)) (ToBriDocM BriDocNumbered) where
  docHandleComms Nothing = id
  docHandleComms (Just epAnn) = docHandleComms epAnn

instance DocHandleComms (GHC.RealSrcSpan) (ToBriDocM BriDocNumbered) where
  docHandleComms loc = docHandleComms (GHC.realSrcSpanStart loc)

instance DocHandleComms (GHC.RealSrcLoc) (ToBriDocM BriDocNumbered) where
  docHandleComms loc bdm = do
    bd <- bdm
    i1 <- allocNodeIndex
    pure (i1, BDFlushCommentsPrior loc bd)

instance DocHandleComms (Maybe GHC.RealSrcSpan) (ToBriDocM BriDocNumbered) where
  docHandleComms Nothing bdm = bdm
  docHandleComms (Just loc) bdm = docHandleComms loc bdm

instance DocHandleComms (Maybe GHC.RealSrcLoc) (ToBriDocM BriDocNumbered) where
  docHandleComms Nothing bdm = bdm
  docHandleComms (Just loc) bdm = docHandleComms loc bdm

instance DocHandleComms (GHC.SrcLoc) (ToBriDocM BriDocNumbered) where
  docHandleComms (GHC.RealSrcLoc loc _) bdm = docHandleComms loc bdm
  docHandleComms (GHC.UnhelpfulLoc _) bdm = bdm

instance DocHandleComms (GHC.LocatedA ast) (ToBriDocM BriDocNumbered) where
  docHandleComms (L (GHC.SrcSpanAnn epAnn span) _) bdm = case span of
    GHC.RealSrcSpan s _ -> docHandleComms s $ docHandleComms epAnn bdm
    GHC.UnhelpfulSpan _ -> bdm

instance DocHandleComms (GHC.LocatedL ast) (ToBriDocM BriDocNumbered) where
  docHandleComms (L (GHC.SrcSpanAnn epAnn span) _) bdm = case span of
    GHC.RealSrcSpan s _ -> docHandleComms s $ docHandleComms epAnn bdm
    GHC.UnhelpfulSpan _ -> bdm

instance DocHandleComms (GHC.LocatedC ast) (ToBriDocM BriDocNumbered) where
  docHandleComms (L (GHC.SrcSpanAnn epAnn span) _) bdm = case span of
    GHC.RealSrcSpan s _ -> docHandleComms s $ docHandleComms epAnn bdm
    GHC.UnhelpfulSpan _ -> bdm

instance DocHandleComms (GHC.LocatedN ast) (ToBriDocM BriDocNumbered) where
  docHandleComms (L (GHC.SrcSpanAnn epAnn span) _) bdm = case span of
    GHC.RealSrcSpan s _ -> docHandleComms s $ docHandleComms epAnn bdm
    GHC.UnhelpfulSpan _ -> bdm


instance DocHandleComms ann (ToBriDocM BriDocNumbered)
      => DocHandleComms ann (ToBriDocM [BriDocNumbered]) where
  docHandleComms ann bdm = do
    x <- bdm
    case x of
      [] -> error "docHandleComms empty list"
        -- TODO92
        -- do
        -- el <- docHandleComms ann docEmpty
        -- pure [el]
      (bd1:bdR) -> do
        bd1' <- docHandleComms ann (pure bd1)
        pure (bd1':bdR)

instance DocHandleComms ann (ToBriDocM BriDocNumbered)
      => DocHandleComms ann (ToBriDocM (Seq BriDocNumbered))
 where
  docHandleComms ast bdsm = do
    bds <- bdsm
    case Seq.viewl bds of
      Seq.EmptyL -> return Seq.empty -- TODO: this might be bad. maybe. then again, not really. well.
      bd1 Seq.:< rest -> do
        bd1' <- docHandleComms ast (return bd1)
        return $ bd1' Seq.<| rest


instance DocHandleComms ann (ToBriDocM BriDocNumbered)
      => DocHandleComms ann [ToBriDocM BriDocNumbered] where
  docHandleComms ann bdms = do
    case bdms of
      [] -> error "docHandleComms empty list"
       -- [docHandleComms ann docEmpty]
      (bd1:bdR) -> (docHandleComms ann bd1:bdR)

instance DocHandleComms GHC.EpaLocation (ToBriDocM BriDocNumbered) where
  docHandleComms loc bdm = docHandleComms (GHC.epaLocationRealSrcSpan loc) bdm

instance DocHandleComms GHC.SrcSpan (ToBriDocM BriDocNumbered) where
  docHandleComms (GHC.RealSrcSpan s _) bdm = docHandleComms s bdm
  docHandleComms (GHC.UnhelpfulSpan _) bdm = bdm


-- CLASS ObtainAnnPos ----------------------------------------------------------

class ObtainAnnPos key ann where
  obtainAnnPos :: ann -> key -> Maybe GHC.RealSrcSpan

instance ObtainAnnPos AnnKeywordId GHC.AddEpAnn where
  obtainAnnPos (GHC.AddEpAnn eKW loc) kw = if eKW == kw
    then Just (GHC.epaLocationRealSrcSpan loc)
    else Nothing

instance ObtainAnnPos AnnKeywordId (EpAnn GHC.AnnsModule) where
  obtainAnnPos = \case
    EpAnnNotUsed -> \_kw -> Nothing
    EpAnn _ (GHC.AnnsModule l annList) _ -> \kw ->
      obtainAnnPos l kw <|> obtainAnnPos annList kw

instance ObtainAnnPos AnnKeywordId (Maybe GHC.AddEpAnn) where
  obtainAnnPos Nothing _ = Nothing
  obtainAnnPos (Just addEpAnn) kw = obtainAnnPos addEpAnn kw

instance ObtainAnnPos AnnKeywordId [GHC.AddEpAnn] where
  obtainAnnPos list kw =
    case [ loc | GHC.AddEpAnn eKW loc <- list, eKW == kw ] of
      []   -> Nothing
      locs -> Just (GHC.epaLocationRealSrcSpan $ minimum locs)

instance ObtainAnnPos AnnKeywordId (EpAnn [GHC.AddEpAnn]) where
  obtainAnnPos EpAnnNotUsed _kw = Nothing
  obtainAnnPos (EpAnn _ list _) kw = obtainAnnPos list kw

instance ObtainAnnPos AnnKeywordId (EpAnn GHC.AnnList) where
  obtainAnnPos = \case
    EpAnnNotUsed -> \_kw -> Nothing
    EpAnn _ annList _ -> \kw -> obtainAnnPos annList kw

instance ObtainAnnPos AnnKeywordId GHC.AnnList where
  obtainAnnPos (GHC.AnnList _ op cl addEpAnn _) kw =
    obtainAnnPos op kw <|> obtainAnnPos cl kw <|> obtainAnnPos addEpAnn kw

instance ObtainAnnPos AnnKeywordId (EpAnn GHC.GrhsAnn) where
  obtainAnnPos = \case
    EpAnn _ (GHC.GrhsAnn _ addEpAnn) _ -> obtainAnnPos addEpAnn
    EpAnnNotUsed                       -> \_kw -> Nothing

instance ObtainAnnPos AnnKeywordId GHC.SrcSpanAnnL where
  obtainAnnPos = \case
    GHC.SrcSpanAnn epAnn _ -> obtainAnnPos epAnn

class ObtainAnnPos AnnKeywordId ann => ObtainAnnDeltaPos ann where
  obtainAnnDeltaPos :: ann -> AnnKeywordId -> Maybe GHC.DeltaPos

instance ObtainAnnDeltaPos (EpAnn GHC.AnnsModule) where
  obtainAnnDeltaPos = \case
    EpAnnNotUsed                          -> \_kw -> Nothing
    EpAnn _ (GHC.AnnsModule l annList) epaComms -> \kw -> do
      loc <- GHC.realSrcSpanStart
        <$> (obtainAnnPos l kw <|> obtainAnnPos annList kw)
      let pos = (GHC.srcLocLine loc, GHC.srcLocCol loc)
      pure $ ExactPrint.pos2delta
        (maximum $ (1, 1) :
          [ ExactPrint.ss2posEnd $ GHC.anchor anch
          | L anch _ <- case epaComms of
            EpaCommentsBalanced cs1 cs2 -> cs1 ++ cs2
            EpaComments cs              -> cs
          , let compPos = ExactPrint.ss2posEnd (GHC.anchor anch)
          , compPos <= pos
          ]
        )
        pos

class DocFlushCommsPost ann a where
  docFlushCommsPost :: Bool -> ann -> a -> a

instance DocFlushCommsPost (Maybe GHC.RealSrcLoc) (ToBriDocM BriDocNumbered) where
  docFlushCommsPost shouldMark = \case
    Nothing   -> id
    Just loc -> \bdm -> do
      i1 <- allocNodeIndex
      bd <- bdm
      pure (i1, BDFlushCommentsPost loc shouldMark bd)

instance DocFlushCommsPost (Maybe GHC.RealSrcSpan) (ToBriDocM BriDocNumbered) where
  docFlushCommsPost shouldMark = \case
    Nothing   -> id
    Just loc -> \bdm -> do
      i1 <- allocNodeIndex
      bd <- bdm
      pure (i1, BDFlushCommentsPost (GHC.realSrcSpanEnd loc) shouldMark bd)

instance DocFlushCommsPost GHC.EpaLocation (ToBriDocM BriDocNumbered) where
  docFlushCommsPost shouldMark epaLocation =
    docFlushCommsPost shouldMark (Just $ GHC.epaLocationRealSrcSpan epaLocation)

instance DocFlushCommsPost ann (ToBriDocM BriDocNumbered)
      => DocFlushCommsPost ann (ToBriDocM [BriDocNumbered]) where
  docFlushCommsPost shouldMark loc bdm = do
    bds <- bdm
    case bds of
      [] -> do
        e <- docFlushCommsPost shouldMark loc docEmpty
        pure [e]
      _ -> do
        e <- docFlushCommsPost shouldMark loc (pure $ List.last bds)
        pure (List.init bds ++ [e])

instance DocFlushCommsPost (Maybe GHC.RealSrcLoc) a
      => DocFlushCommsPost (LocatedA ast) a where
  docFlushCommsPost shouldMark (L ann _) =
    docFlushCommsPost shouldMark $ case GHC.locA ann of
      GHC.RealSrcSpan s _ -> Just $ GHC.realSrcSpanEnd s
      GHC.UnhelpfulSpan{} -> Nothing

unknownNodeError
  :: (Data a, Data (GHC.Anno a), Outputable (GHC.Anno a))
  => String
  -> GHC.XRec GhcPs a
  -> ToBriDocM BriDocNumbered
unknownNodeError infoStr ast = do
  mTell [ErrorUnknownNode infoStr ast]
  docLit $ Text.pack "{- BRITTANY ERROR UNHANDLED SYNTACTICAL CONSTRUCT -}"

spacifyDocs :: [ToBriDocM BriDocNumbered] -> [ToBriDocM BriDocNumbered]
spacifyDocs [] = []
spacifyDocs ds = fmap appSep (List.init ds) ++ [List.last ds]

shareDoc :: ToBriDocM a -> ToBriDocM (ToBriDocM a)
shareDoc = fmap pure


obtainListElemStartCommaLocs
  :: LocatedA ast -> (Maybe GHC.RealSrcLoc, Maybe GHC.RealSrcLoc)
obtainListElemStartCommaLocs = \case
  L (GHC.SrcSpanAnn elemEpAnn _) _ -> case elemEpAnn of
    EpAnn anch (GHC.AnnListItem [item]) _ ->
      ( Just $ GHC.realSrcSpanStart $ GHC.anchor anch
        -- yes, we want `realSrcSpanStart span2` here, but have it flow
        -- to the end of bd. We want any comments _before_ the _start_
        -- of the comma to be inserted _after_ the element.
      , Just $ GHC.realSrcSpanStart $ case item of
        GHC.AddCommaAnn   span -> GHC.epaLocationRealSrcSpan span
        GHC.AddSemiAnn    span -> GHC.epaLocationRealSrcSpan span
        GHC.AddVbarAnn    span -> GHC.epaLocationRealSrcSpan span
        GHC.AddRarrowAnn  span -> GHC.epaLocationRealSrcSpan span
        GHC.AddRarrowAnnU span -> GHC.epaLocationRealSrcSpan span
        GHC.AddLollyAnnU  span -> GHC.epaLocationRealSrcSpan span
      )
    EpAnn anch _ _ -> (Just $ GHC.realSrcSpanStart $ GHC.anchor anch, Nothing)
    EpAnnNotUsed   -> (Nothing, Nothing)


docHandleListElemComms
  :: (LocatedA ast -> ToBriDocM BriDocNumbered)
  -> LocatedA ast
  -> ToBriDocM BriDocNumbered
docHandleListElemComms layouter e = case obtainListElemStartCommaLocs e of
  (posStart, posComma) ->
    docHandleComms posStart $ docFlushCommsPost True posComma $ layouter e

docHandleListElemCommsProperPost
  :: (LocatedA ast -> ToBriDocM BriDocNumbered)
  -> [LocatedA ast]
  -> ToBriDocM
       [(Maybe GHC.RealSrcLoc, LocatedA ast, ToBriDocM BriDocNumbered)]
docHandleListElemCommsProperPost layouter es = case es of
  []          -> pure []
  (e1 : rest) -> case obtainListElemStartCommaLocs e1 of
    (posStart, posComma) -> do
      res <- go posComma rest
      pure
        $ ( Nothing
          , e1
          , docHandleComms posStart $ layouter e1
          )
        : res
 where
  go _intoComma []          = pure []
  go intoComma  (e1 : rest) = case obtainListElemStartCommaLocs e1 of
    (posStart, posComma) -> do
      res <- go posComma rest
      pure
        $ ( intoComma
          , e1
          , docHandleComms posStart $ layouter e1
          )
        : res

epaLocationRealSrcSpanStart :: GHC.EpaLocation -> GHC.RealSrcLoc
epaLocationRealSrcSpanStart = GHC.realSrcSpanStart . GHC.epaLocationRealSrcSpan

askLayoutConf
  :: Coercible a (ConfUnpacked a)
  => (CLayoutConfig Identity -> Identity a)
  -> ToBriDocM (ConfUnpacked a)
askLayoutConf f = mAsk <&> _conf_layout .> f .> confUnpack