module Language.Parser (parseModule) where

import Language.AST
import Language.Lexer (Token(..))
import qualified Language.Lexer as L

type ParseResult a = Either String a

parseModule :: String -> String -> ParseResult Module
parseModule name src = do
  let tokens = L.lex src
  stmts <- parseStmts tokens
  Right (Module name stmts)

parseStmts :: [Token] -> ParseResult [Stmt]
parseStmts [TEOF] = Right []
parseStmts [] = Right []
parseStmts tokens = do
  (stmt, rest) <- parseStmt tokens
  remaining <- parseStmts rest
  Right (stmt : remaining)

parseStmt :: [Token] -> ParseResult (Stmt, [Token])
parseStmt (TKeyword "def" : TIdent name : rest) = do
  let (params, afterParams) = parseParams rest
  case afterParams of
    (TEquals : exprTokens) -> do
      (expr, remaining) <- parseExpr exprTokens
      Right (SDef (Ident name) params expr, dropSemi remaining)
    _ -> Left $ "Expected '=' after def " ++ name
parseStmt (TKeyword "assert" : rest) = do
  (expr, remaining) <- parseExpr rest
  Right (SAssert expr, dropSemi remaining)
parseStmt (TKeyword "prove" : TIdent name : TColon : rest) = do
  (expr, remaining) <- parseExpr rest
  Right (SProve (Ident name) expr, dropSemi remaining)
parseStmt (t:_) = Left $ "Unexpected token: " ++ show t

parseParams :: [Token] -> ([Ident], [Token])
parseParams (TLParen : rest) = go rest []
  where
    go (TRParen : r) acc = (reverse acc, r)
    go (TIdent n : r) acc = go r (Ident n : acc)
    go r acc = (reverse acc, r)
parseParams tokens = ([], tokens)

parseExpr :: [Token] -> ParseResult (Expr, [Token])
parseExpr (TKeyword "true" : rest) = Right (ELit True, rest)
parseExpr (TKeyword "false" : rest) = Right (ELit False, rest)
parseExpr (TIdent name : TLParen : rest) = do
  (args, remaining) <- parseArgs rest
  Right (EApp (Ident name) args, remaining)
parseExpr (TIdent name : rest) = Right (EVar (Ident name), rest)
parseExpr (TLParen : rest) = do
  (left, afterLeft) <- parseExpr rest
  case afterLeft of
    (TNand : rest2) -> do
      (right, afterRight) <- parseExpr rest2
      case afterRight of
        (TRParen : remaining) -> Right (ENand left right, remaining)
        _ -> Left "Expected ')' after NAND expression"
    (TRParen : remaining) -> Right (left, remaining)
    _ -> Left "Expected '|' or ')' in expression"
parseExpr tokens = Left $ "Cannot parse expression at: " ++ show (take 3 tokens)

parseArgs :: [Token] -> ParseResult ([Expr], [Token])
parseArgs (TRParen : rest) = Right ([], rest)
parseArgs tokens = do
  (expr, remaining) <- parseExpr tokens
  case remaining of
    (TRParen : rest) -> Right ([expr], rest)
    _ -> do
      (moreArgs, finalRest) <- parseArgs remaining
      Right (expr : moreArgs, finalRest)

dropSemi :: [Token] -> [Token]
dropSemi (TSemicolon : rest) = rest
dropSemi tokens = tokens