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

module Language.SimpleShell.Parser.Expr
  ( ExprParser
  , runExprParser
  , exprP
  , strongTermP
  , strongTermP_
  )
where


import Language.SimpleShell.AST.Expr (Expr(..), TypedExpr)
import Language.SimpleShell.AST.Name (FunName, VarName)
import Language.SimpleShell.AST.SimpleType (SimpleType(..), Typed, FunSig)
import Language.SimpleShell.Parser
  ( lexeme
  , symbol
  , commentFirstChars
  )
import Language.SimpleShell.Parser.Name (nameP, keyword)
import Language.SimpleShell.Parser.Root (RootParser)
import Language.SimpleShell.Parser.SimpleType (forceType)
import Control.Monad.Combinators.FailExpr
  ( Associativity(..)
  , makeExprParser
  , Operator(..)
  )

import Control.Applicative ((<|>))
import Control.Monad (void)
import Control.Monad.Reader (ask, ReaderT, runReaderT)
import Control.Monad.Combinators (manyTill)
import Data.Foldable (asum)
import Data.List (nub)
import Data.Map (Map)
import qualified Data.Map as Map (empty, lookup, fromList)
import Data.Text (Text, unpack)
import Data.Tuple (swap)
import Text.Megaparsec (takeWhile1P, oneOf, many)
import Text.Megaparsec.Char (char)
import qualified Text.Megaparsec.Char.Lexer as L (charLiteral, decimal)


type ExprParser = ReaderT ExprContext RootParser
type Parser = ExprParser


data ExprContext = ExprContext
  { ctxVars :: Map VarName SimpleType
  , ctxFuns :: Map FunName FunSig
  }


runExprParser :: [Typed VarName] -> ExprParser a -> RootParser a
runExprParser decls p =
  if length (nub names) == length names   -- TODO: inefficient
  then runReaderT p (ExprContext vars Map.empty)
  else fail "Duplicate variable name."
  where
    names = map snd decls
    vars = Map.fromList $ map swap decls


lookupVar :: VarName -> Parser SimpleType
lookupVar varname = do
  mt <- Map.lookup varname . ctxVars <$> ask
  case mt of
    Just t  -> return t
    Nothing -> fail $ "Undeclared variable $" ++ unpack varname

lookupFun :: FunName -> Parser FunSig
lookupFun fname = do
  msig <- Map.lookup fname . ctxFuns <$> ask
  case msig of
    Just sig -> return sig
    Nothing  -> fail $ "Undefined function " ++ unpack fname


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."