#include "../src/Loader.hpp"
#include "src/Logger.hpp"
#include "src/Value.hpp"
#include "src/opcodes.hpp"
#include "../src/Module.hpp"

GRINO_ERROR(assertion_error);

extern "C" void lib_collection(grino::Loader& loader)
{
  loader.add_native("ref", [&loader](auto args){
    std::vector<size_t> idxs;
    for (size_t i=1; i<args.size(); i++)
      {
        idxs.push_back(args[i]->as_int());
      }

    if (args[0]->type() == grino::TYPE_STRING)
      {
        auto char_val = std::string(1, args[0]->as_string()[args[1]->as_int()]);
        return grino::Value::make_string(args[0]->loc(), char_val);
      }

    std::function<std::shared_ptr<grino::Value>
                  (std::shared_ptr<grino::Value>, std::vector<size_t>)>
      f = [&](std::shared_ptr<grino::Value> val, std::vector<size_t> indexes){
        if (indexes.empty())
          {
            return val;
          }

        size_t index = indexes[0];
        std::vector<size_t> next;

        for (size_t i=1; i<indexes.size(); i++)
          {
            next.push_back(indexes[i]);
          }

        return f(loader.vm().heap(val->as_ref())->as_array()[index], next);
      };

    return f(args[0], idxs);
  });

  loader.add_native("empty?", [&loader](auto args){
    auto val = args[0];
    if (val->type() == grino::TYPE_STRING)
      {
        std::string str = val->as_string();
        return grino::Value::make_bool(val->loc(), str.empty());
      }

    size_t ref = val->as_ref();
    auto array_val = loader.vm().heap(ref);
    auto& array = array_val->as_array();

    return grino::Value::make_bool(array_val->loc(), array.empty());
  });

  loader.add_native("len", [&loader](auto args){
    auto val = args[0];

    if (val->type() == grino::TYPE_STRING)
      {
        return grino::Value::make_int(val->loc(), val->as_string().size());
      }

    size_t ref = val->as_ref();
    auto array_val = loader.vm().heap(ref);
    auto& array = array_val->as_array();

    return grino::Value::make_int(array_val->loc(), array.size());
  });
}

extern "C" void lib_array(grino::Loader& loader)
{
  loader.add_native("head", [&loader](auto args){
    auto ref_val = args[0];
    size_t ref = ref_val->as_ref();
    auto array_val = loader.vm().heap(ref);
    auto& array = array_val->as_array();

    return array[0];
  });

  loader.add_native("tail", [&loader](auto args){
    auto ref_val = args[0];
    size_t ref = ref_val->as_ref();
    auto array_val = loader.vm().heap(ref);
    auto& array = array_val->as_array();

    grino::val_array_t data;
    for (size_t i=1; i<array.size(); i++)
      {
        data.push_back(array[i]);
      }

    auto res = grino::Value::make_array(array_val->loc(), data);
    size_t addr = loader.vm().heap_size();
    loader.vm().set_heap(addr, res);

    return grino::Value::make_ref(array_val->loc(), addr);
  });

  loader.add_native("cons", [&loader](auto args){
    auto ref_val = args[1];
    size_t ref = ref_val->as_ref();
    auto array_val = loader.vm().heap(ref);
    auto& array = array_val->as_array();

    array.insert(std::begin(array), args[0]);

    auto res = grino::Value::make_array(array_val->loc(), array);
    size_t addr = loader.vm().heap_size();
    loader.vm().set_heap(addr, res);

    return grino::Value::make_ref(array_val->loc(), addr);
  });
}

extern "C" void lib_flow_control(grino::Loader& loader)
{
  loader.add_static("if", [](auto& compiler, auto node, auto& prog, auto& sym) {
    auto cond_node = node->child(1).lock();
    auto then_node = node->child(2).lock();
    auto else_node = node->child(3).lock();

    compiler.compile(cond_node, prog, sym);
    size_t brf = prog.size();
    prog.push_instr(grino::OPCODE_BRF, 0);

    compiler.compile(then_node, prog, sym);
    size_t br = prog.size();
    prog.push_instr(grino::OPCODE_BR, 0);

    prog.set_param(brf, prog.size());
    compiler.compile(else_node, prog, sym);

    prog.set_param(br, prog.size());
  });
}

