#include "compiler.hpp"

#include "common.hpp"
#include "stdlib.hpp"

struct Context {
  Context() {}

  Context(Value&& fname, Array&& constants, Dict&& constants_dict,
          Dict&& variables_dict, Array&& closures, Dict&& closures_dict,
          Context* parent, bool toplevel)
      : fname(std::move(fname)),
        constants(std::move(constants)),
        constants_dict(std::move(constants_dict)),
        variables_dict(std::move(variables_dict)),
        closures(std::move(closures)),
        closures_dict(std::move(closures_dict)),
        maxreg(0),
        parent(parent),
        toplevel(toplevel) {}

  static Result<Context> create() {
    auto fname = TRY(Nil::create());
    auto constants = TRY(Array::create());
    auto constants_dict = TRY(Dict::create());
    auto variables_dict = TRY(Dict::create());
    auto closures = TRY(Array::create());
    auto closures_dict = TRY(Dict::create());

    return Context(std::move(fname), std::move(constants),
                   std::move(constants_dict), std::move(variables_dict),
                   std::move(closures), std::move(closures_dict), 0, true);
  }

  static Result<Context> create(Context& parent) {
    auto fname = TRY(Nil::create());
    auto constants = TRY(Array::create());
    auto constants_dict = TRY(Dict::create());
    auto variables_dict = TRY(Dict::create());
    auto closures = TRY(Array::create());
    auto closures_dict = TRY(Dict::create());

    return Context(std::move(fname), std::move(constants),
                   std::move(constants_dict), std::move(variables_dict),
                   std::move(closures), std::move(closures_dict), &parent,
                   false);
  }

  uint64_t alloc_reg() {
    uint64_t reg = maxreg;
    ++maxreg;
    return reg;
  }

  Result<int64_t> add_const(const Value& val) {
    auto idx = constants_dict.get(val);

    if (idx.has_value()) {
      if (!idx.value().is<Int64>()) return ERROR(TypeMismatch);

      return idx.value().to<Int64>()->value();
    }

    int64_t i = TRY(constants.size());

    constants = TRY(constants.append(val));
    constants_dict = TRY(constants_dict.set(val, TRY(Value::create(i))));
    return i;
  }

  template <class T>
  Result<int64_t> add_const(const T& val)
    requires std::derived_from<T, Object>
  {
    return add_const(TRY(val.copy_value()));
  }

  Result<int64_t> add_var(const Value& sym) {
    auto idx = variables_dict.get(sym);
    if (idx.has_value()) {
      if (!idx.value().is<Int64>()) return ERROR(TypeMismatch);

      return idx.value().to<Int64>()->value();
    }

    int64_t i = maxreg;

    variables_dict = TRY(variables_dict.set(sym, TRY(Value::create(i))));

    maxreg++;
    return i;
  }

  Result<int64_t> update_var(const Value& sym) {
    int64_t i = maxreg;

    variables_dict = TRY(variables_dict.set(sym, TRY(Value::create(i))));

    maxreg++;
    return i;
  }

  Result<int64_t> get_var(const Value& sym) {
    auto idx = variables_dict.get(sym);
    if (idx.has_value()) {
      if (!idx.value().is<Int64>()) return ERROR(TypeMismatch);

      return idx.value().to<Int64>()->value();
    }

    return ERROR(KeyError);
  }

  Result<int64_t> get_closure(const Value& sym) {
    auto idx = closures_dict.get(sym);
    if (idx.has_value()) {
      if (!idx.value().is<Int64>()) return ERROR(TypeMismatch);

      return idx.value().to<Int64>()->value();
    }

    // We need to try and look up a symbol in the parent scope. If there's no
    // parent scope, we will still create a closure, but that closure will be
    // resolved dynamically at runtime.
    if (parent) {
      auto maybe_var = parent->get_var(sym);
      // If parent has no variable with name <sym>, then we need to propagate
      // the closure to it. The variable may be defined somewhere in outer
      // scope.
      if (!maybe_var.has_value()) {
        TRY(parent->get_closure(sym));
      }
    }

    int64_t i = TRY(closures.size());
    closures = TRY(closures.append(sym));
    closures_dict = TRY(closures_dict.set(sym, TRY(Value::create(i))));

    return i;
  }

  Value fname;
  Array constants;
  Dict constants_dict;
  Dict variables_dict;
  Array closures;
  Dict closures_dict;
  uint64_t maxreg;
  Context* parent;
  bool toplevel;
};

Result<void> Expression::add_opcode(Oc opcode, OpArg arg1, OpArg arg2,
                                    OpArg arg3, OpArg arg4) {
  Value oc = Value(TRY(Opcode::create(opcode, arg1, arg2, arg3, arg4)));
  code = TRY(code.append(oc));
  return Result<void>();
}

