{-# 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 qualified GHC.OldList                   as List

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.ToBriDocTools
import           Language.Haskell.Brittany.Internal.Types
import           Language.Haskell.Brittany.Internal.Utils


allowAsOpPar :: LHsExpr GhcPs -> Bool
allowAsOpPar = \case
  L _ HsLam{}     -> True
  L _ HsLamCase{} -> True
  L _ HsCase{}    -> True
  L _ HsDo{}      -> True
  _               -> False

gatherOpTreeE
  :: OpParenMode
  -> Bool
  -> (ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered)
  -> Maybe GHC.RealSrcLoc
  -> Maybe GHC.RealSrcLoc
  -> [(ToBriDocM BriDocNumbered, OpTree)]
  -> LHsExpr GhcPs
  -> ToBriDocM (OpTree, Bool)
gatherOpTreeE hasParen hasComms commWrap locOpen locClose opExprList = \case
  (L _ (OpApp epAnn l1 op1 r1)) -> do
    inner <- callLayouter layout_expr r1
    gatherOpTreeE
      (case hasParen of
        NoParen -> NoParen
        _       -> ParenWithSpace
      )
      (hasComms || hasAnyCommentsBelow epAnn)
      commWrap
      locOpen
      locClose
      ( ( docHandleComms epAnn $ callLayouter layout_expr op1
        , OpLeaf (allowAsOpPar r1) inner
        )
      : opExprList
      )
      l1
  (L _ (HsPar epAnn inner)) | hasParen == NoParen && null opExprList -> 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)
    gatherOpTreeE ParenNoSpace
                  (hasComms || hasAnyCommentsBelow epAnn)
                  (commWrap . docHandleComms epAnn)
                  (mergePoses locOpen spanOpen)
                  (mergePoses locClose spanClose)
                  []
                  inner
  (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 ParenNoSpace
                    (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
      pure (x', y)
    pure
      $ ( OpUnknown hasParen locOpen locClose innerTree numberedRights
        , innerHasComms
        )
  final | hasParen == NoParen && null opExprList -> do
    tree <- commWrap $ callLayouter layout_expr final
    pure (OpLeaf (allowAsOpPar final) tree, hasComms)
  final@(L _ inner) -> do
    numberedLeft   <- commWrap $ callLayouter layout_expr final
    numberedRights <- opExprList `forM` \(x, y) -> do
      x' <- x
      pure (x', y)
    pure
      $ ( OpUnknown
            (case (hasParen, inner) of
              (NoParen, _              ) -> NoParen
              (_      , ExplicitTuple{}) -> ParenWithSpace
              _                          -> hasParen
            )
            locOpen
            locClose
            (OpLeaf (allowAsOpPar final) numberedLeft)
            numberedRights
        , hasComms
        )

gatherOpTreeT
  :: OpParenMode
  -> Bool
  -> (ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered)
  -> Maybe GHC.RealSrcLoc
  -> Maybe GHC.RealSrcLoc
  -> [(ToBriDocM BriDocNumbered, OpTree)]
  -> LHsType GhcPs
  -> ToBriDocM (OpTree, Bool)
gatherOpTreeT hasParen hasComms commWrap locOpen locClose opExprList = \case
  (L _ (HsOpTy NoExtField l1 op1 r1)) -> do
    inner <- callLayouter2 layout_type False r1
    gatherOpTreeT
      (case hasParen of
        NoParen -> NoParen
        _       -> ParenWithSpace
      )
      hasComms
      commWrap
      locOpen
      locClose
      ((docLit $ printRdrNameWithAnns op1, OpLeaf False inner) : opExprList)
      l1
  final@(L _ inner) -> do
    numberedLeft   <- commWrap $ callLayouter2 layout_type False final
    numberedRights <- opExprList `forM` \(x, y) -> do
      x' <- x
      pure (x', y)
    pure
      $ ( OpUnknown
            (case (hasParen, inner) of
              (NoParen, _          ) -> NoParen
              (_      , HsTupleTy{}) -> ParenWithSpace
              _                      -> hasParen
            )
            locOpen
            locClose
            (OpLeaf False 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 NoParen balancedTree
  -- tellDebugMess $ displayOpTree unknownTree
  -- 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, layoutOpTree True b)) docOps
    coreAlternative hasParen
                    locO
                    locC
                    Nothing
                    leftTree
                    sharedOps
                    sharedOps
                    docForceSingleline
  OpKnown NoParen Nothing Nothing fixity treeL docOps
    | Fixity _ (-1) _ <- fixity -> do
      dHead <- shareDoc $ layoutOpTree True treeL
      body  <- forM docOps $ \(op, arg) -> do
        arg' <- shareDoc $ layoutOpTree True arg
        pure (op, arg')
      runFilteredAlternative $ do
        addAlternativeCond allowSinglelinePar
          $ docForceSingleline
          $ docSeq
          $ dHead
          : join
              [ [docSeparator, pure prefix, docSeparator, doc]
              | (prefix, doc) <- body
              ]
        addAlternative $ docPar (docSetBaseY dHead) $ docLines
          [ docCols
              ColTyOpPrefix
              [ appSep $ case prefix of
                (_, BDLit s) | Text.length s == 1 -> docSeq
                  [docLitS " ", pure prefix]
                _ -> pure prefix
              , docEnsureIndent (BrIndentSpecial (length prefix + 1))
              $ docSetBaseY doc
              ]
          | (prefix, doc) <- body
          ]
  OpKnown hasParen locO locC fixity treeL docOps -> do
    let Fixity _ _prec _ = fixity
    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 NoParen _ _ _ 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 case List.last docOps of
            (_, OpLeaf True _) -> docForceParSpacing
            _                  -> docForceSingleline
    coreAlternative hasParen
                    locO
                    locC
                    (Just fixity)
                    treeL
                    sharedOps
                    sharedOpsFlat
                    lastWrap
  OpLeaf _ l -> pure l
 where
  isPrec0 x = getPrec x == 0
  getPrec = \case
    Fixity _ prec _ -> prec
  coreAlternative
    :: OpParenMode
    -> Maybe GHC.RealSrcLoc
    -> Maybe GHC.RealSrcLoc
    -> Maybe Fixity
    -> OpTree
    -> [(ToBriDocM BriDocNumbered, ToBriDocM BriDocNumbered)]
    -> [(ToBriDocM BriDocNumbered, ToBriDocM BriDocNumbered)]
    -> (ToBriDocM BriDocNumbered -> ToBriDocM BriDocNumbered)
    -> ToBriDocM BriDocNumbered
  coreAlternative NoParen _loc0 _locC _fixity treeL [] [] _lastWrap = do
    layoutOpTree True treeL
  coreAlternative ParenNoSpace locO locC _fixity treeL [] [] _lastWrap = do
    docL <- shareDoc $ layoutOpTree True treeL
    docAlt
      [ docSeq
        [ docLitS "("
        , docHandleComms locO $ docForceSingleline docL
        , docHandleComms locC $ docLitS ")"
        ]
      , docForceZeroAdd $ docSetBaseY $ docLines
        [ docSeq
          [ docLitS "("
          , docHandleComms locO $ docAddBaseY (BrIndentSpecial 2) $ docL
          ]
        , docHandleComms locC $ docLitS ")"
        ]
      , docPar
        (docSeq
          [ docLitS "("
          , docHandleComms locO $ docAddBaseY (BrIndentSpecial 2) $ docL
          ]
        )
        (docHandleComms locC $ docLitS ")")
      ]
  coreAlternative ParenWithSpace locO locC _fixity treeL [] [] _lastWrap = do
    docL <- shareDoc $ layoutOpTree True treeL
    docAlt
      [ docSeq
        [ docLitS "("
        , docHandleComms locO $ docForceSingleline docL
        , docHandleComms locC $ docLitS ")"
        ]
      , docForceZeroAdd $ docSetBaseY $ docLines
        [ docSeq
          [ docLitS "("
          , docSeparator
          , docHandleComms locO $ docAddBaseY (BrIndentSpecial 2) $ docL
          ]
        , docHandleComms locC $ docLitS ")"
        ]
      , docPar
        (docSeq
          [ docLitS "("
          , docSeparator
          , docHandleComms locO $ docAddBaseY (BrIndentSpecial 2) $ docL
          ]
        )
        (docHandleComms locC $ docLitS ")")
      ]
  coreAlternative hasParen locO locC fixity treeL sharedOps sharedOpsFlat lastWrap
    = do
      docL         <- shareDoc $ layoutOpTree True treeL
      indentPolicy <- askLayoutConf _lconfig_indentPolicy
      let zeroOps       = null sharedOps
          spaceAfterPar = not zeroOps
          wrapParenIfSl x inner = if x == NoParen
            then docSetParSpacing inner
            else wrapParenSl 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 = case x of
            NoParen        -> docPar innerHead (docLines innerLines)
            ParenWithSpace -> wrapParenMl True innerHead innerLines
            ParenNoSpace   -> wrapParenMl False innerHead innerLines
          wrapParenMl space innerHead innerLines = docAlt
            [ docForceZeroAdd $ docSetBaseY $ docLines
              (  [ docCols
                     ColOpPrefix
                     [ (if spaceAfterPar || space then appSep else id)
                     $ docLit
                     $ Text.pack "("
                     , docHandleComms locO $ innerHead
                     ]
                 ]
              ++ innerLines
              ++ [docHandleComms locC $ docLit $ Text.pack ")"]
              )
            , docPar
              (docCols
                ColOpPrefix
                [ (if spaceAfterPar || space then appSep else id)
                $ 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 (hasParen == NoParen && 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     -> docSetBaseY
        let curIsPrec0 = case fixity of
              Nothing                -> False
              Just (Fixity _ prec _) -> prec == 0
        case sharedOps of
          [(od, ed)] | curIsPrec0 ->
            addAlternativeCond (hasParen == NoParen && isSingleOp)
              $ docSetParSpacing
              $ docPar (docHandleComms locO $ docForceSingleline $ docL)
                       (docSeq [od, docSeparator, singlelineUnlessFree ed])
          _ -> pure ()
        -- > ( one
        -- > + two
        -- > + three
        -- > )
        addAlternativeCond (allowParIns && hasParen == NoParen)
          $ docForceZeroAdd
          $ wrapParenMl
              True
              (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 /= NoParen then docSetBaseY docL else docL)
            ( (if hasParen /= NoParen then sharedOps else sharedOpsFlat)
            <&> \(od, ed) -> docCols ColOpPrefix [appSep od, docSetBaseY ed]
            )