{-# 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
    let lastWrap = if getPrec fixity <= 1
          then docForceParSpacing
          else docForceSingleline
    coreAlternative hasParen
                    locO
                    locC
                    (Just fixity)
                    docL
                    sharedOps
                    sharedOpsFlat
                    lastWrap
  OpLeaf l -> pure l
 where
  isPrec0 x = getPrec x == 0
  getPrec = \case
    Fixity _ prec _ -> prec
  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
      indentPolicy <- askLayoutConf _lconfig_indentPolicy
      let zeroOps = null sharedOps
          wrapParenIfSl x inner = if x
            then wrapParenSl inner
            else docSetParSpacing 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 = docAlt
            [ docForceZeroAdd $ docSetBaseY $ docLines
              (  [ docCols
                     ColOpPrefix
                     [ (if zeroOps then id else appSep) $ docLit $ Text.pack "("
                     , docHandleComms locO $ innerHead
                     ]
                 ]
              ++ innerLines
              ++ [docHandleComms locC $ docLit $ Text.pack ")"]
              )
            , docPar
              (docCols
                ColOpPrefix
                [ (if zeroOps then id else appSep) $ docLit $ Text.pack "("
                , docHandleComms locO $ innerHead
                ]
              )
              ( docLines
              $ innerLines ++ [docHandleComms locC $ docLit $ Text.pack ")"]
              )
            ]

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

      let isSingleOp = length sharedOps == 1 && length sharedOpsFlat == 1

      runFilteredAlternative $ do
        -- > one + two + three
        -- or
        -- > one + two + case x of
        -- >   _ -> three
        addAlternativeCond allowSinglelinePar
          $ wrapParenIfSl hasParen
          $ docSetParSpacing
          $ docSeq
              ([docForceSingleline docL] ++ case splitFirstLast sharedOpsFlat 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
                     ]
                  )
              )
        -- one
        -- + two
        -- + three
        addAlternativeCond (not hasParen && not isSingleOp) $ docPar
          (docHandleComms locO $ docForceSingleline $ docL)
          (docFlushCommsPost False locC $ docLines $ sharedOps <&> \(od, ed) ->
            docCols ColOpPrefix [appSep od, docForceSingleline ed]
          )
        let singlelineUnlessFree = case indentPolicy of
              IndentPolicyLeft     -> docForceSingleline
              IndentPolicyMultiple -> docForceSingleline
              IndentPolicyFree     -> id
        let curIsPrec0 = case fixity of
              Nothing                -> False
              Just (Fixity _ prec _) -> prec == 0
        case sharedOps of
          [(od, ed)] | curIsPrec0 ->
            addAlternativeCond (not hasParen && isSingleOp)
              $ docSetParSpacing
              $ docPar (docHandleComms locO $ docForceSingleline $ docL)
                       (docSeq [od, docSeparator, singlelineUnlessFree ed])
          _ -> pure ()
        -- > ( one
        -- > + two
        -- > + three
        -- > )
        addAlternativeCond (allowParIns && not hasParen)
          $ docForceZeroAdd
          $ wrapParenMl
              (docSetBaseY docL)
              (sharedOps <&> \(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)
            ((if hasParen then sharedOps else sharedOpsFlat) <&> \(od, ed) ->
              docCols ColOpPrefix [appSep od, docSetBaseY ed]
            )