Result<bool> is_primitive_op(Symbol& sym) {
  return TRY(sym.cmp("+")) == 0 || TRY(sym.cmp("-")) == 0 ||
         TRY(sym.cmp("*")) == 0 || TRY(sym.cmp("/")) == 0;
  return false;
}

Result<bool> is_comparison_op(Symbol& sym) {
  return TRY(sym.cmp("<")) == 0 || TRY(sym.cmp("<=")) == 0 ||
         TRY(sym.cmp(">")) == 0 || TRY(sym.cmp(">=")) == 0 ||
         TRY(sym.cmp("=")) == 0 || TRY(sym.cmp("!=")) == 0;
  return false;
}

Result<Value> Compiler::compile(const Value& expr) {
  auto context = TRY(Context::create());

  // If expression is empty - just return nil
  if (expr.is<Nil>()) {
    auto ex = TRY(Expression::create());
    uint64_t reg = context.alloc_reg();
    int64_t c = TRY(context.add_const(TRY(Nil::create())));
    TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)reg}, {1, c}));
    TRY(ex.add_opcode(Oc::Ret, {0, (int64_t)reg}));
    ex.reg = reg;

    Value name = TRY(Nil::create());
    auto fun = TRY(Function::create(name, 0, context.constants, ex.code,
                                    TRY(Array::create())));

    auto mod = TRY(Module::create(name, fun));

    return Value(std::move(mod));
  }

  // Otherwise perform actual compilation of the expression
  if (!TRY(syntax_is_list(expr))) {
    return ERROR(CompilationError);
  }

  auto ex = TRY(compile_body(context, expr));
  TRY(ex.add_opcode(Oc::Ret, {0, (int64_t)ex.reg}));

  Value name = TRY(Nil::create());
  // TRY(debug_print(context.constants));
  // TRY(debug_print(ex.code));

  auto fun = TRY(Function::create(name, 0, context.constants, ex.code,
                                  TRY(Array::create())));

  auto mod = TRY(Module::create(name, fun));

  return Value(std::move(mod));
}

Result<Expression> Compiler::compile_expr(Context& context, const Value& expr) {
  auto unwrapped = TRY(syntax_unwrap(expr));
  switch (unwrapped.tag()) {
    case Tag::Pair:
      return compile_list(context, expr);
    case Tag::Int64:
      return compile_constant(context, expr);
    case Tag::Bool:
      return compile_constant(context, expr);
    case Tag::Symbol:
      return compile_symbol(context, expr);
    case Tag::String:
      return compile_constant(context, expr);
    case Tag::Nil:
      return compile_constant(context, expr);
    case Tag::Float:
      return compile_constant(context, expr);
    case Tag::SrcLoc:
    case Tag::Syntax:
    case Tag::Array:
    case Tag::ByteArray:
    case Tag::Dict:
    case Tag::Opcode:
    case Tag::Function:
    case Tag::StdlibFunction:
    case Tag::Module:
    case Tag::StackFrame:
    case Tag::Error:
    case Tag::Continuation:
    case Tag::Task:
    case Tag::TaskResult:
      return ERROR(TypeMismatch);
  }
  return ERROR(TypeMismatch);
}

Result<Expression> Compiler::compile_primop(Context& context, Symbol& op,
                                            const Value& expr) {
  Value cur = TRY(expr.rest());
  Expression ex = TRY(Expression::create());

  Oc opcode = Oc::Unknown;

  if (TRY(op.cmp("+")) == 0) {
    opcode = Oc::Add;
  } else if (TRY(op.cmp("*")) == 0) {
    opcode = Oc::Mul;
  } else if (TRY(op.cmp("-")) == 0) {
    opcode = Oc::Sub;
  } else if (TRY(op.cmp("/")) == 0) {
    opcode = Oc::Div;
  } else {
    return ERROR(NotImplemented);
  }

  if (cur.is<Nil>()) {
    uint64_t reg = context.alloc_reg();
    int64_t zero = TRY(context.add_const(TRY(Value::create((int64_t)0))));
    TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)reg}, {1, zero}));

    return ex;
  }

  auto subexpr = TRY(cur.first());

  auto comp = TRY(compile_expr(context, subexpr));

  ex.add_code(comp.code);
  uint64_t firstreg = comp.reg;
  uint64_t reg = firstreg;

  cur = TRY(cur.rest());

  while (!cur.is<Nil>()) {
    auto subexpr = TRY(cur.first());

    auto comp = TRY(compile_expr(context, subexpr));

    ex.add_code(comp.code);

    auto rest = TRY(cur.rest());

    uint64_t res = 0;
    if (rest.is<Nil>())
      res = firstreg;
    else
      res = context.alloc_reg();

    TRY(ex.add_opcode(opcode, {0, (int64_t)res}, {0, (int64_t)reg},
                      {0, (int64_t)comp.reg}));

    reg = res;
    cur = std::move(rest);
  }

  context.maxreg = firstreg + 1;
  ex.reg = reg;
  return std::move(ex);
}

