#include "stdlib.hpp"

#include "common.hpp"
#include "pod.hpp"
#include "sourcerange.hpp"
#include "writer.hpp"

typedef Result<StackFrame> (*StdlibFunctionIdPtr)(const StackFrame& stack);

struct StdlibFunctionEntry {
  const char* name;
  StdlibFunctionId fun_id;
  StdlibFunctionIdPtr fun_ptr;
};

Result<StackFrame> stdlib_unknown(const StackFrame& stack) {
  return ERROR(NotImplemented);
}

Result<StackFrame> stdlib_print(const StackFrame& stack) {
  auto params = TRY(stack.get(0));
  auto size = TRY(params.size());

  String str = TRY(String::create(""));

  for (uint64_t i = 0; i < size; i++) {
    Value param = TRY(params.get(i));

    if (i != 0) {
      str = TRY(str.concat(" "));
    }
    if (param.is<String>()) {
      str = TRY(str.concat(*param.to<String>()));
    } else {
      auto s = TRY(write_one(param));
      str = TRY(str.concat(s));
    }
  }

  TRY(print_string(str));

  auto nil = Value(TRY(Nil::create()));
  auto res = TRY(stack.set(0, nil));

  res = TRY(res.ret(0));

  return res;
}

Result<StackFrame> stdlib_println(const StackFrame& stack) {
  auto res = TRY(stdlib_print(stack));
  std::cout << "\n";

  return res;
}

Result<StackFrame> stdlib_prn(const StackFrame& stack) {
  return ERROR(NotImplemented);
}

Result<StackFrame> stdlib_assert(const StackFrame& stack) {
  auto params = TRY(stack.get(0));
  auto size = TRY(params.size());
  for (uint64_t i = 0; i < size; i++) {
    Value param = TRY(params.get(i));
    if (!param.is<Bool>()) return ERROR(AssertionFailed);
    auto v = param.to<Bool>()->value();
    if (!v) return ERROR(AssertionFailed);
  }

  auto nil = Value(TRY(Nil::create()));
  auto res = TRY(stack.set(0, nil));

  res = TRY(res.ret(0));

  return res;
}

Result<StackFrame> stdlib_dict(const StackFrame& stack) {
  auto params = TRY(stack.get(0));
  if (!params.is<Array>()) return ERROR(TypeMismatch);
  Value d = TRY(Dict::create(*params.to<Array>()));

  auto res = TRY(stack.set(0, d));

  res = TRY(res.ret(0));

  return res;
}

Result<StackFrame> stdlib_list(const StackFrame& stack) {
  auto params = TRY(stack.get(0));
  if (!params.is<Array>()) return ERROR(TypeMismatch);
  Value d = TRY(Pair::create(*params.to<Array>()));

  auto res = TRY(stack.set(0, d));

  res = TRY(res.ret(0));

  return res;
}

Result<StackFrame> stdlib_array(const StackFrame& stack) {
  auto params = TRY(stack.get(0));

  auto res = TRY(stack.set(0, params));

  res = TRY(res.ret(0));

  return res;
}

Result<StackFrame> stdlib_get(const StackFrame& stack) {
  auto params = TRY(stack.get(0));
  auto size = TRY(params.size());
  if (size != 2) return ERROR(ArgumentCountMismatch);

  Value collection = TRY(params.get(0));
  Value key = TRY(params.get(1));

  auto val = TRY(collection.get(key));

  auto res = TRY(stack.set(0, val));

  res = TRY(res.ret(0));

  return res;
}

Result<StackFrame> stdlib_set(const StackFrame& stack) {
  auto params = TRY(stack.get(0));
  auto size = TRY(params.size());
  if (size != 3) return ERROR(ArgumentCountMismatch);

  Value collection = TRY(params.get(0));
  Value key = TRY(params.get(1));
  Value value = TRY(params.get(2));

  collection = TRY(collection.set(key, value));

  auto res = TRY(stack.set(0, collection));

  res = TRY(res.ret(0));

  return res;
}

Result<StackFrame> stdlib_srcloc(const StackFrame& stack) {
  auto params = TRY(stack.get(0));
  auto size = TRY(params.size());
  if (size != 6) return ERROR(ArgumentCountMismatch);

  SourceRange sr;
  for (uint64_t i = 0; i < 6; i++) {
    auto val = TRY(params.get(i));
    if (!val.is<Int64>()) return ERROR(TypeMismatch);
    int64_t intval = val.to<Int64>()->value();
    switch (i) {
      case 0:
        sr.start.line = intval;
        break;
      case 1:
        sr.start.column = intval;
        break;
      case 2:
        sr.start.offset = intval;
        break;
      case 3:
        sr.end.line = intval;
        break;
      case 4:
        sr.end.column = intval;
        break;
      case 5:
        sr.end.offset = intval;
        break;
      default:
        break;
    }
  }

  auto val = Value(TRY(SrcLoc::create(sr)));

  auto res = TRY(stack.set(0, val));

  res = TRY(res.ret(0));

  return res;
}