extern "C" void lib_int(grino::Loader& loader)
{
  loader.add_native("+", [](auto args){
    int result = 0;

    for (auto value: args)
      {
        result += value->as_int();
      }

    grino::Loc loc {"???", 0};

    if (args.empty() == false)
      {
        loc = args.front()->loc();
      }

    return grino::Value::make_int(loc, result);
  });

  loader.add_native("*", [](auto args){
    int result = 1;

    for (auto value: args)
      {
        result *= value->as_int();
      }

    grino::Loc loc {"???", 0};

    if (args.empty() == false)
      {
        loc = args.front()->loc();
      }

    return grino::Value::make_int(loc, result);
  });

  loader.add_native("-", [](auto args){
    grino::Loc loc {"???", 0};

    if (args.empty() == false)
      {
        loc = args.front()->loc();
      }
    else
      {
        return grino::Value::make_int(loc, 0);
      }

    int result = args[0]->as_int();

    if (args.size() == 1)
      {
        return grino::Value::make_int(loc, -result);
      }

    for (size_t i=1; i<args.size(); i++)
      {
        result -= args[i]->as_int();
      }

    return grino::Value::make_int(loc, result);
  });

  loader.add_native("/", [](auto args){
    grino::Loc loc {"???", 0};

    if (args.empty() == false)
      {
        loc = args.front()->loc();
      }
    else
      {
        return grino::Value::make_int(loc, 1);
      }

    int result = args[0]->as_int();

    if (args.size() == 1)
      {
        return grino::Value::make_int(loc, result);
      }

    for (size_t i=1; i<args.size(); i++)
      {
        result /= args[i]->as_int();
      }

    return grino::Value::make_int(loc, result);
  });

  loader.add_native("%", [](auto args){
    grino::Loc loc {"???", 0};

    if (args.empty() == false)
      {
        loc = args.front()->loc();
      }
    else
      {
        return grino::Value::make_int(loc, 1);
      }

    int result = args[0]->as_int();

    if (args.size() == 1)
      {
        return grino::Value::make_int(loc, result);
      }

    for (size_t i=1; i<args.size(); i++)
      {
        result %= args[i]->as_int();
      }

    return grino::Value::make_int(loc, result);
  });

  loader.add_native("^", [](auto args){
    grino::Loc loc {"???", 0};

    if (args.empty() == false)
      {
        loc = args.front()->loc();
      }
    else
      {
        return grino::Value::make_int(loc, 1);
      }

    int result = args[0]->as_int();

    if (args.size() == 1)
      {
        return grino::Value::make_int(loc, result);
      }

    for (size_t i=1; i<args.size(); i++)
      {
        result = std::pow(result, args[i]->as_int());
      }

    return grino::Value::make_int(loc, result);
  });
}

extern "C" void lib_float(grino::Loader& loader)
{
  loader.add_native("+.", [](auto args){
    float result = 0;

    for (auto value: args)
      {
        result += value->as_float();
      }

    grino::Loc loc {"???", 0};

    if (args.empty() == false)
      {
        loc = args.front()->loc();
      }

    return grino::Value::make_float(loc, result);
  });

  loader.add_native("*.", [](auto args){
    float result = 1;

    for (auto value: args)
      {
        result *= value->as_float();
      }

    grino::Loc loc {"???", 0};

    if (args.empty() == false)
      {
        loc = args.front()->loc();
      }

    return grino::Value::make_float(loc, result);
  });

  loader.add_native("-.", [](auto args){
    grino::Loc loc {"???", 0};

    if (args.empty() == false)
      {
        loc = args.front()->loc();
      }
    else
      {
        return grino::Value::make_float(loc, 0);
      }

    float result = args[0]->as_float();

    if (args.size() == 1)
      {
        return grino::Value::make_float(loc, -result);
      }

    for (size_t i=1; i<args.size(); i++)
      {
        result -= args[i]->as_float();
      }

    return grino::Value::make_float(loc, result);
  });

  loader.add_native("/.", [](auto args){
    grino::Loc loc {"???", 0};

    if (args.empty() == false)
      {
        loc = args.front()->loc();
      }
    else
      {
        return grino::Value::make_float(loc, 1.0f);
      }

    float result = args[0]->as_float();

    if (args.size() == 1)
      {
        return grino::Value::make_float(loc, result);
      }

    for (size_t i=1; i<args.size(); i++)
      {
        result /= args[i]->as_float();
      }

    return grino::Value::make_float(loc, result);
  });

  loader.add_native("%.", [](auto args){
    grino::Loc loc {"???", 0};

    if (args.empty() == false)
      {
        loc = args.front()->loc();
      }
    else
      {
        return grino::Value::make_float(loc, 1);
      }

    float result = args[0]->as_float();

    if (args.size() == 1)
      {
        return grino::Value::make_float(loc, result);
      }

    for (size_t i=1; i<args.size(); i++)
      {
        result = std::fmod(result, args[i]->as_float());
      }

    return grino::Value::make_float(loc, result);
  });

  loader.add_native("^.", [](auto args){
    grino::Loc loc {"???", 0};

    if (args.empty() == false)
      {
        loc = args.front()->loc();
      }
    else
      {
        return grino::Value::make_float(loc, 1);
      }

    float result = args[0]->as_float();

    if (args.size() == 1)
      {
        return grino::Value::make_int(loc, result);
      }

    for (size_t i=1; i<args.size(); i++)
      {
        result = std::pow(result, args[i]->as_float());
      }

    return grino::Value::make_float(loc, result);
  });
}