Result<Expression> Compiler::compile_comparison(Context& context, Symbol& op,
                                                const Value& expr) {
  Value cur = TRY(expr.rest());
  Expression ex = TRY(Expression::create());

  Oc opcode = Oc::Unknown;
  int64_t cmp_expected = 0;

  if (TRY(op.cmp("<")) == 0) {
    opcode = Oc::Less;
    cmp_expected = 0;
  } else if (TRY(op.cmp("<=")) == 0) {
    opcode = Oc::LessEqual;
    cmp_expected = 0;
  } else if (TRY(op.cmp(">")) == 0) {
    opcode = Oc::LessEqual;
    cmp_expected = 1;
  } else if (TRY(op.cmp(">=")) == 0) {
    opcode = Oc::Less;
    cmp_expected = 1;
  } else if (TRY(op.cmp("=")) == 0) {
    opcode = Oc::Equal;
    cmp_expected = 0;
  } else if (TRY(op.cmp("!=")) == 0) {
    opcode = Oc::Equal;
    cmp_expected = 1;
  } else {
    return ERROR(NotImplemented);
  }

  if (cur.is<Nil>()) {
    return syntax_error(expr, "Comparison must have at least one argument");
  }

  uint64_t result = context.alloc_reg();
  int64_t true_c = TRY(context.add_const(TRY(Bool::create(true))));
  int64_t false_c = TRY(context.add_const(TRY(Bool::create(false))));
  TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)result}, {1, true_c}));

  auto subexpr = TRY(cur.first());

  auto comp = TRY(compile_expr(context, subexpr));

  ex.add_code(comp.code);
  uint64_t firstreg = comp.reg;
  uint64_t reg = firstreg;

  cur = TRY(cur.rest());

  while (!cur.is<Nil>()) {
    auto subexpr = TRY(cur.first());

    auto comp = TRY(compile_expr(context, subexpr));

    ex.add_code(comp.code);

    auto rest = TRY(cur.rest());

    TRY(ex.add_opcode(opcode, {0, (int64_t)reg}, {0, (int64_t)comp.reg},
                      {0, (int64_t)cmp_expected}));

    TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)result}, {1, false_c}));

    reg = comp.reg;
    cur = std::move(rest);
  }

  context.maxreg = result + 1;
  ex.reg = result;
  return std::move(ex);
}

Result<Expression> Compiler::compile_if(Context& context, Symbol& op,
                                        const Value& expr) {
  Value rest = TRY(expr.rest());
  Expression ex = TRY(Expression::create());

  auto num_params = TRY(rest.size());
  if (num_params != 3) {
    return syntax_error(TRY(expr.first()),
                        "\"if\" form must have exactly 3 arguments");
  }

  auto condition = TRY(rest.first());
  auto option1 = TRY(rest.second());
  auto option2 = TRY(rest.third());

  auto condition_comp = TRY(compile_expr(context, condition));

  ex.add_code(condition_comp.code);
  uint64_t firstreg = condition_comp.reg;
  uint64_t reg = firstreg;

  auto option1_comp = TRY(compile_expr(context, option1));

  uint64_t option1_reg = option1_comp.reg;

  context.maxreg = firstreg + 1;

  auto option2_comp = TRY(compile_expr(context, option2));

  int64_t true_const = TRY(context.add_const(TRY(Bool::create(true))));

  TRY(ex.add_opcode(Oc::Equal, {0, (int64_t)firstreg}, {1, (int64_t)true_const},
                    {0, (int64_t)0}));

  TRY(ex.add_opcode(Oc::Jump, {0, (int64_t)TRY(option1_comp.code.size()) + 2}));

  ex.add_code(option1_comp.code);
  TRY(ex.add_opcode(Oc::Jump, {0, (int64_t)TRY(option2_comp.code.size()) + 1}));
  ex.add_code(option2_comp.code);

  uint64_t option2_reg = option2_comp.reg;

  TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)firstreg},
                    {0, (int64_t)option1_reg}));

  context.maxreg = firstreg + 1;
  ex.reg = firstreg;
  return std::move(ex);
}