Result<StackFrame> stdlib_size(const StackFrame& stack) {
  auto params = TRY(stack.get(0));
  auto size = TRY(params.size());

  if (size != 1) return ERROR(ArgumentCountMismatch);

  auto param = TRY(params.get(0));
  auto psize = TRY(param.size());

  auto val = Value(TRY(Int64::create((int64_t)psize)));

  auto res = TRY(stack.set(0, val));

  res = TRY(res.ret(0));

  return res;
}

Result<StackFrame> stdlib_map(const StackFrame& stack_param) {
  auto stack = TRY(stack_param.copy());
  auto params = TRY(stack.get(0));
  auto stack_size = TRY(stack.size());

  auto fun = TRY(params.get(1));

  if (stack_size == 1) {
    // This is the first call, need to set up arguments on the stack
    auto collection = TRY(params.get(0));

    auto num_params = TRY(params.size());
    if (num_params != 2) return ERROR(ArgumentCountMismatch);

    if (!collection.is<Pair>()) return ERROR(TypeMismatch);

    // Result accumulator will contain nil value initially
    auto nil = Value(TRY(Nil::create()));
    stack = TRY(stack.set(1, nil));

    // Remaining elements in the collection
    stack = TRY(stack.set(2, TRY(collection.rest())));

    stack = TRY(stack.set(3, fun));
    stack = TRY(stack.set(4, TRY(collection.first())));

    return TRY(stack.call(fun, 3, 5));
  }

  auto accumulator = TRY(stack.get(1));
  auto remaining = TRY(stack.get(2));
  auto fun_result = TRY(stack.get(3));

  accumulator = Value(TRY(Pair::create(fun_result, accumulator)));

  stack = TRY(stack.set(1, accumulator));

  if (!remaining.is<Nil>()) {
    stack = TRY(stack.set(2, TRY(remaining.rest())));
    stack = TRY(stack.set(3, fun));
    stack = TRY(stack.set(4, TRY(remaining.first())));

    return TRY(stack.call(fun, 3, 5));
  }

  accumulator = TRY(reverse(accumulator));

  auto res = TRY(stack.set(0, accumulator));
  res = TRY(res.ret(0));

  return res;
}

#define STDLIB_FUNCTION(name, id)                                  \
  [(uint64_t)StdlibFunctionId::id] = {#name, StdlibFunctionId::id, \
                                      stdlib_##name}

static StdlibFunctionEntry function_entries[] = {
    STDLIB_FUNCTION(unknown, Unknown),
    STDLIB_FUNCTION(print, Print),
    STDLIB_FUNCTION(println, PrintLn),
    STDLIB_FUNCTION(prn, Prn),
    STDLIB_FUNCTION(assert, Assert),
    STDLIB_FUNCTION(dict, Dict),
    STDLIB_FUNCTION(list, List),
    STDLIB_FUNCTION(array, Array),
    STDLIB_FUNCTION(get, Get),
    STDLIB_FUNCTION(set, Set),
    STDLIB_FUNCTION(srcloc, SrcLoc),
    STDLIB_FUNCTION(size, Size),
    STDLIB_FUNCTION(map, Map),
    [(uint64_t)StdlibFunctionId::Max] = {0, StdlibFunctionId::Max,
                                         stdlib_unknown},
};

// TODO: this just scans through the array of functions linearly every time.
// It doesn't have much effect at runtime, because such lookup is likely to
// happen on compile-time only. But anyway, this is worth speeding up
// eventually.
StdlibFunctionEntry get_function_entry(const char* name) {
  for (uint64_t i = 1; i < (uint64_t)StdlibFunctionId::Max; i++) {
    if (strcmp(function_entries[i].name, name) == 0) {
      return function_entries[i];
    }
  }
  return function_entries[(uint64_t)StdlibFunctionId::Unknown];
}

Result<StdlibFunctionId> get_stdlib_function(const Symbol& name) {
  const uint64_t bufsize = 256;
  char buf[bufsize];
  auto size = TRY(name.size());
  if (size + 1 > bufsize) {
    return ERROR(KeyError);
  }

  for (uint64_t i = 0; i < size; i++) {
    buf[i] = (char)TRY(name[i]);
  }
  buf[size] = 0;

  StdlibFunctionId fun_id = get_function_entry(buf).fun_id;
  if (fun_id == StdlibFunctionId::Unknown) {
    return ERROR(KeyError);
  }

  return fun_id;
}

Result<StackFrame> call_stdlib_function(StdlibFunctionId fun_id,
                                        const StackFrame& stack) {
  if (fun_id == StdlibFunctionId(0) || fun_id >= StdlibFunctionId::Max) {
    return ERROR(KeyError);
  }

  return function_entries[(uint64_t)fun_id].fun_ptr(stack);
}

Result<const char*> get_stdlib_function_name(StdlibFunctionId fun_id) {
  if (fun_id == StdlibFunctionId(0) || fun_id >= StdlibFunctionId::Max) {
    return ERROR(KeyError);
  }

  return function_entries[(uint64_t)fun_id].name;
}