#include "compiler.hpp"

struct Context {
  Context() {}
  Context(Value&& env, Array&& constants, Dict&& constants_dict)
      : env(std::move(env)),
        constants(std::move(constants)),
        constants_dict(std::move(constants_dict)),
        maxreg(0) {}

  static Result<Context> create() {
    auto env = TRY(Nil::create());
    auto constants = TRY(Array::create());
    auto constants_dict = TRY(Dict::create());

    return Context(std::move(env), std::move(constants),
                   std::move(constants_dict));
  }

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

  Result<int64_t> add_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 = constants.size();

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

  Value env;
  Array constants;
  Dict constants_dict;
  uint64_t maxreg;
};

Result<Array> add_opcode(Array& opcodes, Oc opcode, OpArg arg1 = {0, 0},
                         OpArg arg2 = {0, 0}, OpArg arg3 = {0, 0},
                         OpArg arg4 = {0, 0}) {
  Value oc = Value(TRY(Opcode::create(opcode, arg1, arg2, arg3, arg4)));
  return opcodes.append(oc);
}

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<Value> Compiler::compile(Value& expr) {
  auto context = TRY(Context::create());

  TRY(compile_expr(context, expr));
  return ERROR(NotImplemented);
}
Result<Expression> Compiler::compile_expr(Context& context, Value& expr) {
  switch (expr.tag()) {
    case Tag::Pair:
      return compile_list(context, *expr.to<Pair>());
    case Tag::Int64:
      return compile_int64(context, *expr.to<Int64>());
    case Tag::Nil:
    case Tag::Bool:
    case Tag::Float:
    case Tag::String:
    case Tag::Symbol:
    case Tag::Syntax:
    case Tag::Array:
    case Tag::ByteArray:
    case Tag::Dict:
    case Tag::Opcode:
      return ERROR(TypeMismatch);
  }
  return ERROR(TypeMismatch);
}

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

  uint64_t reg = context.alloc_reg();
  std::cout << "load-immediate r" << reg << ", " << 0 << "\n";
  code = TRY(add_opcode(code, Oc::LoadImmediate, {0, (int64_t)reg}, {0, 0}));

  while (!cur.is<Nil>()) {
    Pair& pair = *cur.to<Pair>();
    auto subexpr = TRY(pair.first());

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

    uint64_t res = context.alloc_reg();
    std::cout << "add r" << res << ", r" << reg << ", r" << comp.reg << "\n";

    reg = res;
    cur = TRY(pair.rest());
  }

  Expression ex = Expression(reg, TRY(Array::create()));
  return ex;
}

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

  if (first.is<Symbol>()) {
    Symbol& sym = *first.to<Symbol>();
    if (TRY(is_primitive_op(sym))) {
      return compile_primop(context, sym, expr);
    }
  }
  return ERROR(TypeMismatch);
}

Result<Expression> Compiler::compile_int64(Context& context, Int64& value) {
  uint64_t reg = context.alloc_reg();
  std::cout << "loadc r" << reg << ", " << value.value() << "\n";

  Expression expr = Expression(reg, TRY(Array::create()));
  return expr;
}