module Parser where

import Lexer
import Research
import Text.Megaparsec
import Text.Megaparsec.Char
import Data.Void

type Parser = Parsec Void [LToken]
data SaplingAST
  = SVar String
  | SInt Int
  | SStr String
  | SList [SaplingAST]
  | SFunc String [String] SaplingAST
  | SApp SaplingAST [SaplingAST]
  | TLeaf
  | TStem SaplingAST
  | TFork SaplingAST SaplingAST
  deriving (Show, Eq, Ord)

parseSapling :: String -> SaplingAST
parseSapling input = case runParser saplingLexer "" input of
  Left err -> error "RIP"
  Right tokens -> case runParser parseExpression "" tokens of
    Left err -> error "RIP"
    Right ast -> ast

scnParser :: Parser ()
scnParser = skipMany (satisfy isNewline)

parseExpression :: Parser SaplingAST
parseExpression = choice
  [ try parseFunction
  , try parseApplication
  , parseTreeTerm
  , parseLiteral
  , parseListLiteral
  ]

parseFunction :: Parser SaplingAST
parseFunction = do
  LIdentifier name <- satisfy isIdentifier
  args <- many (satisfy isIdentifier)
  satisfy (== LAssign)
  body <- parseExpression
  return (SFunc name (map getIdentifier args) body)

parseApplication :: Parser SaplingAST
parseApplication = do
  func <- parseAtomic
  args <- many parseAtomic
  return (SApp func args)

getIdentifier :: LToken -> String
getIdentifier (LIdentifier name) = name
getIdentifier _ = error "Expected identifier"

parseTreeTerm :: Parser SaplingAST
parseTreeTerm = do
  base <- parseTreeLeafOrParenthesized
  rest <- many parseTreeLeafOrParenthesized
  pure $ foldl combine base rest
  where
    combine acc next = case acc of
      TLeaf     -> TStem next
      TStem t   -> TFork t next
      TFork _ _ -> TFork acc next

parseTreeLeafOrParenthesized :: Parser SaplingAST
parseTreeLeafOrParenthesized = choice
  [ between (satisfy (== LOpenParen)) (satisfy (== LCloseParen)) parseTreeTerm
  , satisfy isKeywordT *> pure TLeaf
  ]

foldTree :: [SaplingAST] -> SaplingAST
foldTree [] = TLeaf
foldTree [x] = x
foldTree (x:y:rest) = TFork x (foldTree (y:rest))

parseAtomic :: Parser SaplingAST
parseAtomic = choice
  [ parseVar
  , parseLiteral
  , parseListLiteral
  , between (satisfy (== LOpenParen)) (satisfy (== LCloseParen)) parseExpression
  ]

parseLiteral :: Parser SaplingAST
parseLiteral = choice
  [ parseIntLiteral
  , parseStrLiteral
  ]

parseListLiteral :: Parser SaplingAST
parseListLiteral = do
  satisfy (== LOpenBracket) 
  elements <- sepEndBy parseExpression scnParser
  satisfy (== LCloseBracket)
  return (SList elements)

parseVar :: Parser SaplingAST
parseVar = do 
  LIdentifier name <- satisfy isIdentifier
  return (SVar name)

parseIntLiteral :: Parser SaplingAST
parseIntLiteral = do
  LIntegerLiteral value <- satisfy isIntegerLiteral
  return (SInt value)

parseStrLiteral :: Parser SaplingAST
parseStrLiteral = do
  LStringLiteral value <- satisfy isStringLiteral
  return (SStr value)

isKeywordT       (LKeywordT)         = True
isKeywordT                        _  = False

isIdentifier     (LIdentifier     _) = True
isIdentifier                      _  = False

isIntegerLiteral (LIntegerLiteral _) = True
isIntegerLiteral                  _  = False

isStringLiteral  (LStringLiteral  _) = True
isStringLiteral                   _  = False

isNewline        (LNewline)          = True
isNewline                         _  = False