#include "VM.hpp"
#include "Function.hpp"
#include "Code.hpp"
#include "lib/Opcodes.hpp"

namespace jk
{
  /*explicit*/ VM::VM()
  {
  }

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

  void VM::execute(std::shared_ptr<Program> program)
  {
    Frame frame;
    frame.program = program;
    m_frames.push_back(frame);
    execute();
  }

  void VM::execute()
  {
    while (m_pc < m_frames.back().program->size())
      {
        Instr instr = m_frames.back().program->get(m_pc);

        switch (instr.opcode)
          {
          case OPCODE_MK_FUNCTION: {
            auto code = program()->constant(pop());

            auto prog = code->as_code()->program();

            std::shared_ptr<Function> function;

            if (prog)
              {
                function = std::make_shared<Function>(prog);
              }
            else
              {
                 function = std::make_shared<Function>
                  (code->as_code()->foreign());
              }

            auto value = Value::make_function(function);

            size_t addr = push_heap(value);
            auto ref = Value::make_ref(addr);

            push(program()->push_constant(ref));

            m_pc++;

          } break;

          case OPCODE_PUSH_CONST: {
            push(*instr.param);
            m_pc++;
          } break;

          case OPCODE_LOAD: {
            auto value = m_frames.back().locals[*instr.param];
            push(program()->push_constant(value));
            m_pc++;
          } break;

          case OPCODE_STORE: {
            auto idx = pop();

            m_frames.back().locals[*instr.param] =
              program()->constant(idx);

            m_pc++;
          } break;

          case OPCODE_LOAD_GLOBAL: {
            auto value = m_globals[*instr.param];
            push(program()->push_constant(value));
            m_pc++;
          } break;

          case OPCODE_RET: {
            size_t pc = m_frames.back().ret_addr;
            size_t stack_sz = m_frames.back().stack_sz;
            auto prog = m_frames.back().program;

            param_t ret = pop();
            auto ret_val = prog->constant(ret);
            assert(ret_val);

            while (m_stack.size() > stack_sz)
              {
                pop();
              }

            m_frames.pop_back();

            push(program()->push_constant(ret_val));

            m_pc = pc + 1;
          } break;

          case OPCODE_CALL: {
            std::vector<std::shared_ptr<Value>> args;

            for (size_t i=0; i<*instr.param; i++)
              {
                args.insert(std::begin(args), program()->constant(pop()));
              }

            auto ref_val = program()->constant(pop());
            auto ref = ref_val->as_ref();
            auto fun_val = heap(ref);
            auto fun = fun_val->as_function();

            if (auto prog = fun->program();
                prog)
              {
                Frame frame;
                frame.program = prog;
                frame.ret_addr = m_pc;
                frame.stack_sz = m_stack.size();
                m_pc = 0;

                for (size_t i=0; i<args.size(); i++)
                  {
                    frame.locals[i] = args[i];
                  }

                m_frames.push_back(frame);
              }
            else
              {
                auto result = fun->call(args);
                push(program()->push_constant(result));
                m_pc++;
              }
          } break;

          default:
            std::cerr << "cannot execute unknown opcode "
                      << OpcodeTypeStr[instr.opcode]
                      << std::endl;
            abort();
          }
      }
  }

  std::string VM::string() const
  {
    std::stringstream ss;

    for (size_t i=0; i<m_stack.size(); i++)
      {
        ss << i << "\t" << m_stack[i] << std::endl;
      }

    return ss.str();
  }

  size_t VM::push_heap(std::shared_ptr<Value> value)
  {
    m_heap.push_back(value);
    return m_heap.size() - 1;
  }

  std::shared_ptr<Value> VM::heap(size_t addr)
  {
    assert(addr < m_heap.size());
    return m_heap[addr];
  }

  void VM::set_global(size_t addr, std::shared_ptr<Value> value)
  {
    m_globals[addr] = value;
  }

  void VM::push(param_t param)
  {
    m_stack.push_back(param);
  }

  param_t VM::pop()
  {
    auto param = m_stack.back();
    m_stack.pop_back();
    return param;
  }

  std::shared_ptr<Program> VM::program() const
  {
    return m_frames.back().program;
  }
}