extern "C" void lib_cmp(grino::Loader& loader)
{
  loader.add_native("<", [](auto args){
    if (args[0]->type() == grino::TYPE_INT)
      {
        int lhs = args[0]->as_int();
        int rhs = args[1]->as_int();

        return grino::Value::make_bool(args[0]->loc(), lhs < rhs);
      }
    else if (args[0]->type() == grino::TYPE_FLOAT)
      {
        float lhs = args[0]->as_float();
        float rhs = args[1]->as_float();

        return grino::Value::make_bool(args[0]->loc(), lhs < rhs);
      }

    assert(0);
  });

  loader.add_native("<=", [](auto args){
    if (args[0]->type() == grino::TYPE_INT)
      {
        int lhs = args[0]->as_int();
        int rhs = args[1]->as_int();

        return grino::Value::make_bool(args[0]->loc(), lhs <= rhs);
      }
    else if (args[0]->type() == grino::TYPE_FLOAT)
      {
        float lhs = args[0]->as_float();
        float rhs = args[1]->as_float();

        return grino::Value::make_bool(args[0]->loc(), lhs <= rhs);
      }

    assert(0);
  });

  loader.add_native(">", [](auto args){
    if (args[0]->type() == grino::TYPE_INT)
      {
        int lhs = args[0]->as_int();
        int rhs = args[1]->as_int();

        return grino::Value::make_bool(args[0]->loc(), lhs > rhs);
      }
    else if (args[0]->type() == grino::TYPE_FLOAT)
      {
        float lhs = args[0]->as_float();
        float rhs = args[1]->as_float();

        return grino::Value::make_bool(args[0]->loc(), lhs > rhs);
      }

    assert(0);
  });

  loader.add_native(">=", [](auto args){
    if (args[0]->type() == grino::TYPE_INT)
      {
        int lhs = args[0]->as_int();
        int rhs = args[1]->as_int();

        return grino::Value::make_bool(args[0]->loc(), lhs >= rhs);
      }
    else if (args[0]->type() == grino::TYPE_FLOAT)
      {
        float lhs = args[0]->as_float();
        float rhs = args[1]->as_float();

        return grino::Value::make_bool(args[0]->loc(), lhs >= rhs);
      }

    assert(0);
  });

  loader.add_native("ne?", [](auto args){
    auto lhs = args[0];
    auto rhs = args[1];

    return grino::Value::make_bool(args[0]->loc(), !lhs->equals(*rhs));
  });
}

