{-# LANGUAGE FlexibleInstances   #-}
{-# LANGUAGE PatternGuards       #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -fno-warn-overlapping-patterns -fno-warn-incomplete-patterns #-}

-- | The contents of this module originate from module
--  [HSE.Bracket](https://github.com/ndmitchell/hlint/blob/master/src/HSE/Bracket.hs)
--  in Neil Mitchell's HLint

module Language.Haskell.Exts.Bracket
  ( Brackets(..)
  , paren
  , transformBracket
  , rebracket1
  , appsBracket
  ) where

import           Control.Monad.Trans.State
import           Data.Data
import           Data.Default
import           Data.Generics.Uniplate.Data
import           Language.Haskell.Exts.Syntax
import           Language.Haskell.Exts.Util.Internal


class Brackets a where
    remParen :: a -> Maybe a -- remove one paren, or Nothing if there is no paren
    addParen :: a -> a -- write out a paren

    -- | Is this item lexically requiring no bracketing ever
    --   i.e. is totally atomic
    isAtom :: a -> Bool

    -- | Is the child safe free from brackets in the parent position.
    --   Err on the side of caution, True = don't know
    needBracket :: Int -> a -> a -> Bool


instance (Data l, Default l) => Brackets (Exp l) where
    remParen (Paren _ x) = Just x
    remParen _           = Nothing
    addParen = Paren def

    isAtom x = case x of
        Paren{}          -> True
        Tuple{}          -> True
        List{}           -> True
        LeftSection{}    -> True
        RightSection{}   -> True
        TupleSection{}   -> True
        RecConstr{}      -> True
        ListComp{}       -> True
        EnumFrom{}       -> True
        EnumFromTo{}     -> True
        EnumFromThen{}   -> True
        EnumFromThenTo{} -> True
        _                -> isLexeme x

    -- note: i is the index in children, not in the AST
    needBracket i parent child
        | isAtom child = False
        | InfixApp{} <- parent, App{} <- child = False
        | isSection parent, App{} <- child = False
        | Let{} <- parent, App{} <- child = False
        | ListComp{} <- parent = False
        | List{} <- parent = False
        | Tuple{} <- parent = False
        | If{} <- parent, isAnyApp child = False
        | App{} <- parent, i == 0, App{} <- child = False
        | ExpTypeSig{} <- parent, i == 0, isApp child = False
        | Paren{} <- parent = False
        | isDotApp parent, isDotApp child, i == 1 = False
        | RecConstr{} <- parent = False
        | RecUpdate{} <- parent, i /= 0 = False
        | Case{} <- parent, i /= 0 || isAnyApp child = False
        | Lambda{} <- parent, i == length (universeBi parent :: [Pat l]) - 1 = False -- watch out for PViewPat
        | Do{} <- parent = False
        | otherwise = True


instance Default l => Brackets (Type l) where
    remParen (TyParen _ x) = Just x
    remParen _             = Nothing
    addParen = TyParen def

    isAtom x = case x of
        TyParen{} -> True
        TyTuple{} -> True
        TyList{}  -> True
        TyVar{}   -> True
        TyCon{}   -> True
        _         -> False

    needBracket _ parent child
        | isAtom child = False
-- a -> (b -> c) is not a required bracket, but useful for documentation about arity etc.
--        | TyFun{} <- parent, i == 1, TyFun{} <- child = False
        | TyFun{} <- parent, TyApp{} <- child = False
        | TyTuple{} <- parent = False
        | TyList{} <- parent = False
        | TyInfix{} <- parent, TyApp{} <- child = False
        | TyParen{} <- parent = False
        | otherwise = True


instance Default l => Brackets (Pat l) where
    remParen (PParen _ x) = Just x
    remParen _            = Nothing
    addParen = PParen def

    isAtom x = case x of
        PParen{}    -> True
        PTuple{}    -> True
        PList{}     -> True
        PRec{}      -> True
        PVar{}      -> True
        PApp _ _ [] -> True
        PWildCard{} -> True
        _           -> False

    needBracket _ parent child
        | isAtom child = False
        | PTuple{} <- parent = False
        | PList{} <- parent = False
        | PInfixApp{} <- parent, PApp{} <- child = False
        | PParen{} <- parent = False
        | otherwise = True


-- | Add a Paren around something if it is not atomic
paren :: (Data l, Default l) => Exp l -> Exp l
paren x = if isAtom x then x else addParen x


-- | Descend, and if something changes then add/remove brackets appropriately
descendBracket :: (Data l, Default l) => (Exp l -> (Bool, Exp l)) -> Exp l -> Exp l
descendBracket op x = descendIndex g x
    where
        g i y = if a then f i b else b
            where (a,b) = op y

        f i (Paren _ y) | not $ needBracket i x y = y
        f i y           | needBracket i x y = addParen y
        f _ y           = y


transformBracket :: (Data l, Default l) => (Exp l -> Maybe (Exp l)) -> Exp l -> Exp l
transformBracket op = snd . g
    where
        g = f . descendBracket g
        f x = maybe (False,x) ((,) True) (op x)


-- | Add/remove brackets as suggested needBracket at 1-level of depth
rebracket1 :: (Data l, Default l) => Exp l -> Exp l
rebracket1 = descendBracket (\x -> (True,x))


-- a list of application, with any necessary brackets
appsBracket :: (Data l, Default l) => [Exp l] -> Exp l
appsBracket = foldl1 (\x -> rebracket1 . App def x)

descendIndex :: Data a => (Int -> a -> a) -> a -> a
descendIndex f x = flip evalState 0 $ flip descendM x $ \y -> do
    i <- get
    modify (+1)
    return $ f i y