#include "fun.hpp"

#define NUM_BINOP(OP)\
  auto lhs = args[0];                                                   \
  auto rhs = args[1];                                                   \
  bool res = false;                                                     \
                                                                        \
  if (lhs->type() == TYPE_INT)                                          \
    {                                                                   \
      res = std::get<int>(lhs->value()) OP std::get<int>(rhs->value());  \
    }                                                                   \
                                                                        \
  if (lhs->type() == TYPE_FLOAT)                                        \
    {                                                                   \
      res = std::get<float>(lhs->value()) OP std::get<float>(rhs->value()); \
    }                                                                   \
                                                                        \
  return std::make_shared<Constant>(TYPE_BOOL, res, loc)


namespace fkstd
{
  STDRET assert_eq(Loc loc, Module&, STDARGS args)
  {
    auto oracle = args[0];
    auto expr = args[1];

    if (!oracle->equals(*expr))
      {
        std::stringstream ss;
        ss << "assertion failed: expected '"
           << oracle->string()
           << "', got '"
           << expr->string()
           << "'";

        loc.error<assert_error>(LOG_ERROR, ss.str());
      }

    return std::make_shared<Constant>(TYPE_BOOL, true, loc);
  }

  STDRET println(Loc loc, Module&, STDARGS args)
  {
    std::string sep;

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

    if (args.size() > 0)
      {
        loc = args.front()->loc();
      }

    std::cout << std::endl;

    return std::make_shared<Constant>(TYPE_INT, 0, loc);
  }

  STDRET add_int(Loc loc, Module&, STDARGS args)
  {
    int result = 0;

    for (auto arg: args)
      {
        result += std::get<int>(arg->value());
      }

    return std::make_shared<Constant>(TYPE_INT, result, loc);
  }

  STDRET sub_int(Loc loc, Module&, STDARGS args)
  {
    if (args.empty())
      {
        return std::make_shared<Constant>(TYPE_INT, 0, loc);
      }

    if (args.size() == 1)
      {
        return std::make_shared<Constant>(TYPE_INT,
                                          -std::get<int>(args[0]->value()),
                                          loc);
      }

    int result = std::get<int>(args[0]->value());

    for (size_t i=1; i<args.size(); i++)
      {
        result -= std::get<int>(args[i]->value());
      }

    return std::make_shared<Constant>(TYPE_INT, result, loc);
  }

  STDRET mul_int(Loc loc, Module&, STDARGS args)
  {
    int result = 1;

    for (auto arg: args)
      {
        result *= std::get<int>(arg->value());
      }

    return std::make_shared<Constant>(TYPE_INT, result, loc);
  }

  STDRET div_int(Loc loc, Module&, STDARGS args)
  {
    if (args.empty())
      {
        return std::make_shared<Constant>(TYPE_INT, 1, loc);
      }

    if (args.size() == 1)
      {
        return std::make_shared<Constant>(TYPE_INT,
                                          1/std::get<int>(args[0]->value()),
                                          loc);
      }

    int result = std::get<int>(args[0]->value());

    for (size_t i=1; i<args.size(); i++)
      {
        result /= std::get<int>(args[i]->value());
      }

    return std::make_shared<Constant>(TYPE_INT, result, loc);
  }

  STDRET mod_int(Loc loc, Module&, STDARGS args)
  {
    if (args.empty())
      {
        return std::make_shared<Constant>(TYPE_INT, 1, loc);
      }

    int result = std::get<int>(args[0]->value());

    for (size_t i=1; i<args.size(); i++)
      {
        result %= std::get<int>(args[i]->value());
      }

    return std::make_shared<Constant>(TYPE_INT, result, loc);
  }

  STDRET pow_int(Loc loc, Module&, STDARGS args)
  {
    if (args.empty())
      {
        return std::make_shared<Constant>(TYPE_INT, 1, loc);
      }

    int result = std::get<int>(args[0]->value());

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

    return std::make_shared<Constant>(TYPE_INT, result, loc);
  }

  STDRET add_float(Loc loc, Module&, STDARGS args)
  {
    if (args.empty())
      {
        return std::make_shared<Constant>(TYPE_FLOAT, 0.0f, loc);
      }

    float result = std::get<float>(args[0]->value());

    for (size_t i=1; i<args.size(); i++)
      {
        result = result + std::get<float>(args[i]->value());
      }

    return std::make_shared<Constant>(TYPE_FLOAT, result, loc);
  }

  STDRET sub_float(Loc loc, Module&, STDARGS args)
  {
    if (args.empty())
      {
        return std::make_shared<Constant>(TYPE_FLOAT, 0.0f, loc);
      }

    float result = std::get<float>(args[0]->value());

    for (size_t i=1; i<args.size(); i++)
      {
        result = result - std::get<float>(args[i]->value());
      }

    if (args.size() == 1)
      {
        return std::make_shared<Constant>(TYPE_FLOAT, -result, loc);
      }

    return std::make_shared<Constant>(TYPE_FLOAT, result, loc);
  }

