#include "compiler.hpp"

#include "common.hpp"

struct Context {
  Context() {}

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

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

    return Context(std::move(env), std::move(constants),
                   std::move(constants_dict), std::move(variables_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;
  }

  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.insert(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);
  }

  Value env;
  Array constants;
  Dict constants_dict;
  Dict variables_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());
  // TRY(debug_print(context.constants));
  // TRY(debug_print(ex.code));

  auto fun = TRY(Function::create(name, 0, context.constants, ex.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::Bool:
      return compile_bool(context, *expr.to<Bool>());
    case Tag::Symbol:
      return compile_symbol(context, *expr.to<Symbol>());
    case Tag::Nil:
    case Tag::Float:
    case Tag::String:
    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 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;
  }

  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_if(Context& context, Symbol& op,
                                        Pair& expr) {
  Value first = TRY(expr.rest());
  Expression ex = TRY(Expression::create());
  if (first.is<Nil>() || !first.is<Pair>()) {
    return ERROR(CompilationError);
  }

  Pair& first_pair = *first.to<Pair>();

  Value second = TRY(first_pair.rest());

  if (second.is<Nil>() || !second.is<Pair>()) {
    return ERROR(CompilationError);
  }

  Pair& second_pair = *second.to<Pair>();

  Value third = TRY(second_pair.rest());

  if (third.is<Nil>() || !third.is<Pair>()) {
    return ERROR(CompilationError);
  }

  Pair& third_pair = *third.to<Pair>();

  auto condition = TRY(first_pair.first());

  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 = TRY(second_pair.first());
  auto option1_comp = TRY(compile_expr(context, option1));

  uint64_t option1_reg = option1_comp.reg;

  context.maxreg = firstreg + 1;

  auto option2 = TRY(third_pair.first());
  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::JumpNotEqual, {0, (int64_t)firstreg},
                    {1, (int64_t)true_const},
                    {0, (int64_t)option1_comp.code.size() + 2}));

  ex.add_code(option1_comp.code);
  TRY(ex.add_opcode(Oc::Jump, {0, (int64_t)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_lambda(Context& context, Symbol& op,
                                            Pair& expr) {
  Context ctx = TRY(Context::create());

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

  if (first.is<Nil>() || !first.is<Pair>()) {
    return ERROR(CompilationError);
  }

  Pair& first_pair = *first.to<Pair>();

  auto param = TRY(first_pair.first());
  if (param.is<Nil>()) {
    return ERROR(CompilationError);
  }

  uint64_t arity = 0;
  while (!param.is<Nil>()) {
    if (!param.is<Pair>()) {
      return ERROR(CompilationError);
    }
    Pair& param_pair = *param.to<Pair>();

    auto param_first = TRY(param_pair.first());

    if (!param_first.is<Symbol>()) {
      return ERROR(CompilationError);
    }

    int64_t reg = TRY(ctx.add_var(param_first));
    std::cout << "reg: " << reg << "\n";

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

  Value second = TRY(first_pair.rest());

  if (second.is<Nil>() || !second.is<Pair>()) {
    return ERROR(CompilationError);
  }

  Pair& second_pair = *second.to<Pair>();

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

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

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

  int64_t c = TRY(context.add_const(TRY(fun.copy_value())));

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

  ex.reg = reg;

  return ex_res;
}

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

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

  int64_t maxreg = context.maxreg;

  while (!cur.is<Nil>()) {
    if (!cur.is<Pair>()) {
      return ERROR(CompilationError);
    }
    Pair& cur_pair = *cur.to<Pair>();
    auto expr_val = TRY(cur_pair.first());

    debug_print(expr_val);

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

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

    cur = TRY(cur_pair.rest());

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

  return ex_res;
}

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);
    } else if (TRY(sym.cmp("if")) == 0) {
      return compile_if(context, sym, expr);
    } else if (TRY(sym.cmp("lambda")) == 0) {
      return compile_lambda(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_value())));

  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, Symbol& value) {
  Expression ex = TRY(Expression::create());

  auto maybe_reg = context.get_var(TRY(value.copy_value()));

  if (maybe_reg.has_error()) {
    return ERROR(CompilationError);
  }

  auto var_reg = maybe_reg.value();

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

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

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

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

  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);
}