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

module Language.Haskell.Brittany.Internal.Transformations.T2_Floating where

import qualified Data.Generics.Uniplate.Direct as Uniplate
import qualified GHC.OldList as List
import           GHC                            ( GenLocated(L) )

import           Language.Haskell.Brittany.Internal.Components.BriDoc
import           Language.Haskell.Brittany.Internal.Prelude
import           Language.Haskell.Brittany.Internal.Utils



-- note that this is not total, and cannot be with that exact signature.
mergeIndents :: BrIndent -> BrIndent -> BrIndent
mergeIndents BrIndentNone x = x
mergeIndents x BrIndentNone = x
mergeIndents (BrIndentSpecial i) (BrIndentSpecial j) =
  BrIndentSpecial (max i j)
mergeIndents _ _ = error "mergeIndents"


transformSimplifyFloating :: BriDoc -> BriDoc
transformSimplifyFloating = stepBO .> stepFull
  -- note that semantically, stepFull is completely sufficient.
  -- but the bottom-up switch-to-top-down-on-match transformation has much
  -- better complexity.
  -- UPDATE: by now, stepBO does more than stepFull; for semantic equivalence
  --         the push/pop cases would need to be copied over
 where
  descendCommsPrior = transformDownMay $ \case
    -- prior floating in
    BDFlushCommentsPrior loc1 (BDFlushCommentsPrior loc2 x) ->
      Just $ BDFlushCommentsPrior (max loc1 loc2) x
    BDFlushCommentsPrior loc1 (BDFlushCommentsPost loc2 shouldMark2 x) ->
      Just $ BDFlushCommentsPost loc2 shouldMark2 (BDFlushCommentsPrior loc1 x)
    BDFlushCommentsPrior loc1 (BDPar ind line indented) ->
      Just $ BDPar ind (BDFlushCommentsPrior loc1 line) indented
    BDFlushCommentsPrior loc1 (BDSeq (l : lr)) ->
      Just $ BDSeq (BDFlushCommentsPrior loc1 l : lr)
    BDFlushCommentsPrior loc1 (BDLines (l : lr)) ->
      Just $ BDLines (BDFlushCommentsPrior loc1 l : lr)
    BDFlushCommentsPrior loc1 (BDCols sig (l : lr)) ->
      Just $ BDCols sig (BDFlushCommentsPrior loc1 l : lr)
    BDFlushCommentsPrior loc1 (BDAddBaseY indent x) ->
      Just $ BDAddBaseY indent $ BDFlushCommentsPrior loc1 x
    BDFlushCommentsPrior loc1 (BDDebug s x) ->
      Just $ BDDebug s $ BDFlushCommentsPrior loc1 x
    BDFlushCommentsPrior loc1 (BDQueueComments comms1 x) ->
      Just $ BDQueueComments comms1 $ BDFlushCommentsPrior loc1 x
    _ -> Nothing
  descendCommsPost = transformDownMay $ \case
    -- post floating in
    BDFlushCommentsPost loc1 mark1 (BDFlushCommentsPost loc2 mark2 x) ->
      Just $ BDFlushCommentsPost (max loc1 loc2) (mark1 || mark2) x
    BDFlushCommentsPost loc1 shouldMark1 (BDPar ind line indented) ->
      Just $ BDPar ind line $ BDFlushCommentsPost loc1 shouldMark1 indented
    BDFlushCommentsPost loc1 shouldMark1 (BDSeq list) ->
      Just
        $ BDSeq
        $ List.init list
        ++ [BDFlushCommentsPost loc1 shouldMark1 $ List.last list]
    BDFlushCommentsPost loc1 shouldMark1 (BDLines list) ->
      Just
        $ BDLines
        $ List.init list
        ++ [BDFlushCommentsPost loc1 shouldMark1 $ List.last list]
    BDFlushCommentsPost loc1 shouldMark1 (BDCols sig cols) ->
      Just
        $ BDCols sig
        $ List.init cols
        ++ [BDFlushCommentsPost loc1 shouldMark1 $ List.last cols]
    BDFlushCommentsPost loc1 shouldMark1 (BDIndentLevelPop x) ->
      Just $ BDIndentLevelPop (BDFlushCommentsPost loc1 shouldMark1 x)
    BDFlushCommentsPost loc1 shouldMark1 (BDIndentLevelPushCur x) ->
      Just $ BDIndentLevelPushCur (BDFlushCommentsPost loc1 shouldMark1 x)
    BDFlushCommentsPost loc1 shouldMark1 (BDBaseYPushCur x) ->
      Just $ BDBaseYPushCur (BDFlushCommentsPost loc1 shouldMark1 x)
    BDFlushCommentsPost loc1 shouldMark1 (BDAddBaseY indent x) ->
      Just $ BDAddBaseY indent $ BDFlushCommentsPost loc1 shouldMark1 x
    BDFlushCommentsPost loc1 shouldMark1 (BDDebug s x) ->
      Just $ BDDebug s $ BDFlushCommentsPost loc1 shouldMark1 x
    _ -> Nothing
  descendQueueComments = transformDownMay $ \case
    -- queue comments floating in
    BDQueueComments comms1 (BDQueueComments comms2 x) ->
      Just $ BDQueueComments (mergeOn (\(L l _) -> l) comms1 comms2) x
    BDQueueComments comms1 (BDPar ind line indented) ->
      Just $ BDPar ind (BDQueueComments comms1 line) indented
    BDQueueComments comms1 (BDSeq (l : lr)) ->
      Just $ BDSeq (BDQueueComments comms1 l : lr)
    BDQueueComments comms1 (BDLines (l : lr)) ->
      Just $ BDLines (BDQueueComments comms1 l : lr)
    BDQueueComments comms1 (BDCols sig (l : lr)) ->
      Just $ BDCols sig (BDQueueComments comms1 l : lr)
    BDQueueComments comms1 (BDAddBaseY indent x) ->
      Just $ BDAddBaseY indent $ BDQueueComments comms1 x
    BDQueueComments comms1 (BDDebug s x) ->
      Just $ BDDebug s $ BDQueueComments comms1 x
    _ -> Nothing
  descendBYPush = transformDownMay $ \case
    BDBaseYPushCur (BDCols sig cols@(_ : _)) ->
      Just $ BDCols sig (BDBaseYPushCur (List.head cols) : List.tail cols)
    BDBaseYPushCur (BDDebug s x) -> Just $ BDDebug s (BDBaseYPushCur x)
    _ -> Nothing
  descendEntryDelta = transformDownMay $ \case
    -- entry delta floating in
    BDEntryDelta dp (BDPar ind line indented) ->
      Just $ BDPar ind (BDEntryDelta dp line) indented
    BDEntryDelta dp (BDSeq (l : lr)) ->
      Just $ BDSeq (BDEntryDelta dp l : lr)
    BDEntryDelta dp (BDLines (l : lr)) ->
      Just $ BDLines (BDEntryDelta dp l : lr)
    BDEntryDelta dp (BDCols sig (l : lr)) ->
      Just $ BDCols sig (BDEntryDelta dp l : lr)
    BDEntryDelta dp (BDAddBaseY indent x) ->
      Just $ BDAddBaseY indent $ BDEntryDelta dp x
    BDEntryDelta dp (BDDebug s x) ->
      Just $ BDDebug s $ BDEntryDelta dp x
    _ -> Nothing
  descendILPush = transformDownMay $ \case
    BDIndentLevelPushCur (BDCols sig cols@(_ : _)) ->
      Just $ BDCols sig (BDIndentLevelPushCur (List.head cols) : List.tail cols)
    BDIndentLevelPushCur (BDDebug s x) ->
      Just $ BDDebug s (BDIndentLevelPushCur x)
    _ -> Nothing
  descendILPop = transformDownMay $ \case
    BDIndentLevelPop (BDCols sig cols@(_ : _)) ->
      Just $ BDCols sig (List.init cols ++ [BDIndentLevelPop (List.last cols)])
    BDIndentLevelPop (BDDebug s x) -> Just $ BDDebug s (BDIndentLevelPop x)
    _ -> Nothing
  descendAddB = transformDownMay $ \case
    BDAddBaseY BrIndentNone x -> Just x
    -- AddIndent floats into Lines.
    BDAddBaseY _ind (BDLines lines) ->
      Just $ BDLines $ lines
    -- AddIndent floats into last column
    BDAddBaseY ind (BDCols sig cols) ->
      Just $ BDCols sig $ List.init cols ++ [BDAddBaseY ind $ List.last cols]
    -- merge AddIndent and Par
    BDAddBaseY ind1 (BDPar ind2 line indented) ->
      Just $ BDPar (mergeIndents ind1 ind2) line indented
    BDAddBaseY ind (BDFlushCommentsPrior loc x) ->
      Just $ BDFlushCommentsPrior loc (BDAddBaseY ind x)
    BDAddBaseY ind (BDFlushCommentsPost loc shouldMark x) ->
      Just $ BDFlushCommentsPost loc shouldMark (BDAddBaseY ind x)
    BDAddBaseY ind (BDQueueComments comms x) ->
      Just $ BDQueueComments comms (BDAddBaseY ind x)
    BDAddBaseY ind (BDSeq list) ->
      Just $ BDSeq $ List.init list ++ [BDAddBaseY ind (List.last list)]
    BDAddBaseY _ lit@BDLit{} -> Just $ lit
    BDAddBaseY _ind (BDBaseYPushCur x) ->
      Just $ BDBaseYPushCur x
    -- TODO92 We have several rules here in conflict with each other.
    --        Unless I forget some detail related to some elements being able
    --        to float in further, we probably should define some
    --        intended "floating" order that all rules should adhere.
    --        According to that order, we should either have
    --          A (B x) -> B (A x)
    --        or
    --          B (A x) -> A (B x)
    --        but not both.
    --        I have disable the following conflict for
    --        BDAddBaseY vs BDEntryDelta for now.
    -- BDAddBaseY ind (BDEntryDelta dp) -> Just $ BDEntryDelta dp (BDAddBaseY ind x)
    BDAddBaseY ind (BDDebug s x) -> Just $ BDDebug s (BDAddBaseY ind x)
    BDAddBaseY ind (BDIndentLevelPop x) ->
      Just $ BDIndentLevelPop (BDAddBaseY ind x)
    BDAddBaseY ind (BDIndentLevelPushCur x) ->
      Just $ BDIndentLevelPushCur (BDAddBaseY ind x)
    BDAddBaseY ind (BDEnsureIndent ind2 x) ->
      Just $ BDEnsureIndent (mergeIndents ind ind2) x
    _ -> Nothing
  stepBO :: BriDoc -> BriDoc
  stepBO = -- traceFunctionWith "stepBO" (show . briDocToDocWithAnns) (show . briDocToDocWithAnns) $
           transformUp f
   where
    f = \case
      BDSeq xs -> BDSeq (dropWhile (\case BDEmpty -> True; _ -> False) xs)
      x@BDFlushCommentsPrior{} -> descendCommsPrior x
      x@BDFlushCommentsPost{} -> descendCommsPost x
      x@BDQueueComments{} -> descendQueueComments x
      x@BDAddBaseY{} -> descendAddB x
      x@BDBaseYPushCur{} -> descendBYPush x
      x@BDEntryDelta{} -> descendEntryDelta x
      x@BDIndentLevelPushCur{} -> descendILPush x
      x@BDIndentLevelPop{} -> descendILPop x
      x -> x
  stepFull = -- traceFunctionWith "stepFull" (show . briDocToDocWithAnns) (show . briDocToDocWithAnns) $
             Uniplate.rewrite $ \case
    BDSeq (BDEmpty : x) -> Just (BDSeq x)
    -- TODO92 This could be rewritten, so that we re-use the
    -- (BriDoc -> Maybe Bridoc) part of the descend* functions instead of
    -- copying them here (incompletely).
    BDAddBaseY BrIndentNone x -> Just $ x
    -- AddIndent floats into Lines.
    BDAddBaseY _ind (BDLines lines) ->
      Just $ BDLines lines
    -- AddIndent floats into last column
    BDAddBaseY ind (BDCols sig cols) ->
      Just $ BDCols sig $ List.init cols ++ [BDAddBaseY ind $ List.last cols]
    BDAddBaseY ind (BDSeq list) ->
      Just $ BDSeq $ List.init list ++ [BDAddBaseY ind (List.last list)]
    -- merge AddIndent and Par
    BDAddBaseY ind1 (BDPar ind2 line indented) ->
      Just $ BDPar (mergeIndents ind1 ind2) line indented
    BDAddBaseY _ lit@BDLit{} -> Just $ lit
    BDAddBaseY _ind (BDBaseYPushCur x) ->
      Just $ BDBaseYPushCur x
    -- EnsureIndent float-in
    -- BDEnsureIndent indent (BDCols sig (col:colr)) ->
    --   Just $ BDCols sig (BDEnsureIndent indent col : (BDAddBaseY indent <$> colr))
    -- not sure if the following rule is necessary; tests currently are
    -- unaffected.
    -- BDEnsureIndent indent (BDLines lines) ->
    --   Just $ BDLines $ BDEnsureIndent indent <$> lines
    -- flush-prior floating in
    BDFlushCommentsPrior loc1 (BDFlushCommentsPrior loc2 x) ->
      Just $ BDFlushCommentsPrior (max loc1 loc2) x
    BDFlushCommentsPrior loc1 (BDFlushCommentsPost loc2 shouldMark2 x) ->
      Just $ BDFlushCommentsPost loc2 shouldMark2 (BDFlushCommentsPrior loc1 x)
    BDFlushCommentsPrior loc (BDSeq (l : lr)) ->
      Just $ BDSeq (BDFlushCommentsPrior loc l : lr)
    BDFlushCommentsPrior loc (BDLines (l : lr)) ->
      Just $ BDLines (BDFlushCommentsPrior loc l : lr)
    BDFlushCommentsPrior loc (BDCols sig (l : lr)) ->
      Just $ BDCols sig (BDFlushCommentsPrior loc l : lr)
    -- flush-post floating in
    BDFlushCommentsPost loc1 mark1 (BDPar ind line indented) ->
      Just $ BDPar ind line $ BDFlushCommentsPost loc1 mark1 indented
    BDFlushCommentsPost loc1 mark1 (BDSeq list) ->
      Just $ BDSeq $ List.init list ++ [BDFlushCommentsPost loc1 mark1 $ List.last list]
    BDFlushCommentsPost loc1 mark1 (BDLines list) ->
      Just $ BDLines $ List.init list ++ [BDFlushCommentsPost loc1 mark1 $ List.last list]
    BDFlushCommentsPost loc1 mark1 (BDCols sig list) ->
      Just $ BDCols sig $ List.init list ++ [BDFlushCommentsPost loc1 mark1 $ List.last list]
    BDFlushCommentsPost loc1 mark1 (BDIndentLevelPop x) ->
      Just $ BDIndentLevelPop (BDFlushCommentsPost loc1 mark1 x)
    BDFlushCommentsPost loc1 mark1 (BDIndentLevelPushCur x) ->
      Just $ BDIndentLevelPushCur (BDFlushCommentsPost loc1 mark1 x)
    BDFlushCommentsPost loc1 mark1 (BDBaseYPushCur x) ->
      Just $ BDBaseYPushCur (BDFlushCommentsPost loc1 mark1 x)
    BDFlushCommentsPost loc1 mark1 (BDAddBaseY indent x) ->
      Just $ BDAddBaseY indent $ BDFlushCommentsPost loc1 mark1 x
    BDFlushCommentsPost loc1 mark1 (BDDebug s x) ->
      Just $ BDDebug s $ BDFlushCommentsPost loc1 mark1 x
    -- entry-delta floating in
    BDEntryDelta dp (BDPar ind line indented) ->
      Just $ BDPar ind (BDEntryDelta dp line) indented
    BDEntryDelta dp (BDSeq (l : lr)) ->
      Just $ BDSeq (BDEntryDelta dp l : lr)
    BDEntryDelta dp (BDLines (l : lr)) ->
      Just $ BDLines (BDEntryDelta dp l : lr)
    BDEntryDelta dp (BDCols sig (l : lr)) ->
      Just $ BDCols sig (BDEntryDelta dp l : lr)
    BDEntryDelta dp (BDAddBaseY indent x) ->
      Just $ BDAddBaseY indent $ BDEntryDelta dp x
    BDEntryDelta dp (BDDebug s x) ->
      Just $ BDDebug s $ BDEntryDelta dp x
    
    _ -> Nothing
