content: module Interpreter where
  
  import Control.Monad (foldM)
  import Control.Monad.Except
  import Control.Monad.State
  import Data.IORef (newIORef, readIORef, writeIORef)
  import Data.Map qualified as M
  import Data.Maybe (fromMaybe)
  import Error
  import Expr
  import State
  import Stmt hiding (initialiser)
  
  -- Evaluating!
  
  eval :: Expr -> Interpreter Object Object
  eval (LiteralExpr obj) = return obj
  eval (BinaryExpr (Binary left op right)) =
    op <$> eval left <*> eval right >>= liftEither
  eval (UnaryExpr (Unary op expr)) =
    eval expr >>= liftEither . op
  eval (GroupingExpr expr) = eval expr
  eval (AssignExpr (Assign ident value)) = do
    obj <- eval value
    ctx <- get
    assigned <- liftIO $ assign ident obj ctx
    ctx' <- liftEither assigned
    put ctx'
    return obj
  eval (VarExpr ident) = do
    ctx <- get
    resolved <- liftIO $ resolve ident ctx
    liftEither resolved
  eval (LogicalExpr (Logical first op rest)) = do
    obj <- eval first
    if (op == Or && truthy obj)
      || (op == And && not (truthy obj))
      then return obj
      else eval rest
  eval (CallExpr (Call callee args)) = do
    callee' <- eval callee
    case callee' of
      FunctionObject (Function arity _ f) -> do
        if length args /= arity
          then
            throwError $
              TypeError
                (show callee' ++ "() takes " ++ show arity ++ " positional arguments but " ++ show (length args) ++ " were given")
          else do
            args' <- mapM eval args
            f args'
      ClassObject (Expr.Class name initialiser methods super) -> do
        properties <- liftIO $ newIORef M.empty
        let obj = InstanceObject $ Instance (Expr.Class name initialiser methods super) properties
        case initialiser of
          Just i -> do
            if length args /= arity i
              then
                throwError $
                  TypeError
                    (show callee' ++ " constructor takes " ++ show (arity i) ++ " positional arguments but " ++ show (length args) ++ " were given")
              else do
                bound <- liftIO $ bind i "this" obj
                bound' <- case super of
                  Just s -> liftIO $ bind bound "super" (ClassObject s)
                  Nothing -> return bound
                args' <- mapM eval args
                _ <- fun bound' args' -- could catch bad init behaviour here?
                return obj
          Nothing -> return obj
      _ -> throwError $ TypeError (show callee' ++ " object is not a callable")
  eval (GetExpr (Get expr member)) = do
    obj <- eval expr
    case obj of
      InstanceObject (Instance cls properties) -> do
        curr <- liftIO $ readIORef properties
        case M.lookup member curr of
          Just value -> return value
          Nothing -> case lookupMethod cls of
            Just (cls', method) -> do
              bound <- liftIO $ bind method "this" obj
              bound' <- case superClass cls' of
                Just s -> liftIO $ bind bound "super" (ClassObject s)
                Nothing -> return bound
              return $ FunctionObject bound'
            Nothing -> throwError $ AttributeError $ show cls ++ " object has no attribute " ++ member
      ClassObject cls -> case lookupMethod cls of
        Just (cls', method) -> do
          bound <- case superClass cls' of
            Just s -> liftIO $ bind method "super" (ClassObject s)
            Nothing -> return method
          return $ FunctionObject bound
        Nothing -> throwError $ AttributeError $ show cls ++ " object has no attribute " ++ member
      _ -> throwError $ TypeError "only instances have properties"
    where
      lookupMethod cls =
        case M.lookup member (Expr.methods cls) of
          Just method -> Just (cls, method)
          Nothing -> do
            super <- superClass cls
            lookupMethod super
  eval (SetExpr (Set expr member setter)) = do
    obj <- eval expr
    case obj of
      InstanceObject (Instance _ properties) -> do
        value <- eval setter
        curr <- liftIO $ readIORef properties
        let updated = M.insert member value curr
        liftIO $ writeIORef properties updated
        return value
      _ -> throwError $ TypeError "only instances have properties"
  
  -- Interpreting!
  
  data Result = None | Return Object deriving (Eq)
  
  interpret :: Stmt -> Interpreter Object Result
  interpret (BlockStmt stmts) = do
    modify scope
    res <- exec stmts
    modify descope
    return res
    where
      exec [] = return None
      exec [s] = interpret s
      exec (s : rest) = do
        res <- interpret s
        if res /= None
          then return res
          else exec rest
  interpret (ExprStmt expr) = do
    _ <- eval expr
    return None
  interpret (PrintStmt expr) = do
    obj <- eval expr
    liftIO $ print obj
    return None
  interpret (VariableStmt (Variable ident value)) = do
    obj <- maybe (return Nil) eval value
    ctx <- get
    ctx' <- liftIO $ declare ident obj ctx
    put ctx'
    return None
  interpret (IfStmt (If condition ifBranch elseBranch)) = do
    obj <- eval condition
    if truthy obj
      then interpret ifBranch
      else maybe (return None) interpret elseBranch
  interpret (WhileStmt (While condition loop)) = do
    cond <- eval condition
    if truthy cond
      then do
        res <- interpret loop
        if res /= None
          then return res
          else interpret (WhileStmt (While condition loop))
      else return None
  interpret (ForStmt (For forInitialiser forCondition forIncrement forLoop)) =
    let condition' = fromMaybe (LiteralExpr (BooleanObject True)) forCondition
        body' = case forIncrement of
          Just inc -> BlockStmt [forLoop, ExprStmt inc]
          Nothing -> forLoop
        desugared = WhileStmt (While condition' body')
     in case forInitialiser of
          Just init' -> interpret (BlockStmt [init', desugared])
          Nothing -> interpret desugared
  interpret (ReturnStmt r) = maybe (return (Return Nil)) (fmap Return . eval) r
  interpret (FunStmt fun@Fun {Stmt.name}) = do
    ctx <- get
    function' <- liftIO $ function fun ctx
    ctx' <- liftIO $ declare name (FunctionObject function') ctx
    put ctx'
    return None
  interpret (ClassStmt (Stmt.Class cls parent definitions)) = do
    ctx <- get
    methods <-
      liftIO $
        mapM
          ( \fun@Fun {Stmt.name} -> do
              fn <- function fun ctx
              return (name, fn)
          )
          definitions
  
    let mapping = M.fromList methods
    let (initialiser, filteredMapping) = case M.lookup "init" mapping of
          Just i -> (Just i, M.delete "init" mapping)
          Nothing -> (Nothing, mapping)
  
    superClass <- maybe (return Nothing) evalSuper parent
  
    ctx' <- liftIO $ declare cls (ClassObject (Expr.Class cls initialiser filteredMapping superClass)) ctx
    put ctx'
    return None
    where
      evalSuper p = do
        p' <- eval (VarExpr p)
        case p' of
          ClassObject c -> return $ Just c
          _ -> throwError $ TypeError "Superclass must be a class"
  
  function :: Fun -> Context Object -> IO Function
  function (Fun _ parameters block) ctx = do
    closureRef <- liftIO $ newIORef ctx
    return
      Function
        { arity = length parameters,
          closureRef = closureRef,
          fun = \args -> do
            closure <- liftIO $ readIORef closureRef
            injected <-
              liftIO $
                foldM
                  (\acc (param, arg) -> declare param arg acc)
                  (scope closure)
                  (zip parameters args)
  
            curr <- get
            put injected
            res <- interpret block
            modified <- get
            liftIO $ writeIORef closureRef (descope modified)
            put curr
            case res of
              Return obj -> return obj
              None -> return Nil
        }
  
  bind :: Function -> String -> Object -> IO Function
  bind fun@Function {closureRef} name obj = do
    closure <- liftIO $ readIORef closureRef
    injected <- declare name obj closure
    liftIO $ writeIORef closureRef injected
    return fun
  