Result<Expression> Compiler::compile_when(Context& context, Symbol& op,
                                          const Value& expr) {
  Value rest = TRY(expr.rest());
  Expression ex = TRY(Expression::create());

  auto num_params = TRY(rest.size());
  if (num_params < 2) {
    return syntax_error(TRY(expr.first()),
                        "\"when\" form must have at least 2 arguments");
  }

  auto condition = TRY(rest.first());
  auto body = TRY(rest.rest());

  auto condition_comp = TRY(compile_expr(context, condition));

  ex.add_code(condition_comp.code);
  uint64_t firstreg = condition_comp.reg;
  uint64_t reg = firstreg;

  auto body_comp = TRY(compile_body(context, body));
  uint64_t body_reg = body_comp.reg;

  context.maxreg = firstreg + 1;

  int64_t true_const = TRY(context.add_const(TRY(Bool::create(true))));
  int64_t nil_const = TRY(context.add_const(TRY(Nil::create())));

  TRY(ex.add_opcode(Oc::Equal, {0, (int64_t)firstreg}, {1, (int64_t)true_const},
                    {0, (int64_t)1}));

  TRY(ex.add_opcode(Oc::Jump, {0, (int64_t)3}));
  TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)firstreg}, {1, (int64_t)nil_const}));
  TRY(ex.add_opcode(Oc::Jump, {0, (int64_t)TRY(body_comp.code.size()) + 2}));

  ex.add_code(body_comp.code);

  TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)firstreg}, {0, (int64_t)body_reg}));

  context.maxreg = firstreg + 1;
  ex.reg = firstreg;
  return std::move(ex);
}
Result<Expression> Compiler::compile_and(Context& context, Symbol& op,
                                         const Value& expr) {
  Value param = TRY(expr.rest());
  Expression ex = TRY(Expression::create());

  if (param.is<Nil>()) {
    return syntax_error(TRY(expr.first()),
                        "\"and\" form must have at least one parameter");
  }

  uint64_t prev_reg = 0;
  uint64_t result = context.alloc_reg();
  int64_t true_c = TRY(context.add_const(TRY(Bool::create(true))));
  int64_t false_c = TRY(context.add_const(TRY(Bool::create(false))));
  TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)result}, {1, false_c}));

  bool is_first = true;

  while (!param.is<Nil>()) {
    auto rest = TRY(param.rest());
    Value param_val = TRY(param.first());

    auto param_ex = TRY(compile_expr(context, param_val));

    if (!is_first) {
      TRY(ex.add_opcode(Oc::Equal, {0, (int64_t)prev_reg}, {1, (int64_t)true_c},
                        {0, (int64_t)0}));
      TRY(ex.add_opcode(Oc::Jump, {0, (int64_t)TRY(param_ex.code.size()) + 2}));
    }
    prev_reg = param_ex.reg;
    TRY(ex.add_code(param_ex.code));

    param = std::move(rest);
    is_first = false;
  }
  TRY(ex.add_opcode(Oc::Equal, {0, (int64_t)prev_reg}, {1, (int64_t)true_c},
                    {0, (int64_t)0}));
  TRY(ex.add_opcode(Oc::Jump, {0, (int64_t)2}));
  TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)result}, {1, true_c}));

  context.maxreg = result + 1;
  ex.reg = result;
  return std::move(ex);
}

Result<Expression> Compiler::compile_or(Context& context, Symbol& op,
                                        const Value& expr) {
  Value param = TRY(expr.rest());
  Expression ex = TRY(Expression::create());
  if (param.is<Nil>()) {
    return syntax_error(TRY(expr.first()),
                        "\"or\" form must have at least one parameter");
  }

  uint64_t prev_reg = 0;
  uint64_t result = context.alloc_reg();
  int64_t true_c = TRY(context.add_const(TRY(Bool::create(true))));
  int64_t false_c = TRY(context.add_const(TRY(Bool::create(false))));
  TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)result}, {1, true_c}));

  bool is_first = true;

  while (!param.is<Nil>()) {
    auto rest = TRY(param.rest());
    Value param_val = TRY(param.first());

    auto param_ex = TRY(compile_expr(context, param_val));

    if (!is_first) {
      TRY(ex.add_opcode(Oc::Equal, {0, (int64_t)prev_reg}, {1, (int64_t)true_c},
                        {0, (int64_t)1}));
      TRY(ex.add_opcode(Oc::Jump, {0, (int64_t)TRY(param_ex.code.size()) + 2}));
    }
    prev_reg = param_ex.reg;
    TRY(ex.add_code(param_ex.code));

    param = std::move(rest);
    is_first = false;
  }
  TRY(ex.add_opcode(Oc::Equal, {0, (int64_t)prev_reg}, {1, (int64_t)true_c},
                    {0, (int64_t)1}));
  TRY(ex.add_opcode(Oc::Jump, {0, (int64_t)2}));
  TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)result}, {1, false_c}));

  context.maxreg = result + 1;
  ex.reg = result;
  return std::move(ex);
}

