#include "expression_build.h"
#include "decompiler/Function/Function.h"
#include "decompiler/IR2/Form.h"
#include "decompiler/IR2/FormStack.h"
#include "decompiler/util/DecompilerTypeSystem.h"
#include "common/goos/PrettyPrinter.h"

namespace decompiler {

/*!
 * The main expression building pass.
 */
bool convert_to_expressions(
    Form* top_level_form,
    FormPool& pool,
    Function& f,
    const std::vector<std::string>& arg_names,
    const std::unordered_map<std::string, LocalVarOverride>& var_override_map,
    const DecompilerTypeSystem& dts) {
  assert(top_level_form);

  try {
    // create the root expression stack for the function
    FormStack stack(true);
    // and add all entries
    for (auto& entry : top_level_form->elts()) {
      entry->push_to_stack(f.ir2.env, pool, stack);
    }

    // rewrite the stack to get the correct final value
    std::vector<FormElement*> new_entries;
    if (f.type.last_arg() != TypeSpec("none")) {
      auto return_var = f.ir2.atomic_ops->end_op().return_var();
      new_entries = rewrite_to_get_var(stack, pool, return_var, f.ir2.env);
      auto reg_return_type =
          f.ir2.env.get_types_after_op(f.ir2.atomic_ops->ops.size() - 1).get(return_var.reg());
      if (!dts.ts.tc(f.type.last_arg(), reg_return_type.typespec())) {
        // we need to cast the final value.
        auto to_cast = new_entries.back();
        new_entries.pop_back();
        auto cast = pool.alloc_element<CastElement>(f.type.last_arg(),
                                                    pool.alloc_single_form(nullptr, to_cast));
        new_entries.push_back(cast);
      }
    } else {
      // or just get all the expressions
      new_entries = stack.rewrite(pool, f.ir2.env);
    }

    // if we are a totally empty function, insert a placeholder so we don't have to handle
    // the zero element case ever.
    if (new_entries.empty()) {
      new_entries.push_back(pool.alloc_element<EmptyElement>());
    }

    // turn us back into a form.
    top_level_form->clear();
    for (auto x : new_entries) {
      top_level_form->push_back(x);
    }

    // and sanity check for tree errors.
    for (auto x : top_level_form->elts()) {
      assert(x->parent_form == top_level_form);
    }

  } catch (std::exception& e) {
    f.warnings.expression_build_warning("In {}: {}", f.guessed_name.to_string(), e.what());
    lg::warn("In {}: {}", f.guessed_name.to_string(), e.what());
    return false;
  }

  // set argument names to some reasonable defaults. these will be used if the user doesn't
  // give us anything more specific.
  if (f.guessed_name.kind == FunctionName::FunctionKind::GLOBAL) {
    f.ir2.env.set_remap_for_function(f.type.arg_count() - 1);
  } else if (f.guessed_name.kind == FunctionName::FunctionKind::METHOD) {
    if (f.guessed_name.method_id == GOAL_NEW_METHOD) {
      f.ir2.env.set_remap_for_new_method(f.type.arg_count() - 1);
    } else {
      f.ir2.env.set_remap_for_method(f.type.arg_count() - 1);
    }
  }

  // get variable names from the user.
  f.ir2.env.map_args_from_config(arg_names, var_override_map);

  // override variable types from the user.

  std::unordered_map<std::string, TypeSpec> retype;
  for (auto& remap : var_override_map) {
    if (remap.second.type) {
      retype[remap.first] = dts.parse_type_spec(*remap.second.type);
    }
  }
  f.ir2.env.set_retype_map(retype);

  return true;
}
}  // namespace decompiler