{-# 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 ++ "," ++ show 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)
  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, BriDocNumbered)]
    -> 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 (OpLeaf 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 (OpLeaf 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
               (OpLeaf val)
          (InfixL, InfixL) ->
            go (StackElem fixityS ((c, opDoc) : cs) : stackR)
               inputR
               (OpLeaf 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 False Nothing Nothing

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

remSuperfluousParens :: Int -> OpTree -> OpTree
remSuperfluousParens outerFixity = \case
  x@OpLeaf{}    -> x
  x@OpUnknown{} -> x
  OpKnown paren locO locC fixity c cs ->
    OpKnown
      (paren && outerFixity > fixLevel fixity)
      locO
      locC
      fixity
      (remSuperfluousParens (fixLevel fixity) c)
      [ (op, remSuperfluousParens (fixLevel fixity) tree) | (op, tree) <- cs ]
  where fixLevel (Fixity _ i _) = i

hardcodedFixity :: Bool -> String -> Maybe Fixity
hardcodedFixity allowUnqualify = \case
  "."         -> 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

  -- 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)