Result<Expression> Compiler::compile_not(Context& context, Symbol& op,
                                         const Value& expr) {
  Value first = TRY(expr.rest());
  Expression ex = TRY(Expression::create());
  if (TRY(first.size()) != 1) {
    return syntax_error(TRY(expr.first()),
                        "\"not\" form must have exactly one parameter");
  }

  auto first_expr = TRY(first.first());

  uint64_t result = context.alloc_reg();
  int64_t true_c = TRY(context.add_const(TRY(Bool::create(true))));
  int64_t false_c = TRY(context.add_const(TRY(Bool::create(false))));
  TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)result}, {1, true_c}));

  auto comp = TRY(compile_expr(context, first_expr));
  TRY(ex.add_code(comp.code));

  TRY(ex.add_opcode(Oc::Equal, {0, (int64_t)comp.reg}, {1, (int64_t)true_c},
                    {0, (int64_t)1}));
  TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)result}, {1, false_c}));

  context.maxreg = result + 1;
  ex.reg = result;
  return std::move(ex);
}

Result<Expression> Compiler::compile_def(Context& context, Symbol& op,
                                         const Value& expr) {
  uint64_t expr_size = TRY(expr.size());
  if (expr_size < 2) {
    return syntax_error(TRY(expr.first()),
                        "\"def\" form must have at least one parameter");
  }

  if (expr_size > 3) {
    return syntax_error(TRY(expr.first()),
                        "Too many parameters for \"def\" form");
  }

  auto varname = TRY(expr.second());
  auto varname_unwrapped = TRY(syntax_unwrap(varname));
  if (!varname_unwrapped.is<Symbol>()) {
    return syntax_error(varname,
                        "Variable name in \"def\" form must be a symbol");
  }

  Expression ex = TRY(Expression::create());

  int64_t gname = TRY(context.add_const(varname_unwrapped));
  uint64_t maxreg = context.maxreg;

  // No value for the variable is specified - just load "nil" into it
  if (expr_size == 2) {
    int64_t c = TRY(context.add_const(Value(TRY(Nil::create()))));
    TRY(ex.add_opcode(Oc::GlobalStore, {1, (int64_t)gname}, {1, (int64_t)c}));
    TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)maxreg}, {1, (int64_t)c}));
    ex.reg = maxreg;
    context.maxreg = maxreg + 1;
    return std::move(ex);
  }

  auto var_expr = TRY(expr.third());

  auto comp = TRY(compile_expr(context, var_expr));
  ex.add_code(comp.code);

  TRY(ex.add_opcode(Oc::GlobalStore, {1, (int64_t)gname},
                    {0, (int64_t)comp.reg}));

  ex.reg = comp.reg;
  context.maxreg = maxreg + 1;

  return std::move(ex);
}

Result<Expression> Compiler::compile_quote(Context& context, Symbol& op,
                                           const Value& expr) {
  uint64_t expr_size = TRY(expr.size());
  if (expr_size != 2) {
    return syntax_error(TRY(expr.first()),
                        "\"quote\" form must have exactly one parameter");
  }

  auto quoted = TRY(expr.second());
  quoted = TRY(syntax_unwrap_all(quoted));

  uint64_t maxreg = context.maxreg;
  Expression ex = TRY(Expression::create());

  int64_t quoted_const = TRY(context.add_const(quoted));
  TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)maxreg}, {1, (int64_t)quoted_const}));

  ex.reg = maxreg;
  context.maxreg = maxreg + 1;

  return std::move(ex);
}

Result<Expression> Compiler::compile_syntax(Context& context, Symbol& op,
                                            const Value& expr) {
  uint64_t expr_size = TRY(expr.size());
  if (expr_size != 2) {
    return syntax_error(TRY(expr.first()),
                        "\"syntax\" form must have exactly one parameter");
  }

  auto quoted = TRY(expr.second());

  uint64_t maxreg = context.maxreg;
  Expression ex = TRY(Expression::create());

  int64_t quoted_const = TRY(context.add_const(quoted));
  TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)maxreg}, {1, (int64_t)quoted_const}));

  ex.reg = maxreg;
  context.maxreg = maxreg + 1;

  return std::move(ex);
}

