{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE ViewPatterns #-}
module Language.Haskell.Brittany.Internal.Components.OpTree where

import           Language.Haskell.Brittany.Internal.Prelude

import qualified Data.Text                     as Text
import           GHC.Types.Fixity               ( Fixity(Fixity)
                                                , FixityDirection
                                                  ( InfixL
                                                  , InfixN
                                                  , InfixR
                                                  )
                                                )
import           GHC.Types.SourceText           ( SourceText(NoSourceText) )
import qualified Safe
import qualified Data.Char

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



displayOpTree :: OpTree -> String
displayOpTree = \case
  OpUnknown p _ _ leftTree rs ->
    (  "(OpUnknown "
    ++ show p
    ++ " "
    ++ displayOpTree leftTree
    ++ " ["
    ++ intercalate
         ","
         [ ("(" ++ showOp op ++ "," ++ displayOpTree x ++ ")") | (op, x) <- rs ]
    ++ "])"
    )
  OpKnown p _ _ fixity tree ops ->
    (  "OpKnown "
    ++ show p
    ++ " "
    ++ showOutputable fixity
    ++ " ("
    ++ displayOpTree tree
    ++ ")"
    ++ join [ (showOp op ++ "(" ++ displayOpTree t ++ ")") | (op, t) <- ops ]
    )
  OpLeaf (x, _) -> show x
 where
  showOp :: BriDocNumbered -> String
  showOp = \case
    (_, BDFlushCommentsPrior _ x) -> showOp x
    (_, BDQueueComments _ x     ) -> showOp x
    (_, BDLit x                 ) -> Text.unpack x
    -- (_, BDFlushCommentsPrior _ (_, BDFlushCommentsPrior _ (_, x)))
    --   | trace (show $ toConstr x) False -> "meow"
    (i, _                       ) -> show i


-- lookupFixities :: Monad m => OpTree -> m OpTree
-- lookupFixities = \case
--   OpNode par Nothing opDoc chldrn -> do
--     pure $ OpNode par (hardcodedFixity (Text.unpack opDoc)) opDoc chldrn
--   x@OpNode{} -> pure x
--   x@OpLeaf{} -> pure x

data ReformatParenMode
  = ReformatParenModeKeep  -- don't modify parens at all
  | ReformatParenModeClean -- remove unnecessary parens
  | ReformatParenModeAll   -- add superfluous parens everywhere

-- [(Bool, Fixity, Text, [OpTree])]

-- a == b + c || d * e /= f
-- _ a
-- == a, _ b
-- == a, + b, _ c
-- == a, + b c
-- == a (+ b c)
-- || (== a (+ b c)), _ d
-- || (== a (+ b c)), * d, _ e
-- || (== a (+ b c)), * d e
-- || (== a (+ b c)), /= (* d e), _ f
-- || (== a (+ b c)), /= (* d e) f
-- || (== a (+ b c)) (/= (* d e) f)

data StackElem = StackElem Fixity [(OpTree, BriDocNumbered)]
type Stack = [StackElem]

balanceOpTree :: Bool -> OpTree -> ([BrittanyError], OpTree)
balanceOpTree allowUnqualify = \case
  x@OpLeaf{} -> ([], x)
  OpKnown paren locO locC fixity@(Fixity _ (-1) _) left rest ->
    let
        (warnsLeft, balancedLeft) = balanceOpTree allowUnqualify left
        opRes =
          [ (op, balanceOpTree allowUnqualify argTree) | (op, argTree) <- rest ]
    in  ( warnsLeft ++ [ w | (_, (warns, _)) <- opRes, w <- warns ]
        , OpKnown paren
                  locO
                  locC
                  fixity
                  balancedLeft
                  [ (op, balanced) | (op, (_, balanced)) <- opRes ]
        )
  x@OpKnown{} -> ([], x)
  x@(OpUnknown paren locO locC left rest) ->
    let (warns, balancedLeft) = balanceOpTree allowUnqualify left
    in  case go [] rest balancedLeft of
          Right (leaf@OpLeaf{}) -> (warns, leaf)
          Right (OpKnown _paren _ _ fixity c cs) ->
            (warns, OpKnown paren locO locC fixity c cs)
          Right t -> (warns, t)
          Left moreWarns ->
            ( (  warns
              ++ [ LayoutWarning ("Fixity of operator not known: " ++ w)
                 | w <- moreWarns
                 ]
              )
            , x
            )
 where
  -- singleton :: BriDocNumbered -> StackElem
  -- singleton x = StackElem (Fixity NoSourceText 99 InfixN) (OpLeaf x) []
  go :: Stack -> [(BriDocNumbered, OpTree)] -> OpTree -> Either [String] OpTree
  go [] [] _ = Left []
  go [StackElem fxty cs] [] c =
    let (e1, eops) = shiftOps cs c in Right (known fxty e1 eops)
  go (StackElem fxty cs : StackElem fixity cs2 : rest) [] c =
    -- go (StackElem fixity (OpKnown False fxty (reverse cs) : cs2) : rest) []
    let (e1, eops) = shiftOps cs c
    in  go (StackElem fixity cs2 : rest) [] (known fxty e1 eops)
  go stack input@((opDoc, val) : inputR) c = case stack of
    [] -> do
      fxty <- docFixity opDoc
      go [StackElem fxty [(c, opDoc)]] inputR val
    (StackElem fixityS cs : stackR) -> do
      let Fixity _ precS dirS = fixityS
      fxty@(Fixity _ prec dir) <- docFixity opDoc
      case compare prec precS of
        GT -> go (StackElem fxty [(c, opDoc)] : stack) inputR val
        LT -> do
          let (e1, eops) = shiftOps cs c
          go stackR input (known fixityS e1 eops)
        EQ -> case (dir, dirS) of
          (InfixR, InfixR) ->
            go (StackElem fixityS ((c, opDoc) : cs) : stackR) inputR val
          (InfixL, InfixL) ->
            go (StackElem fixityS ((c, opDoc) : cs) : stackR) inputR val
          _ -> Left []
  docFixity :: BriDocNumbered -> Either [String] Fixity
  docFixity (_, x) = case x of
    BDLit (Text.unpack -> s) -> case hardcodedFixity allowUnqualify s of
      Just f  -> Right f
      Nothing -> Left [s]
    BDFlushCommentsPrior _ d -> docFixity d
    BDQueueComments      _ d -> docFixity d
    _                        -> Left []
  shiftOps
    :: [(OpTree, BriDocNumbered)]
    -> OpTree
    -> (OpTree, [(BriDocNumbered, OpTree)])
  shiftOps ops final = case reverse ops of
    [] -> (final, [])
    ((e1, o1) : rest) ->
      ( e1
      , let (finalOp, list) =
              mapAccumL (\op (nextE, nextOp) -> (nextOp, (op, nextE))) o1 rest
        in  list ++ [(finalOp, final)]
      )
  known = OpKnown NoParen Nothing Nothing

addAllParens :: OpParenMode -> OpTree -> OpTree
addAllParens topLevelParen = \case
  x@OpLeaf{}    -> x
  x@OpUnknown{} -> x
  OpKnown _paren locO locC fixity c cs ->
    OpKnown topLevelParen
            locO
            locC
            fixity
            (addAllParens ParenWithSpace c)
            [ (op, addAllParens ParenWithSpace tree) | (op, tree) <- cs ]

remSuperfluousParens :: Int -> OpTree -> OpTree
remSuperfluousParens outerFixity = \case
  x@OpLeaf{} -> x
  OpUnknown _ locO locC c@(OpLeaf doc) [] | isLit doc ->
    OpUnknown NoParen locO locC c []
  OpUnknown _ locO locC c@(OpUnknown ParenWithSpace _ _ _ _) [] ->
    OpUnknown NoParen locO locC (remSuperfluousParens 11 c) []
  OpUnknown _ locO locC c@(OpUnknown ParenNoSpace _ _ _ _) [] ->
    OpUnknown NoParen locO locC (remSuperfluousParens 11 c) []
  x@OpUnknown{} -> x
  OpKnown paren locO locC fixity c cs ->
    OpKnown
      -- We do not support removing superfluous parens around
      -- function types yet:
      (if outerFixity > fixLevel fixity || fixLevel fixity < 0
        then paren
        else NoParen
      )
      locO
      locC
      fixity
      (remSuperfluousParens (fixLevel fixity) c)
      [ (op, remSuperfluousParens (fixLevel fixity) tree) | (op, tree) <- cs ]
 where
  fixLevel (Fixity _ i _) = i
  isLit = \case
    (_, BDFlushCommentsPrior _ x ) -> isLit x
    (_, BDFlushCommentsPost _ _ x) -> isLit x
    (_, BDQueueComments _ x      ) -> isLit x
    (_, BDEntryDelta _ x         ) -> isLit x
    (_, BDForceAlt _ x           ) -> isLit x
    (_, BDDebug _ x              ) -> isLit x
    (_, BDAddBaseY _ x           ) -> isLit x
    (_, BDBaseYPushCur x         ) -> isLit x
    (_, BDIndentLevelPushCur x   ) -> isLit x
    (_, BDIndentLevelPop x       ) -> isLit x
    (_, BDLit{}                  ) -> True
    _                              -> False


hardcodedFixity :: Bool -> String -> Maybe Fixity
hardcodedFixity allowUnqualify = \case
  --
  "->"          -> Just $ Fixity NoSourceText (-1) InfixR
  "."           -> Just $ Fixity NoSourceText 9 InfixR
  "!!"          -> Just $ Fixity NoSourceText 9 InfixL
  "**"          -> Just $ Fixity NoSourceText 8 InfixR
  "^"           -> Just $ Fixity NoSourceText 8 InfixR
  "^^"          -> Just $ Fixity NoSourceText 8 InfixR
  "*"           -> Just $ Fixity NoSourceText 7 InfixL
  "/"           -> Just $ Fixity NoSourceText 7 InfixL
  "`quot`"      -> Just $ Fixity NoSourceText 7 InfixL
  "`rem`"       -> Just $ Fixity NoSourceText 7 InfixL
  "`div`"       -> Just $ Fixity NoSourceText 7 InfixL
  "`mod`"       -> Just $ Fixity NoSourceText 7 InfixL
  "+"           -> Just $ Fixity NoSourceText 6 InfixL
  "-"           -> Just $ Fixity NoSourceText 6 InfixL
  ":"           -> Just $ Fixity NoSourceText 5 InfixR
  "=="          -> Just $ Fixity NoSourceText 4 InfixN
  "/="          -> Just $ Fixity NoSourceText 4 InfixN
  "<"           -> Just $ Fixity NoSourceText 4 InfixN
  "<="          -> Just $ Fixity NoSourceText 4 InfixN
  ">"           -> Just $ Fixity NoSourceText 4 InfixN
  ">="          -> Just $ Fixity NoSourceText 4 InfixN
  "&&"          -> Just $ Fixity NoSourceText 3 InfixR
  "||"          -> Just $ Fixity NoSourceText 2 InfixR
  ">>="         -> Just $ Fixity NoSourceText 1 InfixL
  ">>"          -> Just $ Fixity NoSourceText 1 InfixL
  "=<<"         -> Just $ Fixity NoSourceText 1 InfixR
  "$"           -> Just $ Fixity NoSourceText 0 InfixR
  "`seq`"       -> Just $ Fixity NoSourceText 0 InfixR
  "$!"          -> Just $ Fixity NoSourceText 0 InfixR
  "!"           -> Just $ Fixity NoSourceText 9 InfixL
  "//"          -> Just $ Fixity NoSourceText 9 InfixL
  "<>"          -> Just $ Fixity NoSourceText 6 InfixR
  "<+>"         -> Just $ Fixity NoSourceText 5 InfixR
  "<$"          -> Just $ Fixity NoSourceText 4 InfixL
  "$>"          -> Just $ Fixity NoSourceText 4 InfixL
  "<$>"         -> Just $ Fixity NoSourceText 4 InfixL
  "<&>"         -> Just $ Fixity NoSourceText 1 InfixL
  "&"           -> Just $ Fixity NoSourceText 1 InfixL
  "<*>"         -> Just $ Fixity NoSourceText 4 InfixL
  "<**>"        -> Just $ Fixity NoSourceText 4 InfixL
  "*>"          -> Just $ Fixity NoSourceText 4 InfixL
  "<*"          -> Just $ Fixity NoSourceText 4 InfixL
  "`elem`"      -> Just $ Fixity NoSourceText 4 InfixN
  "`notElem`"   -> Just $ Fixity NoSourceText 4 InfixN
  "++"          -> Just $ Fixity NoSourceText 5 InfixR
  "%"           -> Just $ Fixity NoSourceText 7 InfixL
  "<|>"         -> Just $ Fixity NoSourceText 3 InfixL
  ".&."         -> Just $ Fixity NoSourceText 7 InfixL
  ".|."         -> Just $ Fixity NoSourceText 5 InfixL
  "`xor`"       -> Just $ Fixity NoSourceText 6 InfixL
  "`shift`"     -> Just $ Fixity NoSourceText 8 InfixL
  "`rotate`"    -> Just $ Fixity NoSourceText 8 InfixL
  "`shiftL`"    -> Just $ Fixity NoSourceText 8 InfixL
  "`shiftR`"    -> Just $ Fixity NoSourceText 8 InfixL
  "`rotateL`"   -> Just $ Fixity NoSourceText 8 InfixL
  "`rotateR`"   -> Just $ Fixity NoSourceText 8 InfixL
  "+#"          -> Just $ Fixity NoSourceText 6 InfixL
  ".^."         -> Just $ Fixity NoSourceText 6 InfixL
  ".>>."        -> Just $ Fixity NoSourceText 8 InfixL
  ".<<."        -> Just $ Fixity NoSourceText 8 InfixL
  "!>>."        -> Just $ Fixity NoSourceText 8 InfixL
  "!<<."        -> Just $ Fixity NoSourceText 8 InfixL
  ">=>"         -> Just $ Fixity NoSourceText 1 InfixR
  "<=<"         -> Just $ Fixity NoSourceText 1 InfixR
  "&&&"         -> Just $ Fixity NoSourceText 3 InfixR
  "<$!>"        -> fixity InfixL 4

  "~"           -> fixity InfixN 4
  "~~"          -> fixity InfixN 4
  ":~:"         -> Just $ Fixity NoSourceText 4 InfixN
  ":~~:"        -> Just $ Fixity NoSourceText 4 InfixN
  ":+:"         -> fixity InfixR 5
  ":*:"         -> fixity InfixR 6
  ":.:"         -> fixity InfixR 7
  ":|"          -> fixity InfixR 5

  -- non-base from random sources.
  "<|"          -> Just $ Fixity NoSourceText 5 InfixR
  "|>"          -> Just $ Fixity NoSourceText 5 InfixL
  "><"          -> Just $ Fixity NoSourceText 5 InfixR
  "$+$"         -> Just $ Fixity NoSourceText 5 InfixL
  "\\\\"        -> Just $ Fixity NoSourceText 5 InfixN
  ".>"          -> Just $ Fixity NoSourceText 9 InfixL
  ":?"          -> Just $ Fixity NoSourceText 7 InfixN
  ":-"          -> Just $ Fixity NoSourceText 9 InfixR
  ".:"          -> Just $ Fixity NoSourceText 9 InfixR
  "$!!"         -> fixity InfixR 0
  "<$!!>"       -> fixity InfixL 4

  -- lens, not complete!
  "<|"          -> fixity InfixR 5
  "|>"          -> fixity InfixL 5
  "%~"          -> fixity InfixR 4
  ".~"          -> fixity InfixR 4
  "?~"          -> fixity InfixR 4
  "<.~"         -> fixity InfixR 4
  "<?~"         -> fixity InfixR 4
  "+~"          -> fixity InfixR 4
  "*~"          -> fixity InfixR 4
  "-~"          -> fixity InfixR 4
  "//~"         -> fixity InfixR 4
  "^~"          -> fixity InfixR 4
  "^^~"         -> fixity InfixR 4
  "**~"         -> fixity InfixR 4
  "||~"         -> fixity InfixR 4
  "&&~"         -> fixity InfixR 4
  ".="          -> fixity InfixN 4
  "%="          -> fixity InfixN 4
  "?="          -> fixity InfixN 4
  "+="          -> fixity InfixN 4
  "-="          -> fixity InfixN 4
  "*="          -> fixity InfixN 4
  "//="         -> fixity InfixN 4
  "^="          -> fixity InfixN 4
  "^^="         -> fixity InfixN 4
  "**="         -> fixity InfixN 4
  "&&="         -> fixity InfixN 4
  "||="         -> fixity InfixN 4
  "<~"          -> fixity InfixR 2
  "<.="         -> fixity InfixN 4
  "<?="         -> fixity InfixN 4
  "<>~"         -> fixity InfixR 4
  "<>="         -> fixity InfixN 4
  "^.."         -> fixity InfixL 8
  "^?"          -> fixity InfixL 8
  "^?!"         -> fixity InfixL 8
  "^@.."        -> fixity InfixL 8
  "^@?"         -> fixity InfixL 8
  "^@?!"        -> fixity InfixL 8
  "^."          -> fixity InfixL 8
  "^@."         -> fixity InfixL 8
  "<."          -> fixity InfixR 9
  ".>"          -> fixity InfixR 9
  "<.>"         -> fixity InfixR 9
  "@@~"         -> fixity InfixR 4
  "@@="         -> fixity InfixR 4
  "&~"          -> fixity InfixL 1
  "??"          -> fixity InfixL 1


  -- certain other operators
  -- "%."        -> Just $ Fixity NoSourceText _ _
  -- "%=="       -> Just $ Fixity NoSourceText _ _
  -- "%=>%"      -> Just $ Fixity NoSourceText _ _
  -- "&*"        -> Just $ Fixity NoSourceText _ _
  -- "&/"        -> Just $ Fixity NoSourceText _ _
  -- "&="        -> Just $ Fixity NoSourceText _ _
  -- "-->"       -> Just $ Fixity NoSourceText _ _
  -- ".*?"       -> fixity _ _
  -- ".+."       -> fixity _ _
  -- ".-."       -> fixity _ _
  ".&."         -> fixity InfixR 1
  ".&&."        -> fixity InfixR 1
  ".||."        -> fixity InfixR 1
  "==>"         -> fixity InfixR 0
  "=/="         -> fixity InfixN 4
  "==="         -> fixity InfixN 4
  ".:!"         -> fixity InfixL 9
  ".:?"         -> fixity InfixL 9
  -- ".:>"       -> fixity _ _
  -- ".:>?"      -> fixity _ _
  "<.>"         -> fixity InfixR 7
  "</>"         -> fixity InfixR 5
  "<?>"         -> fixity InfixL 9
  -- "~"         -> fixity _ _
  "==="         -> fixity InfixN 4
  "!?"          -> fixity InfixL 9
  "%=="         -> fixity InfixN 3
  ".*"          -> fixity InfixR 8
  ".**"         -> fixity InfixR 8
  ".***"        -> fixity InfixR 8
  ":?-"         -> fixity InfixN 1
  "::-"         -> fixity InfixN 0
  "&!"          -> fixity InfixL 1

  -- quickcheck (-state-machine)
  ":&&"         -> fixity InfixL 9
  ":||"         -> fixity InfixL 9
  ":&&:"        -> fixity InfixL 9
  ":=>"         -> fixity InfixL 9
  ":=="         -> fixity InfixL 9
  ":/="         -> fixity InfixL 9
  ":<"          -> fixity InfixL 9
  ":<="         -> fixity InfixL 9
  ":>"          -> fixity InfixL 9
  ":>="         -> fixity InfixL 9
  ":->"         -> fixity InfixL 9
  ".=="         -> fixity InfixN 5
  "./"          -> fixity InfixN 5
  ".<"          -> fixity InfixN 5
  ".<="         -> fixity InfixN 5
  ".>"          -> fixity InfixN 5
  ".>="         -> fixity InfixN 5
  "`member`"    -> fixity InfixN 8
  "`notMember`" -> fixity InfixN 8
  ".//"         -> fixity InfixL 4
  ".&&"         -> fixity InfixR 3
  ".||"         -> fixity InfixR 2
  ".=>"         -> fixity InfixR 1

  -- servant
  ":>"          -> fixity InfixR 4
  ":<|>"        -> fixity InfixR 3
  ":-"          -> fixity InfixL 0

  -- postgresql-simple
  ":."          -> fixity InfixR 3 -- this has wildly different fixities in different libraries

  -- ?
  -- ":-:"       -> fixity _ _
  -- ":."        -> fixity _ _
  -- ":/:"       -> fixity _ _
  -- "::->"      -> fixity _ _
  -- ":<:"       -> fixity _ _
  -- ":<=:"      -> fixity _ _
  -- ":=:"       -> fixity _ _
  -- ":>:"       -> fixity _ _
  -- ":>=:"      -> fixity _ _

  str           -> case (Safe.headMay str, Safe.lastMay str) of
    (Just '\'', _) -> hardcodedFixity False (drop 1 str)
    (Just '`', Just '`') -> Just $ Fixity NoSourceText 9 InfixL
    (Just c, _) | Data.Char.isAlpha c && allowUnqualify -> hardcodedFixity False
      $ dropWhile (\x -> (Data.Char.isAlpha x || x == '.')) str
    _ -> Nothing
  where fixity a b = Just (Fixity NoSourceText b a)