extern "C" void lib_bool(grino::Loader& loader)
{
  loader.add_native("not", [](auto args){
    return grino::Value::make_bool(args[0]->loc(), !args[0]->as_bool());
  });

  loader.add_static("and", [](auto& compiler, auto node,
                              auto& program,
                              auto& sym){
    std::vector<size_t> to_false;

    for (size_t i=1; i<node->size(); i++)
      {
        compiler.compile(node->child(i).lock(), program, sym);
        to_false.push_back(program.size());
        program.push_instr(grino::OPCODE_BRF, 0 /* to false */);
      }

    size_t addr = program
      .push_constant(grino::Value::make_bool(node->loc(), true));
    program.push_instr(grino::OPCODE_LOAD_CONST, addr);
    size_t to_end = program.size();
    program.push_instr(grino::OPCODE_BR, 0 /* to end */);

    for (auto a: to_false)
      {
        program.set_param(a, program.size());
      }

    addr = program
      .push_constant(grino::Value::make_bool(node->loc(), false));
    program.push_instr(grino::OPCODE_LOAD_CONST, addr);

    program.set_param(to_end, program.size());
  });

  loader.add_static("or", [](auto& compiler, auto node,
                             auto& program, auto& sym){
    std::vector<size_t> to_true;

    for (size_t i=1; i<node->size(); i++)
      {
        compiler.compile(node->child(i).lock(), program, sym);
        program.push_instr(grino::OPCODE_NOT);
        to_true.push_back(program.size());
        program.push_instr(grino::OPCODE_BRF, 0 /* to true */);
      }

    size_t addr = program
      .push_constant(grino::Value::make_bool(node->loc(), false));
    program.push_instr(grino::OPCODE_LOAD_CONST, addr);
    size_t to_end = program.size();
    program.push_instr(grino::OPCODE_BR, 0 /* to end */);

    for (auto a: to_true)
      {
        program.set_param(a, program.size());
      }

    addr = program
      .push_constant(grino::Value::make_bool(node->loc(), true));
    program.push_instr(grino::OPCODE_LOAD_CONST, addr);

    program.set_param(to_end, program.size());
  });
}

extern "C" void lib_assert(grino::Loader& loader)
{
  loader.add_native("assert", [](auto args){

    for (auto value: args)
      {
        if (!value->as_bool())
          {
            grino::Logger logger;
            logger.log<assertion_error>(grino::LOG_ASSERT,
                                        value->loc(),
                                        "assertion failed");
          }
      }

    return grino::Value::make_bool(args.front()->loc(), true);
  });

  loader.add_native("assert=", [](auto args){
    auto lhs = args.front();

    for (size_t i=1; i<args.size(); i++)
      {
        auto rhs = args[i];

        if (!lhs->equals(*rhs))
          {
            grino::Logger logger;
            std::stringstream ss;

            ss << "'" << lhs->string() << "' is not equals to '"
               << rhs->string() << "'";

            logger.log<assertion_error>(grino::LOG_ASSERT,
                                        lhs->loc(),
                                        ss.str());

          }
      }

    return grino::Value::make_bool(args.front()->loc(), true);
  });
}

extern "C" void lib_string(grino::Loader& loader)
{
  loader.add_native("dup", [](auto args) {
    std::string value = args[0]->as_string();
    size_t count = args[1]->as_int();
    std::string result;

    for (size_t i=0; i<count; i++)
      {
        result += value;
      }

    return grino::Value::make_string(args[0]->loc(), result);
  });

  loader.add_native("cat", [](auto args) {
    std::string result;

    for (auto arg: args)
      {
        result += arg->as_string();
      }

    return grino::Value::make_string(args[0]->loc(), result);
  });
}

extern "C" void lib(grino::Loader& loader)
{
  lib_assert(loader);
  lib_int(loader);
  lib_float(loader);
  lib_cmp(loader);
  lib_bool(loader);
  lib_flow_control(loader);
  lib_collection(loader);
  lib_array(loader);
  lib_string(loader);

  loader.add_native("dump", [](auto args){
    std::string sep;
    std::stringstream ss;

    for (auto arg: args)
      {
        ss << sep << arg->string();
        sep = " ";
      }

    std::cout << ss.str() << std::endl;

    return grino::Value::make_nil(args.back()->loc());
  });

  loader.add_static("set!", [](auto& compiler,
                               auto node,
                               auto& program,
                               auto& sym){
    std::string ident = node->child(1).lock()->repr();
    auto expr = node->child(2).lock();

    auto entry = sym.find(ident, compiler.scope());

    if (entry)
      {
        compiler.compile(expr, program, sym);
        program.push_instr(grino::OPCODE_STORE_LOCAL, entry->addr);
      }

    program.push_value(grino::Value::make_nil(node->loc()));
  });

  loader.add_native("eq?", [](auto args){
    auto lhs = args.front();

    for (size_t i=1; i<args.size(); i++)
      {
        auto rhs = args[i];

        if (!lhs->equals(*rhs))
          {
            return grino::Value::make_bool(rhs->loc(), false);
          }
      }

    return grino::Value::make_bool(args.front()->loc(), true);
  });
}