  STDRET mul_float(Loc loc, Module&, STDARGS args)
  {
    if (args.empty())
      {
        return std::make_shared<Constant>(TYPE_FLOAT, 1.0f, loc);
      }

    float result = std::get<float>(args[0]->value());

    for (size_t i=1; i<args.size(); i++)
      {
        result = result * std::get<float>(args[i]->value());
      }

    return std::make_shared<Constant>(TYPE_FLOAT, result, loc);
  }

  STDRET div_float(Loc loc, Module&, STDARGS args)
  {
    if (args.empty())
      {
        return std::make_shared<Constant>(TYPE_FLOAT, 1.0f, loc);
      }

    float result = std::get<float>(args[0]->value());

    for (size_t i=1; i<args.size(); i++)
      {
        result = result / std::get<float>(args[i]->value());
      }

    if (args.size() == 1)
      {
        return std::make_shared<Constant>(TYPE_FLOAT, 1.0f/result, loc);
      }

    return std::make_shared<Constant>(TYPE_FLOAT, result, loc);
  }

  STDRET mod_float(Loc loc, Module&, STDARGS args)
  {
    if (args.empty())
      {
        return std::make_shared<Constant>(TYPE_FLOAT, 1.0f, loc);
      }

    float result = std::get<float>(args[0]->value());

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

    return std::make_shared<Constant>(TYPE_FLOAT, result, loc);
  }

  STDRET pow_float(Loc loc, Module&, STDARGS args)
  {
    if (args.empty())
      {
        return std::make_shared<Constant>(TYPE_FLOAT, 1.0f, loc);
      }

    float result = std::get<float>(args[0]->value());

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

    return std::make_shared<Constant>(TYPE_FLOAT, result, loc);
  }

  STDRET lt(Loc loc, Module&, STDARGS args)
  {
    NUM_BINOP(<);
  }

  STDRET le(Loc loc, Module&, STDARGS args)
  {
    NUM_BINOP(<=);
  }

  STDRET gt(Loc loc, Module&, STDARGS args)
  {
    NUM_BINOP(>);
  }

  STDRET ge(Loc loc, Module&, STDARGS args)
  {
    NUM_BINOP(>=);
  }

  STDRET eq(Loc loc, Module&, STDARGS args)
  {
    NUM_BINOP(==);
  }

  STDRET ne(Loc loc, Module&, STDARGS args)
  {
    NUM_BINOP(!=);
  }

  STDRET bool_not(Loc loc, Module&, STDARGS args)
  {
    bool val = std::get<bool>(args[0]->value());
    return std::make_shared<Constant>(TYPE_BOOL, !val, loc);
  }

  STDRET array_ref(Loc, Module& mod, STDARGS args)
  {
    auto ref_val = args[0]->value();
    size_t ref = std::get<size_t>(ref_val);

    std::shared_ptr<Array> val =
      std::get<std::shared_ptr<Array>>(mod.vm()->load_global(ref));

    for (size_t i=1; i<args.size() - 1; i++)
      {
        int k = std::get<int>(args[i]->value());

        auto ref_val = val->at(k)->value();
        size_t ref = std::get<size_t>(ref_val);

        auto arr =
          std::get<std::shared_ptr<Array>>(mod.vm()->load_global(ref));

        val = arr;

      }

    return val->at(std::get<int>(args.back()->value()));
  }

  STDRET array_set(Loc, Module& mod, STDARGS args)
  {
    auto ref_val = args[0]->value();
    size_t ref = std::get<size_t>(ref_val);

    int index = std::get<int>(args[1]->value());

    std::shared_ptr<Array> val =
      std::get<std::shared_ptr<Array>>(mod.vm()->load_global(ref));

    val->set(index, args[2]);
    return args[2];
  }

  STDRET string_cat(Loc loc, Module&, STDARGS args)
  {
    if (args.empty())
      {
        return std::make_shared<Constant>(TYPE_STRING, "", loc);
      }

    std::string value = std::get<std::string>(args[0]->value());

    for (size_t i=1; i<args.size(); i++)
      {
        value += std::get<std::string>(args[i]->value());
      }

    return std::make_shared<Constant>(TYPE_STRING, value, loc);
  }

  STDRET string_dup(Loc loc, Module&, STDARGS args)
  {
    std::string str = std::get<std::string>(args[0]->value());
    int n = std::get<int>(args[1]->value());
    std::string res;

    for (int i=0; i<n; i++)
      {
        res += str;
      }

    return std::make_shared<Constant>(TYPE_STRING, res, loc);
  }

  STDRET char_at(Loc loc, Module&, STDARGS args)
  {
    std::string str = std::get<std::string>(args[0]->value());
    int n = std::get<int>(args[1]->value());

    return std::make_shared<Constant>(TYPE_STRING,
                                      std::string(1, str[n]),
                                      loc);
  }
}