#include "Compiler.hpp"
#include "Lambda.hpp"

namespace fk
{
  /*explicit*/ Compiler::Compiler(std::shared_ptr<SymTable> sym)
    : m_sym { sym }
  {
  }

  /*virtual*/ Compiler::~Compiler()
  {
  }

  void Compiler::add_macro(std::string const& name,
                           std::shared_ptr<NativeMacro> macro)
  {
    m_macros[name] = macro;
  }

  std::shared_ptr<Program> Compiler::compile(std::shared_ptr<Node> node)
  {
    auto prog = std::make_shared<Program>();
    compile_prog(node, prog);
    return prog;
  }

  void Compiler::compile_prog(std::shared_ptr<Node> node,
                              std::shared_ptr<Program> prog)
  {
    switch (node->type())
      {
      case NODE_MODULE: {
        for (size_t i=0; i<node->size(); i++)
          {
            compile_prog(node->child(i), prog);
            prog->add(OP_POP);
          }
      } break;

      case NODE_BODY: {
        for (size_t i=0; i<node->size(); i++)
          {
            compile_prog(node->child(i), prog);

            if (i != node->size() - 1)
              {
                prog->add(OP_POP);
              }
          }
      } break;

      case NODE_LAMBDA: {
        auto params = node->child(0);
        auto body = node->child(1);

        m_sym->enter_scope(node);

        for (size_t i=0; i<params->size(); i++)
          {
            std::string ident = params->child(i)->repr();
            m_sym->declare_local(ident, i, node->loc());
          }

        Compiler compiler {m_sym};
        for (auto e: m_macros)
          {
            compiler.add_macro(e.first, e.second);
          }

        auto program = compiler.compile(body);
        program->add(OP_RET);

        m_sym->leave_scope();

        auto constant = std::make_shared<Constant>(TYPE_PROGRAM,
                                                   program,
                                                   node->loc());
        prog->load_const(constant);
        prog->add(OP_MAKE_FUNCTION, params->size());
      } break;

      case NODE_CALL: {
        std::string ident = node->child(0)->repr();

        if (auto itr=m_macros.find(ident);
            itr != std::end(m_macros))
          {
            auto macro = itr->second;
            macro->call(*this, node, prog);
          }
        else
          {
            for (size_t i=0; i<node->size(); i++)
              {
                compile_prog(node->child(i), prog);
              }

            if (node->child(0)->type() == NODE_LAMBDA
                || node->child(0)->type() == NODE_CALL)
              {
                prog->add(OP_CALL_REF, node->size() - 1);
              }
            else if (node->child(0)->type() == NODE_IDENT)
              {
                if (auto entry = m_sym->find(node->child(0)->repr());
                    entry && entry->is_global() == false)
                  {
                    size_t arity = entry->node()->child(0)->size();
                    if (arity != node->size() - 1)
                      {
                        std::stringstream ss;
                        ss << "arity mismatch: " << node->child(0)->repr()
                           << " expect '"
                           << arity
                           << "' arguments, got '"
                           << node->size() - 1
                           << "'";

                        node->loc().error<compile_error>(LOG_ERROR, ss.str());
                      }

                    prog->add(OP_CALL_REF, node->size() - 1);
                  }
                else
                  {
                    prog->add(OP_CALL_NATIVE, node->size() - 1);
                  }
              }
          }
      } break;

      case NODE_IDENT: {
        auto entry = m_sym->find(node->repr());

        if (!entry)
          {
            std::stringstream ss;
            ss << "'"
               << node->repr()
               << "' is undefined";

            node->loc().error<compile_error>(LOG_ERROR, ss.str());
          }

        if (!entry->is_global())
          {
            prog->add(OP_LOAD_LOCAL, entry->addr());
          }
        else
          {
            prog->add(OP_LOAD_GLOBAL, entry->addr());
          }

      } break;

      case NODE_INT: {
        prog->load_const(std::make_shared<Constant>(TYPE_INT,
                                                    stoi(node->repr()),
                                                    node->loc()));
      } break;

      case NODE_FLOAT: {
        prog->load_const(std::make_shared<Constant>(TYPE_FLOAT,
                                                    stof(node->repr()),
                                                    node->loc()));
      } break;

      case NODE_BOOL: {
        prog->load_const(std::make_shared<Constant>(TYPE_BOOL,
                                                    node->repr() == "true",
                                                    node->loc()));
      } break;

      case NODE_STRING: {
        std::string str = node->repr();
        prog->load_const(std::make_shared<Constant>(TYPE_STRING,
                                                    str.substr(1, str.size() - 2),
                                                    node->loc()));
      } break;

      default: {
        std::stringstream ss;
        ss << "cannot compile expression '"
           << (NodeTypeStr[node->type()] + strlen("NODE_"))
           << "'";

        node->loc().error<compile_error>(LOG_ERROR, ss.str());
      } break;
      }
  }
}