Result<Expression> Compiler::compile_fn(Context& context, Symbol& op,
                                        const Value& expr) {
  Context ctx = TRY(Context::create(context));

  auto rest = TRY(expr.rest());

  Value name = TRY(Nil::create());

  auto maybe_name = TRY(expr.second());
  auto maybe_name_unwrapped = TRY(syntax_unwrap(maybe_name));

  if (maybe_name_unwrapped.is<Symbol>()) {
    name = TRY(maybe_name_unwrapped.copy());
    ctx.fname = TRY(name.copy());

    rest = TRY(rest.rest());
  }

  auto param = TRY(rest.first());

  if (!TRY(syntax_is_list(param))) {
    return syntax_error(param, "Function parameters must be a list");
  }

  uint64_t arity = 0;
  while (!TRY(syntax_is_nil(param))) {
    auto param_name = TRY(param.first());
    auto param_name_unwrapped = TRY(syntax_unwrap(param_name));

    if (!param_name_unwrapped.is<Symbol>()) {
      return syntax_error(param_name, "Parameter name must be a symbol");
    }

    int64_t reg = TRY(ctx.add_var(param_name_unwrapped));

    param = TRY(param.rest());
    arity++;
  }

  Value body = TRY(rest.rest());

  auto ex = TRY(compile_body(ctx, body));

  TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)0}, {0, (int64_t)ex.reg}));
  TRY(ex.add_opcode(Oc::Ret, {0, (int64_t)0}));

  auto fun = TRY(Function::create(name, arity, ctx.constants, ex.code,
                                  TRY(Array::create())));

  // std::cout << "--------------- LAMBDA " << arity << "\n";
  // TRY(debug_print(expr));
  // TRY(debug_print(ctx.constants));
  // TRY(debug_print(ex.code));

  Expression ex_res = TRY(Expression::create());

  if (TRY(ctx.closures.size()) == 0) {
    int64_t c = TRY(context.add_const(TRY(fun.copy())));

    if (context.toplevel && !name.is<Nil>()) {
      int64_t gname = TRY(context.add_const(name));
      TRY(ex_res.add_opcode(Oc::GlobalStore, {1, (int64_t)gname},
                            {1, (int64_t)c}));
    }

    uint64_t reg = context.alloc_reg();
    TRY(ex_res.add_opcode(Oc::Mov, {0, (int64_t)reg}, {1, (int64_t)c}));

    context.maxreg = reg + 1;
    ex_res.reg = reg;

    return ex_res;
  }

  int64_t c = TRY(context.add_const(TRY(fun.copy())));
  uint64_t reg = context.alloc_reg();
  TRY(ex_res.add_opcode(Oc::Mov, {0, (int64_t)reg}, {1, (int64_t)c}));

  auto closures_size = TRY(ctx.closures.size());
  for (uint64_t i = 0; i < closures_size; i++) {
    Value sym = TRY(ctx.closures.get(i));

    // If this is a regular variable from outer scope - it is present in a
    // register, and as such we can just move it to the correct place
    auto maybe_var_reg = context.get_var(sym);
    if (maybe_var_reg.has_value()) {
      int64_t var_reg = maybe_var_reg.value();
      uint64_t vr = context.alloc_reg();
      TRY(ex_res.add_opcode(Oc::Mov, {0, (int64_t)vr}, {0, (int64_t)var_reg}));
      continue;
    }

    // Otherwise this can be a variable we get from the intermediate closure
    // (such as in case of 3 nested functions where a var we close over is
    // defined in the outermost function.
    auto maybe_closure = context.get_closure(sym);
    if (maybe_closure.has_value() && !context.toplevel) {
      int64_t closure_id = maybe_closure.value();
      uint64_t vr = context.alloc_reg();
      TRY(ex_res.add_opcode(Oc::ClosureLoad, {0, (int64_t)vr},
                            {0, (int64_t)closure_id}));
      continue;
    }

    int64_t c = TRY(context.add_const(sym));

    uint64_t vr = context.alloc_reg();
    TRY(ex_res.add_opcode(Oc::GlobalLoad, {0, (int64_t)vr}, {1, (int64_t)c}));
  }

  TRY(ex_res.add_opcode(
      Oc::MakeClosure, {0, (int64_t)reg},
      {0, (int64_t)reg + (int64_t)TRY(ctx.closures.size()) + 1}));

  if (context.toplevel && !name.is<Nil>()) {
    int64_t gname = TRY(context.add_const(name));
    TRY(ex_res.add_opcode(Oc::GlobalStore, {1, (int64_t)gname},
                          {0, (int64_t)reg}));
  }

  context.maxreg = reg + 1;
  ex_res.reg = reg;

  return ex_res;
}

