{-# LANGUAGE NoImplicitPrelude #-}

module Language.Haskell.Brittany.Internal.ToBriDoc.OpTree where

import qualified Data.Text                     as Text
import           GHC                            ( GenLocated(L) )
import           GHC.Hs
import           GHC.Types.Fixity               ( Fixity(Fixity) )
import qualified GHC.Types.SrcLoc              as GHC

import           Language.Haskell.Brittany.Internal.Components.BriDoc
import           Language.Haskell.Brittany.Internal.Components.OpTree
import           Language.Haskell.Brittany.Internal.Config.Types
import           Language.Haskell.Brittany.Internal.Prelude
import           Language.Haskell.Brittany.Internal.S3_ToBriDocTools
import           Language.Haskell.Brittany.Internal.Types
import           Language.Haskell.Brittany.Internal.Utils


gatherOpTreeE
  :: Bool
  -> Bool
  -> (ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered)
  -> Maybe GHC.RealSrcLoc
  -> Maybe GHC.RealSrcLoc
  -> [(ToBriDocM BriDocNumbered, ToBriDocM BriDocNumbered)]
  -> LHsExpr GhcPs
  -> ToBriDocM (OpTree, Bool)
gatherOpTreeE hasParen hasComms commWrap locOpen locClose opExprList = \case
  (L _ (OpApp epAnn l1 op1 r1)) ->
    gatherOpTreeE
      hasParen
      (hasComms || hasAnyCommentsBelow epAnn)
      commWrap
      locOpen
      locClose
      ( ( docHandleComms epAnn $ callLayouter layout_expr op1
        , callLayouter layout_expr r1
        )
      : opExprList
      )
      l1
  (L _ (HsPar epAnn inner)) -> do
    let AnnParen _ spanOpen spanClose = anns epAnn
    let mergePoses locMay span = case locMay of
          Nothing  -> Just (epaLocationRealSrcSpanStart span)
          Just loc -> Just (min (epaLocationRealSrcSpanStart span) loc)
    (innerTree, innerHasComms) <-
      gatherOpTreeE True
                    (hasComms || hasAnyCommentsBelow epAnn)
                    (commWrap . docHandleComms epAnn)
                    (mergePoses locOpen spanOpen)
                    (mergePoses locClose spanClose)
                    []
                    inner
    if null opExprList
      then pure (innerTree, innerHasComms)
      else do
        numberedRights <- opExprList `forM` \(x, y) -> do
          x' <- x
          y' <- y
          pure (x', y')
        pure
          $ ( OpUnknown hasParen locOpen locClose innerTree numberedRights
            , innerHasComms
            )
  final -> do
    numberedLeft   <- commWrap $ callLayouter layout_expr final
    numberedRights <- opExprList `forM` \(x, y) -> do
      x' <- x
      y' <- y
      pure (x', y')
    pure
      $ ( OpUnknown hasParen
                    locOpen
                    locClose
                    (OpLeaf $ numberedLeft)
                    numberedRights
        , hasComms
        )

gatherOpTreeT
  :: Bool
  -> Bool
  -> (ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered)
  -> Maybe GHC.RealSrcLoc
  -> Maybe GHC.RealSrcLoc
  -> [(ToBriDocM BriDocNumbered, ToBriDocM BriDocNumbered)]
  -> LHsType GhcPs
  -> ToBriDocM (OpTree, Bool)
gatherOpTreeT hasParen hasComms commWrap locOpen locClose opExprList = \case
  (L _ (HsOpTy NoExtField l1 op1 r1)) ->
    gatherOpTreeT
      hasParen
      hasComms
      commWrap
      locOpen
      locClose
      ((docLit $ printRdrNameWithAnns op1, callLayouter layout_type r1) : opExprList)
      l1
  (L _ (HsParTy epAnn inner)) -> do
    let AnnParen _ spanOpen spanClose = anns epAnn
    let mergePoses locMay span = case locMay of
          Nothing  -> Just (epaLocationRealSrcSpanStart span)
          Just loc -> Just (min (epaLocationRealSrcSpanStart span) loc)
    (innerTree, innerHasComms) <-
      gatherOpTreeT True
                    (hasComms || hasAnyCommentsBelow epAnn)
                    (commWrap . docHandleComms epAnn)
                    (mergePoses locOpen spanOpen)
                    (mergePoses locClose spanClose)
                    []
                    inner
    if null opExprList
      then pure (innerTree, innerHasComms)
      else do
        numberedRights <- opExprList `forM` \(x, y) -> do
          x' <- x
          y' <- y
          pure (x', y')
        pure
          $ ( OpUnknown hasParen locOpen locClose innerTree numberedRights
            , innerHasComms
            )
  final -> do
    numberedLeft   <- commWrap $ callLayouter layout_type final
    numberedRights <- opExprList `forM` \(x, y) -> do
      x' <- x
      y' <- y
      pure (x', y')
    pure
      $ ( OpUnknown hasParen
                    locOpen
                    locClose
                    (OpLeaf $ numberedLeft)
                    numberedRights
        , hasComms
        )

processOpTree :: (OpTree, Bool) -> ToBriDocM BriDocNumbered
processOpTree (unknownTree, hasComments) = do
  enabled          <- askLayoutConf _lconfig_fixityAwareOps
  refactorMode     <- askLayoutConf _lconfig_operatorParenthesisRefactorMode
  allowOpUnqualify <- askLayoutConf _lconfig_operatorAllowUnqualify
  let (warns, balancedTree) = if enabled
        then balanceOpTree allowOpUnqualify unknownTree
        else ([], unknownTree)
  mTell warns
  let processedTree = case refactorMode of
        PRMKeep     -> balancedTree
        PRMMinimize -> remSuperfluousParens 11 balancedTree
        PRMMaximize -> addAllParens False balancedTree
  -- tellDebugMess $ displayOpTree balancedTree
  -- tellDebugMess $ displayOpTree processedTree
  layoutOpTree (not hasComments) processedTree

layoutOpTree :: Bool -> OpTree -> ToBriDocM BriDocNumbered
layoutOpTree allowSinglelinePar = \case
  OpUnknown hasParen locO locC leftTree docOps -> do
    let sharedOps = fmap (\(a, b) -> (pure a, pure b)) docOps
    leftDoc <- layoutOpTree True leftTree
    coreAlternative hasParen
                    locO
                    locC
                    Nothing
                    (pure leftDoc)
                    sharedOps
                    sharedOps
                    docForceSingleline
  OpKnown hasParen locO locC fixity treeL docOps -> do
    let Fixity _ _prec _ = fixity
    docL <- shareDoc $ layoutOpTree True treeL
    let flattenList ops = case ops of
          []           -> pure []
          [(op, tree)] -> case treeL of
            OpLeaf{} -> flattenInner op tree
            _        -> do
              treeDoc <- shareDoc $ layoutOpTree True tree
              pure [(pure op, treeDoc)]
          ((op1, tree1@OpLeaf{}) : tR) -> do
            tree1Doc    <- shareDoc $ layoutOpTree True tree1
            flattenRest <- flattenList tR
            pure $ (pure op1, tree1Doc) : flattenRest
          _ -> simpleTransform ops
        flattenInner op = \case
          OpKnown False _ _ _ innerL innerOps | isPrec0 fixity -> do
            flattenList ((op, innerL) : innerOps)
          tree -> do
            treeDoc <- shareDoc $ layoutOpTree True tree
            pure [(pure op, treeDoc)]
        simpleTransform
          :: [(BriDocNumbered, OpTree)]
          -> ToBriDocM [(ToBriDocM BriDocNumbered, ToBriDocM BriDocNumbered)]
        simpleTransform = mapM $ \(op, subTree) -> do
          subDoc <- shareDoc $ layoutOpTree True subTree
          pure (pure op, subDoc)
    sharedOpsFlat <- flattenList docOps
    sharedOps     <- simpleTransform docOps
    coreAlternative hasParen
                    locO
                    locC
                    (Just fixity)
                    docL
                    sharedOps
                    sharedOpsFlat
                    docForceParSpacing
  OpLeaf l -> pure l
 where
  isPrec0 = \case
    Fixity _ prec _ -> prec == 0
  coreAlternative
    :: Bool
    -> Maybe GHC.RealSrcLoc
    -> Maybe GHC.RealSrcLoc
    -> Maybe Fixity
    -> ToBriDocM BriDocNumbered
    -> [(ToBriDocM BriDocNumbered, ToBriDocM BriDocNumbered)]
    -> [(ToBriDocM BriDocNumbered, ToBriDocM BriDocNumbered)]
    -> (ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered)
    -> ToBriDocM BriDocNumbered
  coreAlternative hasParen locO locC fixity docL sharedOps sharedOpsFlat lastWrap
    = do
      let wrapParenIfSl x inner = if x then wrapParenSl inner else inner
          wrapParenSl inner = docAlt
            [ docSeq
              [ docLit $ Text.pack "("
              , docHandleComms locO $ docForceSingleline inner
              , docHandleComms locC $ docLit $ Text.pack ")"
              ]
            , docLines
              [ docSeq [docLitS "(", docHandleComms locO inner]
              , docHandleComms locC $ docLit $ Text.pack ")"
              ]
            ]
          wrapParenMlIf x innerHead innerLines = if x
            then wrapParenMl innerHead innerLines
            else docPar innerHead (docLines innerLines)
          wrapParenMl innerHead innerLines = docSetBaseY $ docLines
            (  [ docCols
                   ColOpPrefix
                   [ appSep $ docLit $ Text.pack "("
                   , docHandleComms locO $ innerHead
                   ]
               ]
            ++ innerLines
            ++ [docHandleComms locC $ docLit $ Text.pack ")"]
            )

      configAllowsParInsert <- askLayoutConf _lconfig_fixityBasedAddAlignParens
      let allowParIns = configAllowsParInsert && case fixity of
            Nothing                -> False
            Just (Fixity _ prec _) -> prec > 0

      runFilteredAlternative $ do
        -- > one + two + three
        -- or
        -- > one + two + case x of
        -- >   _ -> three
        addAlternativeCond allowSinglelinePar $ wrapParenIfSl hasParen $ docSeq
          ([docForceSingleline docL] ++ case splitFirstLast sharedOps of
            FirstLastEmpty -> []
            FirstLastSingleton (od, ed) ->
              [docSeparator, docForceSingleline od, docSeparator, lastWrap ed]
            FirstLast (od1, ed1) ems (odN, edN) ->
              (  [ docSeparator
                 , docForceSingleline od1
                 , docSeparator
                 , docForceSingleline ed1
                 ]
              ++ join
                   [ [ docSeparator
                     , docForceSingleline od
                     , docSeparator
                     , docForceSingleline ed
                     ]
                   | (od, ed) <- ems
                   ]
              ++ [ docSeparator
                 , docForceSingleline odN
                 , docSeparator
                 , lastWrap edN
                 ]
              )
          )
        -- this case rather leads to some unfortunate layouting than to anything
        -- useful; disabling for now. (it interfers with cols stuff.)
        addAlternativeCond (not hasParen) $ docPar
          (docHandleComms locO $ docForceSingleline $ docL)
          (docFlushCommsPost False locC $ docLines $ sharedOps <&> \(od, ed) ->
            docCols ColOpPrefix [appSep od, docForceSingleline ed]
          )
        -- > ( one
        -- > + two
        -- > + three
        -- > )
        addAlternativeCond (allowParIns && not hasParen)
          $ docForceZeroAdd
          $ wrapParenMl
              (docSetBaseY docL)
              (sharedOpsFlat <&> \(od, ed) ->
                docCols ColOpPrefix [appSep od, docSetBaseY ed]
              )
        -- > one
        -- >   + two
        -- >   + three
        addAlternative
          $ wrapParenMlIf
              hasParen
          -- ((if not hasParen then docAddBaseY BrIndentRegular else id) docL)
              (if hasParen then docSetBaseY docL else docL)
              (sharedOpsFlat <&> \(od, ed) ->
                docCols ColOpPrefix [appSep od, docSetBaseY ed]
              )