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 parseExpression "" (lexSapling input) of
  Left err -> error "Failed to parse input"
  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 <- parseAtomicBase
  args <- many parseAtomic
  case func of
    TLeaf | not (null args) && all isTreeTerm args -> fail "Not an application, defer to Tree Calculus"
    _ -> return (SApp func args)

isTreeTerm :: SaplingAST -> Bool
isTreeTerm TLeaf = True
isTreeTerm (TStem _) = True
isTreeTerm (TFork _ _) = True
isTreeTerm _ = False

parseAtomicBase :: Parser SaplingAST
parseAtomicBase = choice
  [ parseVar
  , parseTreeLeaf
  ]

parseTreeLeaf :: Parser SaplingAST
parseTreeLeaf = satisfy isKeywordT *> pure TLeaf

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
  , parseTreeLeafOrParenthesized
  , parseLiteral
  , parseListLiteral
  , between (satisfy (== LOpenParen)) (satisfy (== LCloseParen)) parseExpression
  ]

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

parens :: Parser SaplingAST -> Parser SaplingAST
parens p = do
  satisfy (== LOpenParen)
  result <- p
  satisfy (== LCloseParen)
  return result

parseListLiteral :: Parser SaplingAST
parseListLiteral = do
  satisfy (== LOpenBracket)
  elements <- many parseListItem
  satisfy (== LCloseBracket)
  return (SList elements)

parseListItem :: Parser SaplingAST
parseListItem = parseGroupedItem <|> parseSingleItem

parseGroupedItem :: Parser SaplingAST
parseGroupedItem = do
  satisfy (== LOpenParen) 
  inner <- parseExpression
  satisfy (== LCloseParen)
  return inner

parseSingleItem :: Parser SaplingAST
parseSingleItem = do
  token <- satisfy isListItem 
  case token of
    LIdentifier name -> return (SVar name)
    LKeywordT        -> return TLeaf
    _                -> fail "Unexpected token in list item"

isListItem :: LToken -> Bool
isListItem (LIdentifier _) = True
isListItem LKeywordT = True
isListItem _ = False

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)

parseMulti :: String -> [SaplingAST]
parseMulti input =
  let nonEmptyLines = filter (not . null) (lines input) 
  in map parseSapling nonEmptyLines

-- Boolean Helpers
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