Result<Expression> Compiler::compile_let(Context& context, Symbol& op,
                                         const Value& expr) {
  // Save the variable bindings to restore it later
  Dict saved_vars = TRY(context.variables_dict.copy());

  auto rest = TRY(expr.rest());

  if (TRY(expr.size()) < 2) {
    return syntax_error(TRY(expr.first()),
                        "\"let\" form must have at least one parameter");
  }

  uint64_t maxreg = context.maxreg;

  auto bindings = TRY(rest.first());
  bindings = TRY(syntax_unwrap(bindings));

  Expression ex_res = TRY(Expression::create());

  while (!bindings.is<Nil>()) {
    auto binding = TRY(bindings.first());

    auto binding_name = TRY(binding.first());
    auto binding_name_unwrapped = TRY(syntax_unwrap(binding_name));
    auto binding_expr = TRY(binding.second());

    if (!binding_name_unwrapped.is<Symbol>()) {
      return syntax_error(binding_name, "Binding name must be a symbol");
    }

    int64_t reg = TRY(context.update_var(binding_name_unwrapped));
    auto ex = TRY(compile_expr(context, binding_expr));
    TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)reg}, {0, (int64_t)ex.reg}));
    context.maxreg = reg + 1;

    ex_res.add_code(ex.code);

    bindings = TRY(bindings.rest());
  }

  Value body = TRY(rest.rest());

  auto ex = TRY(compile_body(context, body));

  TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)maxreg}, {0, (int64_t)ex.reg}));
  ex_res.add_code(ex.code);

  // Restore the variables back
  context.variables_dict = std::move(saved_vars);
  ex_res.reg = maxreg;

  return std::move(ex_res);
}

Result<Expression> Compiler::compile_body(Context& context, const Value& expr) {
  auto cur = TRY(expr.copy());

  Expression ex_res = TRY(Expression::create());

  int64_t maxreg = context.maxreg;

  // Body is empty, in which case just replace it with "nil" value
  if (TRY(expr.size()) == 0) {
    int64_t c = TRY(context.add_const(Value(TRY(Nil::create()))));
    TRY(ex_res.add_opcode(Oc::Mov, {0, (int64_t)maxreg}, {1, (int64_t)c}));
    ex_res.reg = maxreg;
    context.maxreg = maxreg + 1;
    return std::move(ex_res);
  }

  cur = TRY(syntax_unwrap(cur));
  while (!cur.is<Nil>()) {
    auto expr_val = TRY(cur.first());

    // debug_print(expr_val);

    auto expr = TRY(compile_expr(context, expr_val));

    TRY(ex_res.add_code(expr.code));

    cur = TRY(cur.rest());

    if (cur.is<Nil>()) {
      ex_res.reg = expr.reg;
    } else {
      context.maxreg = maxreg;
    }
  }

  return ex_res;
}

Result<Expression> Compiler::compile_function_call(Context& context,
                                                   const Value& expr) {
  auto ex = TRY(Expression::create());

  auto first = TRY(expr.first());
  auto first_unwrapped = TRY(syntax_unwrap(first));
  auto param = TRY(expr.rest());

  bool is_self_call = false;
  uint64_t firstreg = context.maxreg;
  if (first_unwrapped.is<Symbol>() &&
      TRY(first_unwrapped.cmp(context.fname)) == 0) {
    int64_t r = context.alloc_reg();
    int64_t c = TRY(context.add_const(TRY(Nil::create())));
    TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)r}, {1, (int64_t)c}));
    is_self_call = true;
  } else {
    auto fun_ex = TRY(compile_expr(context, first));
    TRY(ex.add_code(fun_ex.code));
  }

  int64_t maxreg = context.maxreg;
  while (!param.is<Nil>()) {
    Value param_val = TRY(param.first());

    auto param_ex = TRY(compile_expr(context, param_val));
    TRY(ex.add_code(param_ex.code));

    param = TRY(param.rest());
  }

  if (is_self_call) {
    TRY(ex.add_opcode(Oc::SelfCall, {0, (int64_t)firstreg},
                      {0, (int64_t)context.maxreg}));
  } else {
    TRY(ex.add_opcode(Oc::Call, {0, (int64_t)firstreg},
                      {0, (int64_t)context.maxreg}));
  }

  ex.reg = firstreg;
  context.maxreg = maxreg;
  return ex;
}

Result<Expression> Compiler::compile_list(Context& context, const Value& expr) {
  auto first = TRY(expr.first());
  auto unwrapped = TRY(syntax_unwrap(first));

  if (unwrapped.is<Symbol>()) {
    Symbol& sym = *unwrapped.to<Symbol>();
    if (TRY(is_primitive_op(sym))) {
      return compile_primop(context, sym, expr);
    } else if (TRY(is_comparison_op(sym))) {
      return compile_comparison(context, sym, expr);
    } else if (TRY(sym.cmp("if")) == 0) {
      return compile_if(context, sym, expr);
    } else if (TRY(sym.cmp("when")) == 0) {
      return compile_when(context, sym, expr);
    } else if (TRY(sym.cmp("and")) == 0) {
      return compile_and(context, sym, expr);
    } else if (TRY(sym.cmp("or")) == 0) {
      return compile_or(context, sym, expr);
    } else if (TRY(sym.cmp("not")) == 0) {
      return compile_not(context, sym, expr);
    } else if (TRY(sym.cmp("def")) == 0) {
      return compile_def(context, sym, expr);
    } else if (TRY(sym.cmp("fn")) == 0) {
      return compile_fn(context, sym, expr);
    } else if (TRY(sym.cmp("let")) == 0) {
      return compile_let(context, sym, expr);
    } else if (TRY(sym.cmp("quote")) == 0) {
      return compile_quote(context, sym, expr);
    } else if (TRY(sym.cmp("syntax")) == 0) {
      return compile_syntax(context, sym, expr);
    } else {
      return compile_function_call(context, expr);
    }
  } else if (TRY(syntax_is_list(first))) {
    return compile_function_call(context, expr);
  }

  return ERROR(TypeMismatch);
}

