#include "jsdhash.h"
#include "jsdbgapi.h"

#include "memory_profiler.h"
#include "server_socket.h"

#define SERVER_FILENAME "memory_profiler_server.js"

// TODO: Lots of code here is copied from spidermonkey-playground.cpp; we should
// consolidate it.

/* The error reporter callback. */
static void reportError(JSContext *cx, const char *message, JSErrorReport *report)
{
    fprintf(stderr, "%s:%u:%s\n",
            report->filename ? report->filename : "<no filename>",
            (unsigned int) report->lineno,
            message);
}

static JSClass global_class = {
  "serverGlobal", JSCLASS_GLOBAL_FLAGS,
  JS_PropertyStub, JS_PropertyStub, JS_PropertyStub, JS_PropertyStub,
  JS_EnumerateStub, JS_ResolveStub, JS_ConvertStub, JS_FinalizeStub,
  JSCLASS_NO_OPTIONAL_MEMBERS
};

// This native JS function prints the given string to the console.

static JSBool print(JSContext *cx, JSObject *obj, uintN argc, jsval *argv,
                    jsval *rval)
{
  char *str;

  if (!JS_ConvertArguments(cx, argc, argv, "s", &str))
    return JS_FALSE;

  printf("%s\n", str);

  return JS_TRUE;
}

// Private structure to track the state of tracing the JS heap.

typedef struct TracingState {
  // Keeps track of what objects we've visited so far.
  JSDHashTable visited;

  // Whether the tracing operation is successful or failed.
  JSBool result;

  // Runtime that we're tracing.
  JSRuntime *runtime;

  // Structure required to use JS tracing functions.
  JSTracer tracer;
};

// Static singleton for tracking the state of tracing the JS heap.
static TracingState tracingState;

typedef struct ChildTracingState {
  int num;
  void **things;
  uint32 *kinds;
  JSTracer tracer;
};

static ChildTracingState childTracingState;

// JSTraceCallback to build a hashtable of children.
static void childCountBuilder(JSTracer *trc, void *thing, uint32 kind)
{
  childTracingState.num++;
}

// JSTraceCallback to build a hashtable of children.
static void childBuilder(JSTracer *trc, void *thing, uint32 kind)
{
  *childTracingState.kinds = kind;
  childTracingState.kinds++;
  *childTracingState.things = thing;
  childTracingState.things++;
}

// JSTraceCallback to build a hashtable of existing object references.
static void visitedBuilder(JSTracer *trc, void *thing, uint32 kind)
{
  switch (kind) {
  case JSTRACE_OBJECT:
    JSDHashEntryStub *entry = (JSDHashEntryStub *)
      JS_DHashTableOperate(&tracingState.visited,
                           thing,
                           JS_DHASH_LOOKUP);
    if (JS_DHASH_ENTRY_IS_FREE((JSDHashEntryHdr *)entry)) {
      entry = (JSDHashEntryStub *) JS_DHashTableOperate(&tracingState.visited,
                                                        thing,
                                                        JS_DHASH_ADD);
      if (entry == NULL) {
        JS_ReportOutOfMemory(trc->context);
        tracingState.result = JS_FALSE;
        return;
      }
      entry->key = thing;
      JS_TraceChildren(trc, thing, kind);
    }
    break;
  case JSTRACE_DOUBLE:
    break;
  case JSTRACE_STRING:
    break;
  }
}

static intN countRoots(void *rp, const char *name, void *data)
{
  return JS_MAP_GCROOT_NEXT;
}

static intN rootMapFun(void *rp, const char *name, void *data)
{
  jsval **currIndex = (jsval **) data;
  **currIndex = INT_TO_JSVAL((unsigned int) rp);
  *currIndex++;
  return JS_MAP_GCROOT_NEXT;
}

static JSBool getChildrenInfo(JSContext *cx, JSObject *info,
                              JSObject *target, JSContext *targetCx)
{
  childTracingState.tracer.context = targetCx;
  childTracingState.tracer.callback = childCountBuilder;
  childTracingState.num = 0;
  JS_TraceChildren(&childTracingState.tracer, target, JSTRACE_OBJECT);

  void *things[childTracingState.num];
  uint32 kinds[childTracingState.num];

  childTracingState.things = things;
  childTracingState.kinds = kinds;

  childTracingState.tracer.callback = childBuilder;
  JS_TraceChildren(&childTracingState.tracer, target, JSTRACE_OBJECT);

  int numObjectChildren = 0;
  for (int i = 0; i < childTracingState.num; i++) {
    if (kinds[i] == JSTRACE_OBJECT)
      numObjectChildren++;
  }

  jsval childrenVals[numObjectChildren];

  int currChild = 0;
  for (int i = 0; i < childTracingState.num; i++) {
    if (kinds[i] == JSTRACE_OBJECT) {
      childrenVals[currChild] = INT_TO_JSVAL((unsigned int) things[i]);
      currChild += 1;
    }
  }

  JSObject *children = JS_NewArrayObject(cx, numObjectChildren, childrenVals);
  if (children == NULL) {
    JS_ReportOutOfMemory(cx);
    return JS_FALSE;
  }

  return JS_DefineProperty(cx, info, "children", OBJECT_TO_JSVAL(children),
                           NULL, NULL, JSPROP_ENUMERATE);
}

