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 Debug.Trace (traceShowId)
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 <|> oneOf "_")
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)

fv :: Term -> Set.Set String
fv (Var x) = Set.singleton x
fv (Lam x m) = x `Set.delete` (fv m)
fv (App m1 m2) = (fv m1) `Set.union` (fv m2)
fv (Lit _) = Set.empty

pickFresh :: Set.Set String -> String -> String
pickFresh s = pickFresh'
  where
    pickFresh' n | n `Set.notMember` s = n
    pickFresh' n = pickFresh' $ n ++ "'"

subst n x v@(Var e) 
  | x == e = n
  | otherwise = v
subst n x (App m n') = App (subst n x m) (subst n x n')
subst n x l@(Lam y m)
    | x == y = l
    | y `Set.notMember` (fv n) = Lam y (subst n x m)
    | otherwise = Lam y' (subst n x (subst (Var y') y m))
  where
    y' = pickFresh ((fv n) `Set.union` (fv m)) y
subst _ _ t = t

normalstep :: Term -> State Env (Maybe Term)
normalstep (App m n) = do
  (Env env) <- get
  case m of
    Lam x m -> return $ Just $ subst n x 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 (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 $ Shadow 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 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 env') = runState (eval term) env
        put $ Env $ alter (const $ Just term') "_" $ env'
        case term' of
          (Let _ _) -> pure ()
          _ -> liftIO $ print term'

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