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

module Language.SimpleShell.Parser.Expr
  ( exprP
  )
where


import Language.SimpleShell.Parser (Parser, lexeme, lookupVar, lookupFun)
import Language.SimpleShell.AST.Expr (Expr(..), TypedExpr, VarName, FunName)
import Language.SimpleShell.AST.SimpleType (SimpleType(..))
import Control.Monad.Combinators.FailExpr
  ( Associativity(..)
  , makeExprParser
  , Operator(..)
  )

import Control.Applicative ((<|>))
import Control.Monad.Combinators (manyTill)
import Data.Char (isAlpha)
import Data.Foldable (asum)
import Data.Text (Text)
import Text.Megaparsec (takeWhile1P)
import Text.Megaparsec.Char (string, 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
  <|> lexeme (char '(') *> exprP <* lexeme (char ')')


-- | Parse "strong" term with fixed type.
strongTermP' :: String -> SimpleType -> Parser Expr
strongTermP' errMsg t = do
  (t', e) <- strongTermP
  if t == t'
  then return e
  else fail errMsg


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

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

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


-- TODO
varNameP :: Parser VarName
varNameP = takeWhile1P Nothing isAlpha

-- TODO
funNameP :: Parser FunName
funNameP = takeWhile1P Nothing isAlpha


-- 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
      _ <- lexeme $ char '+'
      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 symbol op sig = Binary assoc $ do
      _ <- lexeme $ string symbol
      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) =
  ( [ unary "!"      Not     $ uniqueSig BoolType BoolType
    , unary "-"      UMinus  $ uniqueSig IntType  IntType
    ]
  , [ unary "length" Length  $ uniqueSig StrType  IntType
    , unary "int"    IntCast $ uniqueSig StrType  IntType
    , unary "str"    StrCast $ uniqueSig IntType  StrType
    ]
  )
  where
    uniqueSig :: SimpleType -> SimpleType -> UnarySig
    uniqueSig t1 t2 t1'
      | t1 == t1' = Just t2
      | otherwise = Nothing

    unary :: Text -> UnaryFun -> UnarySig -> UnaryParser
    unary symbol op sig = do
      _ <- lexeme $ string symbol
      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."