#include "compiler.hpp"

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

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

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

  auto fun = TRY(Function::create(name, 0, constants, code));
  return Value(std::move(fun));
}

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:
    case Tag::Function:
    case Tag::Stack:
      return ERROR(TypeMismatch);
  }
  return ERROR(TypeMismatch);
}

Result<Expression> Compiler::compile_primop(Context& context, Symbol& op,
                                            Pair& 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 {
    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;
  }

  Pair& pair = *cur.to<Pair>();
  auto subexpr = TRY(pair.first());

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

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

  cur = TRY(pair.rest());

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

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

    ex.add_code(comp.code);

    auto rest = TRY(pair.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_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) {
  Expression ex = TRY(Expression::create());
  uint64_t reg = context.alloc_reg();

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

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

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

Result<Value> compile(Value& expr) {
  Compiler c = Compiler();
  return c.compile(expr);
}