module Parser
  ( AstNode(..)
  , parse
  ) where

import qualified Data.Map.Strict as Map
import Lexer (Token(..), TokenType)
import qualified Lexer as L

-- AST
data AstNode
  = Abstraction { parameter :: AstNode, term :: AstNode }
  | Application { lhs :: AstNode, rhs :: AstNode }
  | Identifier String
  | Epsilon
  deriving (Eq)

instance Show AstNode where
  show (Application l r) = "(" ++ show l ++ " " ++ show r ++ ")"
  show (Abstraction p t) = "λ" ++ show p ++ "." ++ show t
  show (Identifier s)    = s
  show Epsilon           = "ε"

reduce :: AstNode -> Map.Map String AstNode -> AstNode
reduce (Abstraction p t) env =
  case p of
    Identifier pid ->
      let env' = Map.delete pid env
      in case t of
           Application tl tr
             | tr == p && not (isFree tl p) -> reduce tl env'
           _ -> Abstraction p (reduce t env')
    _ -> error "Invalid abstraction"

reduce (Application l r) env =
  case l of
    Abstraction p body ->
      case p of
        Identifier pid ->
          let env' = Map.insert pid (reduce r env) env
          in reduce body env'
        _ -> error "Invalid application"
    _ -> Application (reduce l env) (reduce r env)

reduce (Identifier id_) env =
  case Map.lookup id_ env of
    Just x  -> x
    Nothing -> Identifier id_

reduce Epsilon _ = Epsilon

isFree :: AstNode -> AstNode -> Bool
isFree (Abstraction p t) id_ = id_ /= p && isFree t id_
isFree (Application l r) id_ = isFree l id_ || isFree r id_
isFree (Identifier s)    id_ = Identifier s == id_
isFree _                 _   = False

-- Parser state
data ParseState = ParseState
  { tokens :: [Token]
  , idx    :: Int
  }

type ParseResult = Either String AstNode

peekTok :: ParseState -> Either String Token
peekTok ps = case drop (idx ps) (tokens ps) of
  []    -> Left "Out of tokens"
  (t:_) -> Right t

advanceP :: ParseState -> ParseState
advanceP ps = ps { idx = idx ps + 1 }

consume :: ParseState -> TokenType -> Either String (Token, ParseState)
consume ps expected = do
  t <- peekTok ps
  if tokType t /= expected
    then Left $ "Unexpected token. Expected " ++ show expected ++
                ", but found " ++ show (tokType t)
    else Right (t, advanceP ps)

parseIdentifier :: ParseState -> Either String (AstNode, ParseState)
parseIdentifier ps = do
  (t, ps') <- consume ps L.Identifier
  return (Identifier (tokLexeme t), ps')

parseAbstraction :: ParseState -> Either String (AstNode, ParseState)
parseAbstraction ps = do
  (_, ps1) <- consume ps L.Lambda
  (ident, ps2) <- parseIdentifier ps1
  (_, ps3) <- consume ps2 L.Dot
  (t, ps4) <- parseTerm ps3
  return (Abstraction ident t, ps4)

parseAtom :: ParseState -> Maybe (Either String (AstNode, ParseState))
parseAtom ps = case peekTok ps of
  Left e  -> Just (Left ("Error parsing atom. " ++ e))
  Right t -> case tokType t of
    L.Identifier -> Just (parseIdentifier ps)
    L.OpenParen  -> Just (pGT ps)
    _            -> Nothing

pA :: ParseState -> Either String (AstNode, ParseState)
pA ps = do
  (l0, ps') <- parseAtomOrFail ps
  go l0 ps'
  where
    parseAtomOrFail s = case parseAtom s of
      Just r  -> r
      Nothing -> Left "Expected atom"
    go lhs_ ps_ = case parseAtom ps_ of
      Nothing       -> Right (lhs_, ps_)
      Just (Left e) -> Left e
      Just (Right (r, ps_')) -> go (Application lhs_ r) ps_'

pGT :: ParseState -> Either String (AstNode, ParseState)
pGT ps = do
  (_, ps1) <- consume ps L.OpenParen
  (t, ps2) <- parseTerm ps1
  (_, ps3) <- consume ps2 L.CloseParen
  return (t, ps3)

Epsilon :: ParseState -> Either String (AstNode, ParseState)
Epsilon ps = Right (Epsilon, advanceP ps)

parseTerm :: ParseState -> Either String (AstNode, ParseState)
parseTerm ps = do
  t <- peekTok ps
  case tokType t of
    L.EOF        -> Epsilon ps
    L.CloseParen -> Epsilon ps
    L.Lambda     -> parseAbstraction ps
    _            -> pA ps

parse :: [Token] -> Either String AstNode
parse toks = do
  (ast, _) <- parseTerm (ParseState toks 0)
  return ast