{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE FlexibleContexts #-}

module Language.SimpleShell.Parser.Function
  ( functionP
  , functionP0
  )
where


import Language.SimpleShell.AST.Expr (Expr)
import Language.SimpleShell.AST.Function (Function(..))
import Language.SimpleShell.AST.Name (FunName, VarName)
import Language.SimpleShell.AST.SimpleType (SimpleType, FunSig, Typed)
import Language.SimpleShell.Parser.Common.Lexeme (symbol)
import Language.SimpleShell.Parser.Common.Root (RootParser)
import Language.SimpleShell.Parser.Expr
  ( ExprContext(ExprContext)
  , strongTermP
  , strongTermP0_
  )
import Language.SimpleShell.Parser.Name (nameP, keyword)
import Language.SimpleShell.Parser.SimpleType (simpleTypeP, forceType)

import Control.Applicative ((<|>))
import Control.Monad (guard)
import Control.Monad.Reader (ReaderT, withReaderT)
import Data.List (nub)
import Data.Map (Map)
import qualified Data.Map as Map (fromList)
import Data.Text (Text)
import Data.Tuple (swap)
import Text.Megaparsec (MonadParsec, sepBy)


type FunctionParser0 = RootParser
type FunctionParser = ReaderT (Map FunName FunSig) RootParser


functionPrefixP
  :: (MonadParsec e Text m) => m (SimpleType, FunName, [Typed VarName])
functionPrefixP = do
  keyword "function"
  t' <- simpleTypeP
  fname <- nameP
  params <- symbol "(" *> paramP `sepBy` symbol "," <* symbol ")"
  symbol "="
  return (t', fname, params)

functionP :: FunctionParser Function
functionP = do
  (t', fname, params) <- functionPrefixP
  let vars = map snd params
  checkNoDupVarNames vars
  body <- forceType t' (bodyP params) <|> fail "Function return type mismatch."
  return $ Function fname vars body

-- | First-pass function parser.
functionP0 :: FunctionParser0 (FunName, FunSig)
functionP0 = do
  (t', fname, params) <- functionPrefixP
  let ts = map fst params
  bodyP_
  return (fname, (t', ts))


checkNoDupVarNames :: [VarName] -> FunctionParser ()
checkNoDupVarNames vars =
  guard (length (nub vars) == length vars)  -- TODO: inefficient
    <|> fail "Duplicate variable name."


paramP :: (MonadParsec e Text m) => m (Typed VarName)
paramP = (,) <$> simpleTypeP <*> nameP

bodyP :: [Typed VarName] -> FunctionParser (Typed Expr)
bodyP params = withReaderT f strongTermP
  where
    f funs = ExprContext funs (Map.fromList $ map swap params)

bodyP_ :: FunctionParser0 ()
bodyP_ = strongTermP0_