module Compiler.AST.Common where

import Common
import Control.Applicative
import Control.Monad
import Data.List.NonEmpty as NE
import Data.Text

import Compiler.AST.Parser.Common
import Compiler.Lexer
import Compiler.Lexer.Comments
import Parser
import Test.Common

parseComment :: AstParser Comment
parseComment = parseToken "comment" (\case
  TkComment c -> Just c
  _           -> Nothing)

parseToken :: Text -> (TokenRaw -> Maybe a) -> AstParser a
parseToken name fn = ParserM name $ liftModifier $ \s -> case s of
  ((Token tr _ _) : rst) -> pure $ case fn tr of
    Just a  -> (Right a, rst)
    Nothing -> (Left CantHandle, s)
  _ -> pure (Left CantHandle, s)

parseIdentifier :: AstParser Identifier
parseIdentifier =  parseToken "Identifier" (\case
  TkIdentifier i -> Just i
  _              -> Nothing)

parseDelimeter :: Delimeter -> AstParser Delimeter
parseDelimeter dl =  parseToken (toSource dl) (\case
  TkDelimeter l -> if l == dl then Just l else Nothing
  _             -> Nothing)

parseKeyword :: Keyword -> AstParser Keyword
parseKeyword kw =  parseToken (toSource kw) (\case
  TkKeyword k -> if kw == k then Just k else Nothing
  _           -> Nothing)

parseOperator :: AstParser Operator
parseOperator = parseToken "Operator" (\case
  TkOperator k -> Just k
  _            -> Nothing)

parseOperator' :: Operator -> AstParser ()
parseOperator' op = parseToken "Operator" (\case
  TkOperator k -> if k == op then Just () else Nothing
  _            -> Nothing)

mandatoryNewline :: AstParser ()
mandatoryNewline = parseToken "NL" (\case
  TkWhitespace (NewLine _) -> Just ()
  _                        -> Nothing)

whitespace :: AstParser ()
whitespace = void $ optional $ parseToken "WS" (\case
  TkWhitespace (Space _) -> Just ()
  TkWhitespace (Tab _)   -> Just ()
  _                      -> Nothing)

whitespaceNL :: AstParser ()
whitespaceNL = void $ optional $ parseToken "NL" (\case
  TkWhitespace (NewLine _) -> Just ()
  _                        -> Nothing)

surroundWs :: AstParser a -> AstParser a
surroundWs p@(ParserM name _) = nameParser (name <> " surrounded by WS or NL ") $ do
  whitespaceOrNl
  a <- p
  whitespaceOrNl
  pure a

whitespaceOrNl :: AstParser ()
whitespaceOrNl = void $ many $ parseToken "WS" (\case
  TkWhitespace (Space _)   -> Just ()
  TkWhitespace (Tab _)     -> Just ()
  TkWhitespace (NewLine _) -> Just ()
  _                        -> Nothing)

surroundWs_ :: AstParser a -> AstParser ()
surroundWs_ p@(ParserM name _) = nameParser (name <> " surrounded by WS") $ void $ surroundWs p

mandatory :: AstParser a -> AstParser a
mandatory (ParserM name fn) = ParserM ("mandatory " <> name) $ \s -> do
  case (astInput s) of
    [] -> pure (Left $ FatalError $ CustomError $ "Expecting " <> name <> " but ran out of input ", s)
    _ -> fn s  >>= \case
        r@(Right _, _) -> pure r
        (Left fe@(FatalError _), _)    ->  pure (Left fe, s)
        (Left fe@(FatalErrorWithLocation _ _), _)    ->  pure (Left fe, s)
        (Left e, s'') -> case astInput s'' of
          [] -> pure (Left e, s)
          (h: _) -> pure (Left $ FatalErrorWithLocation (tkLocation h) $ CustomError $ "Expecting " <> name <> " but got " <> (pack $ show h), s)

nonEmptyGen :: Gen a -> Gen (NonEmpty a)
nonEmptyGen p = do h <- p; t <- list (linear 0 3) p; pure (h :| t)

wst :: Text
wst = toSource $ TkWhitespace (Space 1)

nlt :: Text
nlt = toSource $ TkWhitespace (NewLine 1)

parseItemListInParen :: AstParser a -> AstParser (Maybe (NonEmpty a))
parseItemListInParen itemParser = do
  surroundWs_ (parseDelimeter DlParenOpen)
  args <- optional itemParser >>= \case
    Just argHead -> do
      argsTail <- many $ do
        surroundWs_ $ parseDelimeter DlComma
        itemParser
      pure $ Just (argHead :| argsTail)
    Nothing -> pure Nothing
  surroundWs_ (parseDelimeter DlParenClose)
  pure args