Result<Expression> Compiler::compile_constant(Context& context,
                                              const Value& value) {
  Expression ex = TRY(Expression::create());
  uint64_t reg = context.alloc_reg();

  auto unwrapped = TRY(syntax_unwrap(value));
  int64_t c = TRY(context.add_const(unwrapped));

  TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)reg}, {1, (int64_t)c}));

  ex.reg = reg;
  return std::move(ex);
}

Result<Expression> Compiler::compile_symbol(Context& context,
                                            const Value& value) {
  Expression ex = TRY(Expression::create());

  auto unwrapped = TRY(syntax_unwrap(value));

  // Symbol may be self-evaluating (e.g. :foo), in which case we just use it as
  // its own value

  if (unwrapped.is<Symbol>()) {
    Symbol& sym = *unwrapped.to<Symbol>();

    if (TRY(sym.size()) > 0 && TRY(sym[0]) == ':') {
      int64_t c = TRY(context.add_const(unwrapped));

      uint64_t reg = context.alloc_reg();
      TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)reg}, {1, (int64_t)c}));

      ex.reg = reg;
      return std::move(ex);
    }
  }

  auto maybe_reg = context.get_var(TRY(unwrapped.copy()));

  if (!maybe_reg.has_error()) {
    auto var_reg = maybe_reg.value();

    uint64_t reg = context.alloc_reg();
    TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)reg}, {0, (int64_t)var_reg}));

    ex.reg = reg;
    return std::move(ex);
  }

  if (unwrapped.is<Symbol>()) {
    auto maybe_stdlib_fun = get_stdlib_function(*unwrapped.to<Symbol>());

    if (!maybe_stdlib_fun.has_error()) {
      auto stdlib_fun = TRY(StdlibFunction::create(maybe_stdlib_fun.value()));
      int64_t c = TRY(context.add_const(stdlib_fun));

      uint64_t reg = context.alloc_reg();
      TRY(ex.add_opcode(Oc::Mov, {0, (int64_t)reg}, {1, (int64_t)c}));

      ex.reg = reg;
      return std::move(ex);
    }
  }

  // Most variable references in function bodies will fall back to loading the
  // accessed variables from parent closures.
  if (!context.toplevel) {
    auto maybe_closure = context.get_closure(TRY(unwrapped.copy()));

    if (!maybe_closure.has_error()) {
      auto var_closure = maybe_closure.value();

      uint64_t reg = context.alloc_reg();
      TRY(ex.add_opcode(Oc::ClosureLoad, {0, (int64_t)reg},
                        {0, (int64_t)var_closure}));

      ex.reg = reg;
      return std::move(ex);
    }
  }

  // Otherwise treat unknown symbol as a global and try to load it from the
  // global scope. This usually happens in toplevel expressions.

  int64_t c = TRY(context.add_const(TRY(unwrapped.copy())));

  uint64_t reg = context.alloc_reg();
  TRY(ex.add_opcode(Oc::GlobalLoad, {0, (int64_t)reg}, {1, (int64_t)c}));

  ex.reg = reg;
  return std::move(ex);
}

Result<Expression> Compiler::syntax_error(const Value& expr, const char* msg) {
  if (expr.is<Syntax>()) {
    Syntax& stx = *expr.to<Syntax>();
    auto srcloc = TRY(stx.srcloc());
    SrcLoc& loc = *srcloc.to<SrcLoc>();
    auto range = loc.sourcerange();
    auto pos = range.start;

    return ERROR_FMT(SyntaxError, _fname, ":", pos.line, ":", pos.column,
                     " Syntax error: ", msg);
  }
  return ERROR_FMT(SyntaxError, _fname, " Syntax error: ", msg);
}

Result<Compiler> Compiler::create(const String& fname) {
  return Compiler(TRY(fname.copy()));
}

Result<Value> compile(const String& fname, const Value& expr) {
  Compiler c = TRY(Compiler::create(fname));
  return c.compile(expr);
}