{-# LANGUAGE LambdaCase #-}

import qualified Data.Map as Map
import Data.Map (Map)
import Control.Monad
import System.IO
data Expr
  = Number Integer
  | Symbol String
  | List [Expr]
  | Lambda [String] Expr Env
  | Macro [String] Expr Env
  | Builtin ([Expr] -> Eval Expr)
  | Thunk Expr Env

instance Show Expr where
  show (Number n) = show n
  show (Symbol s) = s
  show (List xs) = "(" ++ unwords (map show xs) ++ ")"
  show (Lambda _ _ _) = "<lambda>"
  show (Macro _ _ _) = "<macro>"
  show (Builtin _) = "<builtin>"
  show (Thunk _ _) = "<thunk>"

type Env = Map String Expr
type Eval a = Either String a

eval :: Env -> Expr -> Eval Expr
eval env = \case
  Number n -> return (Number n)

  Symbol s ->
    case Map.lookup s env of
      Just (Thunk e eEnv) -> eval eEnv e  
      Just v -> return v
      Nothing -> Left ("Unbound symbol: " ++ s)

  List [] -> return (List [])

  List (f:args) -> do
    func <- eval env f
    apply env func args

  lam@(Lambda {}) -> return lam
  mac@(Macro {})  -> return mac
  bi@(Builtin {}) -> return bi
  Thunk e eEnv    -> eval eEnv e

apply :: Env -> Expr -> [Expr] -> Eval Expr

apply env (Builtin f) args = do
  argVals <- mapM (eval env) args
  f argVals

-- Lazy lambdas
apply env (Lambda params body closure) args =
  if length params /= length args
    then Left "Argument mismatch"
    else
      let thunked = zip params (map (\a -> Thunk a env) args)
          newEnv = Map.union (Map.fromList thunked) closure
       in eval newEnv body

-- Macros (no evaluation of args!)
apply env (Macro params body closure) args =
  if length params /= length args
    then Left "Macro arg mismatch"
    else do
      let newEnv = Map.union (Map.fromList (zip params args)) closure
      expanded <- eval newEnv body
      eval env expanded

apply _ _ _ = Left "this aint a function!"


numOp f [Number a, Number b] = return $ Number (f a b)
numOp _ _ = Left "i said give me 2 numbers!"

builtins :: Env
builtins = Map.fromList
  [ ("+", Builtin (numOp (+)))
  , ("-", Builtin (numOp (-)))
  , ("*", Builtin (numOp (*)))
  , ("/", Builtin (numOp div))
  , ("list", Builtin (\xs -> return $ List xs))
  , ("car", Builtin $ \case
        [List (x:_)] -> return x
        _ -> Left "car error")
  , ("cdr", Builtin $ \case
        [List (_:xs)] -> return (List xs)
        _ -> Left "cdr error")
  ]

special :: Env
special = Map.fromList
  [ ("lambda", Builtin lambdaForm)
  , ("macro", Builtin macroForm)
  , ("if", Builtin ifForm)
  , ("quote", Builtin quoteForm)
  ]

lambdaForm [List params, body] =
  return $ Lambda (map getSym params) body Map.empty
lambdaForm _ = Left "lambda error"

macroForm [List params, body] =
  return $ Macro (map getSym params) body Map.empty
macroForm _ = Left "macro error"

ifForm [cond, t, f] = do
  v <- eval globalEnv cond
  case v of
    Number 0 -> eval globalEnv f
    _ -> eval globalEnv t
ifForm _ = Left "if error"

quoteForm [x] = return x
quoteForm _ = Left "quote error"

getSym (Symbol s) = s
getSym _ = error "unexpected"

tokenize :: String -> [String]
tokenize [] = []
tokenize (c:cs)
  | c `elem` "()" = [c] : tokenize cs
  | c == ' ' = tokenize cs
  | otherwise =
      let (tok, rest) = span (`notElem` " ()") (c:cs)
       in tok : tokenize rest

parse :: [String] -> (Expr, [String])
parse ("(":ts) = go [] ts
  where
    go acc (")":rest) = (List (reverse acc), rest)
    go acc ts =
      let (e, ts') = parse ts
       in go (e:acc) ts'
parse (t:ts)
  | all (`elem` "-0123456789") t = (Number (read t), ts)
  | otherwise = (Symbol t, ts)
parse [] = error "parse error"

readExpr :: String -> Expr
readExpr s = fst $ parse (tokenize s)

globalEnv :: Env
globalEnv = Map.union builtins special

repl :: IO ()
repl = do
  putStr "haslisp> "
  hFlush stdout
  line <- getLine
  unless (line == ":q") $ do
    let expr = readExpr line
    case eval globalEnv expr of
      Left err -> putStrLn ("Error: " ++ err)
      Right val -> print val
    repl

main :: IO ()
main = repl