#include <nan.h>
#include <v8-profiler.h>
#include <stdlib.h>
#include <atomic>
#if defined(_WIN32)
#include <time.h>
#define snprintf _snprintf
#else
#include <sys/time.h>
#endif

using namespace v8;

char filename[256];
bool addTimestamp;
std::atomic<bool> processingOOM(false);

class FileOutputStream : public OutputStream
{
public:
  FileOutputStream(FILE *stream) : stream_(stream) {}
  virtual int GetChunkSize()
  {
    return 65536;
  }
  virtual void EndOfStream() {}
  virtual WriteResult WriteAsciiChunk(char *data, int size)
  {
    const size_t len = static_cast<size_t>(size);
    size_t off = 0;
    while (off < len && !feof(stream_) && !ferror(stream_))
      off += fwrite(data + off, 1, len - off, stream_);
    return off == len ? kContinue : kAbort;
  }

private:
  FILE *stream_;
};

size_t RaiseLimit(void *data, size_t current_heap_limit, size_t initial_heap_limit)
{
  return current_heap_limit + 10u * 1024 * 1024; // 10MiB
}

void OnOOMErrorHandler()
{
  if (processingOOM.exchange(true))
  {
    fprintf(stderr, "FATAL: OnOOMError called more than once.\n");
    exit(2);
  }

  if (addTimestamp)
  {
    // Add timestamp to filename
    time_t rawtime;
    struct tm timeinfo;
    time(&rawtime);
#if defined(_WIN32)
    localtime_s(&timeinfo, &rawtime);
#else
    localtime_r(&rawtime, &timeinfo);
#endif

    char *pch = strstr(filename, ".heapsnapshot");
    if (pch != NULL)
    {
      *pch = '\0';
    }
    strncat(filename, "-%Y%m%dT%H%M%S.heapsnapshot", sizeof(filename) - strlen(filename) - 1);

    char newFilename[256];
    if (strftime(newFilename, sizeof(newFilename), filename, &timeinfo) == 0)
    {
      // strftime failed (buffer too small or format error); keep base filename
      snprintf(newFilename, sizeof(newFilename), "%s", filename);
    }
    strncpy(filename, newFilename, sizeof(filename) - 1);
    filename[sizeof(filename) - 1] = '\0';
  }

  fprintf(stderr, "Generating Heapdump to '%s' now...\n", filename);
  FILE *fp = fopen(filename, "w");
  if (fp == NULL)
  {
    fprintf(stderr, "FATAL: Failed to open '%s' for writing heapdump.\n", filename);
    abort();
  }

  auto *isolate = v8::Isolate::GetCurrent();

  // Capturing a heap snapshot forces a garbage collection which can, in turn,
  // trigger the OOM flow which causes recursion. To prevent this, this callback
  // will raise the heap limit if the GC tries to go down that path again.
  // Normally we would want to add a call to RemoveNearHeapLimitCallback() after
  // we are done, but that is not necessary since we exit() before it matters.
  isolate->AddNearHeapLimitCallback(RaiseLimit, nullptr);

  // Create heapdump, depending on which Node.js version this can differ
  // for now, just support Node.js 7 and higher
  const HeapSnapshot *snap = isolate->GetHeapProfiler()->TakeHeapSnapshot();

  FileOutputStream stream(fp);
  snap->Serialize(&stream, HeapSnapshot::kJSON);
  fclose(fp);

  // Free the heap snapshot memory
  const_cast<HeapSnapshot *>(snap)->Delete();

  fprintf(stderr, "Done! Exiting process now.\n");
  exit(1);
}

#if NODE_VERSION_AT_LEAST(20, 0, 0)
void OnOOMErrorNode20(const char *location, const OOMDetails &details)
{
  // Node20+
  OnOOMErrorHandler();
}
#else
void OnOOMErrorNode18(const char *location, bool is_heap_oom)
{
  // Up until Node18
  OnOOMErrorHandler();
}
#endif

void ParseArgumentsAndSetErrorHandler(const FunctionCallbackInfo<Value> &args)
{
  Isolate *isolate = args.GetIsolate();
#if NODE_VERSION_AT_LEAST(20, 0, 0)
  // Node20+
  isolate->SetOOMErrorHandler(OnOOMErrorNode20);
#else
  // Up until Node18
  isolate->SetOOMErrorHandler(OnOOMErrorNode18);
#endif

// parse JS arguments
// 1: filename
// 2: addTimestamp boolean
#if NODE_VERSION_AT_LEAST(13, 0, 0)
  Local<Context> context = isolate->GetCurrentContext();
  String::Utf8Value fArg(isolate, args[0]->ToString(context).ToLocalChecked());
#elif NODE_VERSION_AT_LEAST(12, 0, 0)
  String::Utf8Value fArg(isolate, args[0]->ToString(isolate));
#elif NODE_VERSION_AT_LEAST(9, 0, 0)
  String::Utf8Value fArg(isolate, args[0]->ToString());
#else
  String::Utf8Value fArg(args[0]->ToString());
#endif
  strncpy(filename, (const char *)(*fArg), sizeof(filename) - 1);
  filename[sizeof(filename) - 1] = '\0';

#if NODE_VERSION_AT_LEAST(12, 0, 0)
  addTimestamp = args[1]->BooleanValue(isolate);
#else
  addTimestamp = args[1]->BooleanValue();
#endif
}

void init(Local<Object> exports)
{
  NODE_SET_METHOD(exports, "call", ParseArgumentsAndSetErrorHandler);
}

NODE_MODULE(NODE_OOM_HEAPDUMP_NATIVE, init)