{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TupleSections #-}

module Language.SimpleShell.Parser.Expr
  ( exprP
  , strongTermP
  , strongTermP_  -- TODO: Unused.
  )
where


import Language.SimpleShell.AST.Expr (Expr(..), TypedExpr)
import Language.SimpleShell.AST.SimpleType (SimpleType(..))
import Language.SimpleShell.Parser
  ( Parser
  , lexeme
  , symbol
  , commentFirstChars
  , lookupVar
  , lookupFun
  )
import Language.SimpleShell.Parser.Name (nameP, keyword)
import Language.SimpleShell.Parser.SimpleType (forceType)
import Control.Monad.Combinators.FailExpr
  ( Associativity(..)
  , makeExprParser
  , Operator(..)
  )

import Control.Applicative ((<|>))
import Control.Monad (void)
import Control.Monad.Combinators (manyTill)
import Data.Foldable (asum)
import Data.Text (Text)
import Text.Megaparsec (takeWhile1P, oneOf, many)
import Text.Megaparsec.Char (char)
import qualified Text.Megaparsec.Char.Lexer as L (charLiteral, decimal)


exprP :: Parser TypedExpr
exprP
   =  makeExprParser weakTermP binaryOperatorTable

weakTermP :: Parser TypedExpr
weakTermP
   =  strongTermP
  <|> unaryOpP
  <|> builtinUnaryFunP
  <|> funP

strongTermP :: Parser TypedExpr
strongTermP
   =  literalP
  <|> varP
  <|> symbol "(" *> exprP <* symbol ")"

-- | Parse a strong term--assuming its correctness--without yielding a result.
--   This basically only checks for matching parentheses.
strongTermP_ :: Parser ()
strongTermP_
   =  void literalP
  <|> varP_
  <|> symbol "(" *> void (many tok) <* symbol ")"
  where
    -- Notes:
    --  * We need to make sure to correctly handle:
    --    * parentheses
    --    * string literals
    --    * comments

    tok :: Parser ()
    tok
       =  void strLitP
      <|> symbol "(" *> void (many tok) <* symbol ")"
      <|> void (lexeme $ takeWhile1P Nothing isBoring)
      <|> void (lexeme $ oneOf commentFirstChars)
      where
        isBoring :: Char -> Bool
        isBoring = not . (`elem` "()\"" ++ commentFirstChars)


-- | Parse strong term with fixed type.
strongTermP' :: String -> SimpleType -> Parser Expr
strongTermP' errMsg t = forceType t strongTermP <|> fail errMsg


literalP :: Parser TypedExpr
literalP
   =  (IntType,)  . IntLiteral  <$> lexeme L.decimal
  <|> (StrType,)  . StrLiteral  <$> strLitP
  <|> (BoolType,) . BoolLiteral <$> boolLitP
  where
    boolLitP
       =  True  <$ keyword "true"
      <|> False <$ keyword "false"

strLitP :: Parser String
strLitP = lexeme $ char '"' *> manyTill L.charLiteral (char '"')

varP :: Parser TypedExpr
varP = do
  _ <- char '$'
  x <- nameP
  t <- lookupVar x
  return (t, Var x)

varP_ :: Parser ()
varP_ = void (char '$') <* nameP

funP :: Parser TypedExpr
funP = do
  fname <- nameP
  (t', ts) <- lookupFun fname
  args <- mapM (strongTermP' "Type mismatch with function signature.") ts
  return (t', FunCall fname args)


-- Binary operators.

type BinaryFun = Expr -> Expr -> Expr
type BinarySig = SimpleType -> SimpleType -> Maybe SimpleType

binaryOperatorTable :: [[Operator Parser TypedExpr]]
binaryOperatorTable =
  [ [ binary AssocR "||" Or  $ sameSig BoolType
    ]
  , [ binary AssocR "&&" And $ sameSig BoolType
    ]
  , [ binary AssocN "==" Eq  $ anyCmpSig
    , binary AssocN "!=" Neq $ anyCmpSig
    , binary AssocN "<=" Le  $ intCmpSig
    , binary AssocN "<"  Lt  $ intCmpSig
    , binary AssocN ">=" Ge  $ intCmpSig
    , binary AssocN ">"  Gt  $ intCmpSig
    ]
  , [ Binary AssocL addP
    , binary AssocL "-"  Sub $ sameSig IntType
    ]
  , [ binary AssocL "*"  Mul $ sameSig IntType
    , binary AssocL "/"  Div $ sameSig IntType
    ]
  ]
  where
    sameSig :: SimpleType -> BinarySig
    sameSig t t1 t2 = if t == t1 && t == t2 then Just t else Nothing

    anyCmpSig t t'
      | t == t'   = Just BoolType
      | otherwise = Nothing

    intCmpSig IntType IntType = Just BoolType
    intCmpSig _       _       = Nothing

    addP :: Parser (TypedExpr -> TypedExpr -> Maybe TypedExpr)
    addP = do
      symbol "+"
      return $ \(t1, x1) (t2, x2) ->
        if t1 == t2
        then
          case t1 of
            StrType -> Just (StrType, Concat x1 x2)
            IntType -> Just (IntType, Add    x1 x2)
            _       -> Nothing
        else Nothing

    binary
      :: Associativity -> Text -> BinaryFun -> BinarySig
      -> Operator Parser TypedExpr
    binary assoc symb op sig = Binary assoc $ do
      symbol symb
      return $ \(t1, e1) (t2, e2) -> fmap (, op e1 e2) $ sig t1 t2


-- Unary operators and builtin unary functions.

type UnaryFun = Expr -> Expr
type UnarySig = SimpleType -> Maybe SimpleType
type UnaryParser = Parser (TypedExpr -> Maybe TypedExpr)

unaryOperators, builtinUnaryFuns :: [UnaryParser]
(unaryOperators, builtinUnaryFuns) =
  ( [ unaryOp  "!"      Not     $ uniqueSig BoolType BoolType
    , unaryOp  "-"      UMinus  $ uniqueSig IntType  IntType
    ]
  , [ unaryFun "length" Length  $ uniqueSig StrType  IntType
    , unaryFun "int"    IntCast $ uniqueSig StrType  IntType
    , unaryFun "str"    StrCast $ uniqueSig IntType  StrType
    ]
  )
  where
    uniqueSig :: SimpleType -> SimpleType -> UnarySig
    uniqueSig t1 t2 t1'
      | t1 == t1' = Just t2
      | otherwise = Nothing

    unaryOp :: Text -> UnaryFun -> UnarySig -> UnaryParser
    unaryOp = unary symbol

    unaryFun :: Text -> UnaryFun -> UnarySig -> UnaryParser
    unaryFun = unary keyword

    unary :: (Text -> Parser ()) -> Text -> UnaryFun -> UnarySig -> UnaryParser
    unary symbPF symb op sig = do
      symbPF symb
      return $ \(t, e) -> fmap (, op e) $ sig t

unaryOpP, builtinUnaryFunP :: Parser TypedExpr
(unaryOpP, builtinUnaryFunP) =
  ( asum $ map (aux "unary operator"         weakTermP  ) unaryOperators
  , asum $ map (aux "builtin unary function" strongTermP) builtinUnaryFuns
  )
  where
    aux :: String -> Parser TypedExpr -> UnaryParser -> Parser TypedExpr
    aux desc argP p = do
      f <- p
      x <- argP
      case f x of
        Just x' -> return x'
        Nothing -> fail $ "Mismatching " ++ desc ++ " signature."