module Main where
import Text.ParserCombinators.Parsec (many1, sepBy1, string, noneOf, char, digit, spaces, sepEndBy, Parser, (<|>), parse, letter, oneOf)
--import Relude hiding ((<|>))
import Prelude
import Data.Foldable (foldl1) 
import qualified Data.Set as Set
import Data.Map (Map, (!?), insert, fromList, alter)
import Control.Monad (forever, void)
import System.IO (hFlush, stdout)
import Control.Monad.State

data Term = Var String | App Term Term | Lam String Term | Let String Term | Lit Int deriving (Eq)
instance Show Term where
    show (Lit n) = show n
    show (Var x) = x
    show (App x y) = "("++show x ++ show y++")"
    show (Lam x y) = "(λ"++x++". "++show y++")"

parens p = char '(' *> spaces *> p <* spaces <* char ')'
var = Var <$> many1 letter
lam = do
  void $ char '\\'
  spaces
  vars <- many1 $ (many1 letter) <* spaces
  void $ char '.'
  spaces
  body <- term
  return (foldr Lam body vars)

elet = do
  void $ string "let"
  spaces
  name <- many1 letter
  spaces
  void $ char '='
  spaces
  body <- term
  return $ Let name body 

lit = Lit <$> read <$> many1 digit
app = foldl1 App <$> sepBy1 (var <|> lit <|> parens term) spaces

term :: Parser Term
term = elet <|> lam <|> app <|> lit <|> var 

newtype Env = Env (Map String Term) deriving (Show)

normalstep :: Term -> State Env (Maybe Term)
normalstep (App m n) = do
  (Env env) <- get
  case m of
    Lam x m -> do
      m' <- return $ evalState (normalstep m) (Env $ alter (const $ Just $ n) x env)
      case m' of
        Nothing -> return $ Just $ m
	Just m' -> return $ Just $ m'
    _ -> do 
         m' <- normalstep m
	 case m' of
           Nothing -> do
	     n' <- normalstep n
	     return $ App m <$> n'
	   Just m' -> return $ Just $ App m' n

normalstep l@(Lam x m) = do
  (Env env) <- get
  --case env !? x of
  --  Nothing -> return $ Lam x <$> evalState (normalstep m) (Env $ alter (const $ Just $ Var x) x env)
  --  Just n -> return $ evalState (normalstep m) e 
  --m' <- return $ evalState (normalstep m) (Env $ alter (const $ Just $ Var x) x env)
  m' <- return $ evalState (normalstep m) (Env $ alter (const $ Just $ Var x) x env) 
  return $ Lam x <$> m'
  --return Nothing

normalstep (Let s n) = do
  (Env env) <- get
  modify (\(Env ls) -> (Env $ alter (const $ Just n) s ls)) 
  return Nothing

normalstep (Var s) = do
  (Env env) <- get
  case env !? s of
    Nothing -> error $ "unbound variable: " ++ s
    Just v@(Var s') | s == s' -> return Nothing 
    Just term -> return $ Just $ term 

normalstep (Lit _) = return Nothing


eval term = do
  env <- get
  let (term', env') = runState (normalstep term) env
  put env'
  case term' of
    Nothing -> return term
    Just term'' -> eval term'' 
  
repl :: StateT Env IO ()
repl = forever $ do
  liftIO $ putStr "Untyped> "
  liftIO $ hFlush stdout
  contents <- liftIO $ getLine
  env <- get
  case parse term "<stdin>" contents of
    Left x -> liftIO $ print x
    Right term -> do
        let (term', env') = runState (eval term) env
        put env'
	case term' of
          (Let _ _) -> pure ()
          _ -> liftIO $ print term'

main = runStateT (repl) $ Env $ fromList []