{-# LANGUAGE FlexibleContexts, FlexibleInstances, KindSignatures, RankNTypes, RecordWildCards,
             StandaloneDeriving, UndecidableInstances #-}
module Main where

import Control.Applicative ((<|>), empty)
import Control.Arrow (first)
import Control.Monad (guard, join)
import Data.Char (isSpace)
import Data.List (isPrefixOf)
import Data.Functor.Compose (Compose(..))
import Data.Functor.Identity (Identity(..))
import Data.Maybe (mapMaybe)
import Data.Semigroup ((<>))
import System.Environment (getArgs)
import Text.Parser.Token (TokenParsing(..), symbol)
import qualified Text.Grampa
import Text.Grampa (TokenParsing(someSpace), LexicalParsing(lexicalComment, lexicalWhiteSpace, someLexicalSpace),
                    GrammarBuilder, ParseResults,
                    fixGrammar, parseComplete,
                    char, identifier, keyword, takeCharsWhile)
import Text.Grampa.ContextFree.SortedMemoizing.Transformer.LeftRecursive (ParserT, lift, tmap)
import qualified Boolean
import Boolean(Boolean(..))

-- |
-- >>> simplifiedSource "True && [comment] x"
-- "[comment] x"
-- >>> simplifiedSource "False || [comment1] (True || [comment2] x)"
-- "[comment1] True "
-- >>> simplifiedSource "False || [^ trailing comment] [leading comment] (True || [comment2] x)"
-- "[leading comment] True "
-- >>> simplifiedSource "False || [^ trailing comment] [leading comment] (True [operator leading] || [comment2] x)"
-- "[leading comment] True "
-- >>> simplifiedSource "([^1][1] True [^2][2] && [^3][3] x [^4][4])[^5][5] || [^6][6] False [^7][7]"
-- "[3] x [^4][^5]"

type Parser = ParserT ((,) [Ignorables])

data AST f = And (f (AST f)) (f (AST f))
           | Or (f (AST f)) (f (AST f))
           | Not (f (AST f))
           | Literal Bool
           | Variable String

deriving instance (Show (f (AST f)), Show (f Bool), Show (f String)) => Show (AST f)

instance Boolean.BooleanDomain (ParsedWrap (AST ParsedWrap)) where
   and = binary And
   or = binary Or
   not = bare . Not
   true = bare (Literal True)
   false = bare (Literal False)

binary :: (ParsedWrap (AST ParsedWrap) -> ParsedWrap (AST ParsedWrap) -> AST ParsedWrap)
       -> ParsedWrap (AST ParsedWrap) -> ParsedWrap (AST ParsedWrap) -> ParsedWrap (AST ParsedWrap)
binary f a b = bare (f a b)

type ParsedWrap = (,) ParsedIgnorables
type NodeWrap = (,) AttachedIgnorables
data AttachedIgnorables = Attached Ignorables AttachedIgnorables Ignorables
                        | Blank
                        | AttachedToOperator Ignorables Ignorables
                        | Parenthesized Ignorables AttachedIgnorables Ignorables
                        deriving Show
data ParsedIgnorables = Trailing Ignorables
                      | OperatorTrailing [Ignorables]
                      | ParenthesesTrailing Ignorables ParsedIgnorables Ignorables
                      deriving Show
type Ignorables = [Either WhiteSpace Comment]
newtype Comment    = Comment{getComment :: String} deriving Show
newtype WhiteSpace = WhiteSpace String deriving Show

type Grammar = Boolean.Boolean (ParsedWrap (AST ParsedWrap))

main :: IO ()
main = do args <- concat <$> getArgs
          let tree = parse args
          print tree
          case tree
             of Right [parsed] -> do let rearranged = completeRearranged mempty parsed
                                     print rearranged
                                     putStrLn (showSource $ simplified rearranged)
                other -> error (show other)

parse :: String -> ParseResults String [ParsedWrap (AST ParsedWrap)]
parse = getCompose . fmap snd . getCompose . Boolean.expr . parseComplete (fixGrammar grammar)

simplifiedSource :: String -> String
simplifiedSource input = case (parse input)
                         of Right [parsed] -> showSource (simplified $ completeRearranged mempty parsed)
                            other -> error (show other)

class ShowSource a where
   showSource :: a -> String
   showsSourcePrec :: Int -> a -> String -> String
   showSource a = showsSourcePrec 0 a mempty

instance ShowSource (NodeWrap (AST NodeWrap)) where
   showsSourcePrec prec (Attached lead Blank follow, node) rest =
      whiteString lead <> showsSourcePrec prec node (whiteString follow <> rest)
   showsSourcePrec prec (Parenthesized lead ws follow, node) rest =
      whiteString lead <> showsSourcePrec 0 (ws, node) (whiteString follow <> rest)

instance ShowSource (AST NodeWrap) where
   showsSourcePrec prec (Or left right) rest
      | prec < 1 = showsSourcePrec 1 left ("||" <> showsSourcePrec 0 right rest)
   showsSourcePrec prec (And left right) rest
      | prec < 2 = showsSourcePrec 2 left ("&&" <> showsSourcePrec 1 right rest)
   showsSourcePrec prec (Not expr) rest
      | prec < 3 = "not" <> showsSourcePrec 2 expr rest
   showsSourcePrec _ (Literal True) rest = "True" <> rest
   showsSourcePrec _ (Literal False) rest = "False" <> rest
   showsSourcePrec _ (Variable name) rest = name <> rest
   showsSourcePrec _ node rest =   "(" <> showsSourcePrec 0 node (")" <> rest)

completeRearranged :: Ignorables -> ParsedWrap (AST ParsedWrap) -> NodeWrap (AST NodeWrap)
completeRearranged ws node
   | ((ws', node'), trailing) <- rearranged ws node = (embed [] ws' trailing, node')

rearranged :: Ignorables -> ParsedWrap (AST ParsedWrap) -> (NodeWrap (AST NodeWrap), Ignorables)
rearranged leftover (Trailing follow, node)
   | (follow', lead') <- splitDirections follow, (lead'', node', follow'') <- rearrangedChildren [] node =
        ((Attached (leftover <> lead'') Blank (follow'' <> follow'), node'), lead')
rearranged leftover (OperatorTrailing [[], follow], node)
   | (follow', lead') <- splitDirections follow,
     (lead'', node', follow'') <- rearrangedChildren [[], lead'] node =
        ((Attached leftover (AttachedToOperator lead'' follow') [], node'), follow'')
rearranged leftover (ParenthesesTrailing lead ws follow, node)
   | (follow', lead') <- splitDirections follow, (ws', node') <- completeRearranged lead (ws, node) =
        ((Parenthesized leftover ws' follow', node'), lead')

embed leading Blank trailing = Attached leading Blank trailing
embed leading (Attached lead inside follow) trailing = embed (leading <> lead) inside (follow <> trailing)
embed leading (AttachedToOperator lead follow) trailing = AttachedToOperator (leading <> lead) (follow <> trailing)
embed leading (Parenthesized lead inside follow) trailing = Parenthesized (leading <> lead) inside (follow <> trailing)

rearrangedChildren :: [Ignorables] -> AST ParsedWrap -> (Ignorables, AST NodeWrap, Ignorables)
rearrangedChildren [left, right] (And a b)
   | (a', follow1) <- rearranged left a,
     (b', follow2) <- rearranged right b = (follow1, And a' b', follow2)
rearrangedChildren [left, right] (Or a b)
   | (a', follow1) <- rearranged left a,
     (b', follow2) <- rearranged right b = (follow1, Or a' b', follow2)
rearrangedChildren [leftover] (Not a)
   | (a', follow) <- rearranged leftover a = ([], Not a', follow)
rearrangedChildren [] (Literal a) = ([], Literal a, [])
rearrangedChildren [] (Variable name) = ([], Variable name, [])

-- | Separates the whitespace and comments that refer to the preceding construct.
splitDirections :: Ignorables -> (Ignorables, Ignorables)
splitDirections = span (either (const True) (isPrefixOf "^" . getComment))

-- | Simplifies the given expression according to the laws of Boolean algebra.
simplified :: NodeWrap (AST NodeWrap) -> NodeWrap (AST NodeWrap)
simplified e@(_, Literal{}) = e
simplified e@(_, Variable{}) = e
simplified (a, Not e) = case simplified e
                        of (b, Literal True) -> (raise a b, Literal False)
                           (b, Literal False) -> (raise a b, Literal True)
                           e' -> (a, Not e')
simplified (a, And l r) = case (simplified l, simplified r)
                          of ((b, Literal False), _) -> (raise a b, Literal False)
                             ((b, Literal True), (c, r')) -> (raise a c, r')
                             (_, (b, Literal False)) -> (raise a b, Literal False)
                             ((b, l'), (c, Literal True)) -> (raise a b, l')
                             (l', r') -> (a, And l' r')
simplified (a, Or l r) =  case (simplified l, simplified r)
                          of ((b, Literal False), (c, r')) -> (raise a c, r')
                             ((b, Literal True), _) -> (raise a b, Literal True)
                             ((b, l'), (c, Literal False)) -> (raise a b, l')
                             (_, (b, Literal True)) -> (raise a b, Literal True)
                             (l', r') -> (a, Or l' r')

raise :: AttachedIgnorables -> AttachedIgnorables -> AttachedIgnorables
raise Blank arg = arg
raise op Blank = op
raise AttachedToOperator{} arg = arg
raise (Parenthesized opl op opr) arg = Parenthesized opl (raise op arg) opr
raise (Attached opl inside opr) (Parenthesized l arg r) =
   Parenthesized (comments opl <> l) (raise inside arg) (r <> comments opr)
-- raise (AttachedToOperator opl opr) (Parenthesized l arg r) = Parenthesized (comments opl <> l) arg (r <> comments opr)
-- raise (AttachedToOperator opl opr) (Attached argl inside argr) =
--    Attached (comments opl <> argl) inside (argr <> comments opr)

comments :: Ignorables -> Ignorables
comments = mapMaybe (either (const Nothing) (Just . Right))

whiteString :: Ignorables -> String
whiteString (Left (WhiteSpace ws) : rest) = ws <> whiteString rest
whiteString (Right (Comment c) : rest) = "[" <> c <> "]" <> whiteString rest
whiteString [] = ""

grammar :: GrammarBuilder Grammar Grammar Parser String
grammar Boolean{..} = Boolean{
   expr= term
         <|> operatorTrailingWhiteSpace [mempty] (Boolean.or <$> term <* symbol "||" <*> expr),
   term= factor
         <|> operatorTrailingWhiteSpace [mempty] (Boolean.and <$> factor <* symbol "&&" <*> term),
   factor= trailingWhiteSpace (keyword "True" *> pure Boolean.true
                                 <|> keyword "False" *> pure Boolean.false
                                 <|> bare . Variable <$> identifier)
           <|> operatorTrailingWhiteSpace [] (keyword "not" *> (Boolean.not <$> factor))
           <|> parenthesizedWhiteSpace (symbol "(" *> expr <* symbol ")")}

bare :: a -> ParsedWrap a
bare a = (Trailing [], a)

operatorTrailingWhiteSpace :: [Ignorables] -> Parser Grammar String (ParsedWrap (AST ParsedWrap))
                           -> Parser Grammar String (ParsedWrap (AST ParsedWrap))
trailingWhiteSpace, parenthesizedWhiteSpace
   :: Parser Grammar String (ParsedWrap (AST ParsedWrap)) -> Parser Grammar String (ParsedWrap (AST ParsedWrap))
trailingWhiteSpace = tmap store
   where store ([ws], (Trailing ws', a)) = (mempty, (Trailing $ ws' <> ws, a))
operatorTrailingWhiteSpace prefix = tmap store
   where store (wss, (Trailing [], a)) = (mempty, (OperatorTrailing (prefix <> wss), a))
parenthesizedWhiteSpace = tmap store
   where store ([ws,ws'], (aws, a)) = ([], (ParenthesesTrailing ws aws ws', a))

instance {-# OVERLAPS #-} TokenParsing (Parser Grammar String) where
   someSpace = someLexicalSpace
   token p = p <* lexicalWhiteSpace

instance {-# OVERLAPS #-} LexicalParsing (Parser Grammar String) where
   lexicalWhiteSpace = tmap (first (\ws-> [concat ws])) $
                       do ws <- takeCharsWhile isSpace
                          lift ([[Left $ WhiteSpace ws]], ())
                          (lexicalComment *> lexicalWhiteSpace <|> pure ())
   lexicalComment = do char '['
                       comment <- takeCharsWhile (/= ']')
                       char ']'
                       lift ([[Right $ Comment comment]], ())
   identifierToken p = do ident <- p
                          guard (ident /= "True" && ident /= "False")
                          lexicalWhiteSpace
                          pure ident