static JSBool getObjInfo(JSContext *cx, JSObject *obj, uintN argc,
                         jsval *argv, jsval *rval)
{
  uint32 id;

  if (!JS_ConvertArguments(cx, argc, argv, "u", &id))
    return JS_FALSE;

  JSDHashEntryStub *entry = (JSDHashEntryStub *)
    JS_DHashTableOperate(&tracingState.visited,
                         (void *) id,
                         JS_DHASH_LOOKUP);
  if (entry == NULL) {
    JS_ReportOutOfMemory(cx);
    return JS_FALSE;
  }

  if (JS_DHASH_ENTRY_IS_BUSY((JSDHashEntryHdr *)entry)) {
    JSObject *info = JS_NewObject(cx, NULL, NULL, NULL);

    JSObject *target = (JSObject *) id;
    JSContext *targetCx = tracingState.tracer.context;
    JSClass *classp = JS_GET_CLASS(targetCx, target);
    if (classp != NULL) {
      // TODO: Should really be using an interned string here or something.
      JSString *name = JS_NewStringCopyZ(cx, classp->name);
      if (name == NULL) {
        JS_ReportOutOfMemory(cx);
        return JS_FALSE;        
      }
      if (!JS_DefineProperty(cx, info, "nativeClass", STRING_TO_JSVAL(name),
                             NULL, NULL, JSPROP_ENUMERATE)) {
        JS_ReportOutOfMemory(cx);
        return JS_FALSE;
      }
    }

    if (!JS_DefineProperty(
          cx, info, "size",
          INT_TO_JSVAL(JS_GetObjectTotalSize(targetCx, target)),
          NULL, NULL, JSPROP_ENUMERATE)) {
      JS_ReportOutOfMemory(cx);
      return JS_FALSE;
    }

    if (!getChildrenInfo(cx, info, target, targetCx))
      return JS_FALSE;

    *rval = OBJECT_TO_JSVAL(info);
  } else
    *rval = JSVAL_NULL;

  return JS_TRUE;
}

static JSBool getGCRoots(JSContext *cx, JSObject *obj, uintN argc,
                         jsval *argv, jsval *rval)
{
  uint32 numRoots = JS_MapGCRoots(tracingState.runtime, countRoots, NULL);
  jsval rootList[numRoots];
  jsval *currIndex = rootList;

  JS_MapGCRoots(tracingState.runtime, rootMapFun, &currIndex);

  JSObject *roots = JS_NewArrayObject(cx, numRoots, rootList);
  if (roots == NULL) {
    JS_ReportError(cx, "Creating array failed.");
    return JS_FALSE;
  }

  *rval = OBJECT_TO_JSVAL(roots);
  return JS_TRUE;
}

static JSFunctionSpec server_global_functions[] = {
  JS_FS("print",          print,              1, 0, 0),
  JS_FS("ServerSocket",   createServerSocket, 0, 0, 0),
  JS_FS("getGCRoots",     getGCRoots,         0, 0, 0),
  JS_FS("getObjectInfo",  getObjInfo,         1, 0, 0),
  JS_FS_END
};

JSBool profileMemory(JSContext *cx, JSObject *obj, uintN argc,
                     jsval *argv, jsval *rval)
{
  if (!JS_DHashTableInit(&tracingState.visited, JS_DHashGetStubOps(),
                         NULL, sizeof(JSDHashEntryStub),
                         JS_DHASH_DEFAULT_CAPACITY(100))) {
    JS_ReportOutOfMemory(cx);
    return JS_FALSE;
  }

  tracingState.runtime = JS_GetRuntime(cx);
  tracingState.result = JS_TRUE;
  tracingState.tracer.context = cx;
  tracingState.tracer.callback = visitedBuilder;
  JS_TraceRuntime(&tracingState.tracer);

  if (!tracingState.result)
    return JS_FALSE;

  JSRuntime *serverRuntime = JS_NewRuntime(8L * 1024L * 1024L);
  if (serverRuntime == NULL) {
    JS_ReportError(cx, "Couldn't create server JS runtime.");
    return JS_FALSE;
  }

  JSContext *serverCx = JS_NewContext(serverRuntime, 8192);
  if (serverCx == NULL) {
    JS_ReportError(cx, "Couldn't create server JS context.");
    return JS_FALSE;
  }
  JS_SetOptions(serverCx, JSOPTION_VAROBJFIX);
  JS_SetVersion(serverCx, JSVERSION_LATEST);
  JS_SetErrorReporter(serverCx, reportError);

  JS_BeginRequest(serverCx);

  JSObject *serverGlobal = JS_NewObject(serverCx, &global_class, NULL, NULL);

  if (serverGlobal == NULL) {
    JS_ReportError(cx, "Couldn't create server JS global.");
    return JS_FALSE;
  }

  if (!JS_InitStandardClasses(serverCx, serverGlobal)) {
    JS_ReportError(cx, "Couldn't init standard classes on server JS global.");
    return JS_FALSE;
  }

  if (!JS_DefineFunctions(serverCx, serverGlobal, server_global_functions)) {
    JS_ReportError(cx, "Couldn't define functions on server JS global.");
    return JS_FALSE;
  }

  FILE *f = fopen(SERVER_FILENAME, "r");
  if (!f) {
    JS_ReportError(cx, "Couldn't open " SERVER_FILENAME);
    return JS_FALSE;
  }

  fseek(f, 0, SEEK_END);
  long fileSize = ftell(f);
  fseek(f, 0, SEEK_SET);
  
  char source[fileSize];
  fread(source, fileSize, 1, f);
  fclose(f);

  jsval serverRval;
  if (!JS_EvaluateScript(serverCx, serverGlobal, source,
                         fileSize, SERVER_FILENAME, 1,
                         &serverRval)) {
    JS_ReportError(cx, "Server failed.");
    return JS_FALSE;
  }

  /* Cleanup. */
  JS_DHashTableFinish(&tracingState.visited);
  JS_EndRequest(serverCx);
  JS_DestroyContext(serverCx);
  JS_DestroyRuntime(serverRuntime);

  *rval = JSVAL_VOID;
  return JS_TRUE;
}