#include "jnj.h"
#include <janet.h>

#define UNUSED __attribute__((unused))

// Struct representing a value retrieved from the J engine.
typedef struct {
  int32_t type;
  I count;
  I rank;
  I *shape;
  V *data;
} j_value_t;

// Free allocated pointers in J Value struct
static void j_value_deinit(j_value_t *j_a) {
  janet_free(j_a->shape);
  janet_free(j_a->data);
}

// Callback when the Janet Abstract Type is GCed
static int janet_jvalue_gc(V *p, UNUSED size_t s) {
  j_value_t *j_a = (j_value_t *)p;
  j_value_deinit(j_a);
  return 0;
}

// Callback for Janet `get`
static int janet_jvalue_get(V *p, Janet key, Janet *out) {
  if (!janet_checktype(key, JANET_NUMBER))
    janet_panicf("Expected integer key for J Value");

  j_value_t *j_a = (j_value_t *)p;
  int32_t k = janet_unwrap_integer(key);

  if (k >= j_a->count) {
    return 0;
  } else {
    switch (j_a->type) {
    case JTBoolean: {
      C *j_data;
      j_data = (C *)j_a->data;
      *out = janet_wrap_integer(j_data[k]);
      return 1;
    }
    case JTInteger: {
      I *j_data;
      j_data = (I *)j_a->data;
      *out = janet_wrap_integer(j_data[k]);
      return 1;
    }
    case JTFloat: {
      D *j_data;
      j_data = (D *)j_a->data;
      *out = janet_wrap_number(j_data[k]);
      return 1;
    }
    case JTChar: {
      const uint8_t *j_data;
      j_data = (const uint8_t *)j_a->data;
      *out = janet_wrap_string(janet_string(j_data, j_a->count));
      return 1;
    }
    default:
      janet_panicf("Can't get value from J Value of type %n", j_a->type);
    }
  }
}

// Callback for Janet `marshal`
static void janet_jvalue_marshal(V *p, JanetMarshalContext *ctx) {
  j_value_t *j_a = (j_value_t *)p;

  janet_marshal_abstract(ctx, p);
  janet_marshal_int(ctx, j_a->type);
  janet_marshal_int64(ctx, j_a->count);
  janet_marshal_int64(ctx, j_a->rank);

  for (int i = 0; i < j_a->rank; i++) {
    janet_marshal_int64(ctx, j_a->shape[i]);
  }

  switch (j_a->type) {
  case JTBoolean:
    for (I i = 0; i < j_a->count; i++) {
      janet_marshal_janet(ctx, janet_wrap_integer(((C *)j_a->data)[i]));
    }
    return;
  case JTInteger:
    for (I i = 0; i < j_a->count; i++) {
      janet_marshal_janet(ctx, janet_wrap_integer(((I *)j_a->data)[i]));
    }
    return;
  case JTFloat:
    for (I i = 0; i < j_a->count; i++) {
      janet_marshal_janet(ctx, janet_wrap_number(((D *)j_a->data)[i]));
    }
    return;
  case JTChar:
    janet_marshal_janet(ctx, janet_wrap_string(janet_string(
                                 (const uint8_t *)j_a->data, j_a->count)));
    return;
  }
}

// Callback for Janet `unmarshal`
static V *janet_jvalue_unmarshal(JanetMarshalContext *ctx) {
  j_value_t *j_a = janet_unmarshal_abstract(ctx, sizeof(j_value_t));
  j_a->type = janet_unmarshal_int(ctx);
  j_a->count = janet_unmarshal_int64(ctx);
  j_a->rank = janet_unmarshal_int64(ctx);

  j_a->shape = janet_malloc(j_a->rank * sizeof(I));
  for (I i = 0; i < j_a->rank; i++) {
    j_a->shape[i] = janet_unmarshal_int64(ctx);
  }

  switch (j_a->type) {
  case JTBoolean: {
    uint8_t *j_data = janet_malloc(j_a->count * sizeof(uint8_t));
    for (I i = 0; i < j_a->count; i++) {
      j_data[i] = janet_unwrap_integer(janet_unmarshal_janet(ctx));
    }
    j_a->data = j_data;
    break;
  }
  case JTInteger: {
    I *j_data = janet_malloc(j_a->count * sizeof(I));
    for (I i = 0; i < j_a->count; i++) {
      j_data[i] = janet_unwrap_integer(janet_unmarshal_janet(ctx));
    }
    j_a->data = j_data;
    break;
  }
  case JTFloat: {
    D *j_data = janet_malloc(j_a->count * sizeof(D));
    for (I i = 0; i < j_a->count; i++) {
      j_data[i] = janet_unwrap_number(janet_unmarshal_janet(ctx));
    }
    j_a->data = j_data;
    break;
  }
  case JTChar: {
    uint8_t *j_data = janet_malloc(j_a->count * sizeof(uint8_t));
    memcpy(j_data, janet_unwrap_string(janet_unmarshal_janet(ctx)),
           j_a->count * sizeof(uint8_t));
    j_a->data = j_data;
    break;
  }
  }
  return j_a;
}

static C *type_tostring(I type) {
  switch (type) {
  case JTBoolean:
    return "boolean";
  case JTInteger:
    return "integer";
  case JTFloat:
    return "float";
  case JTChar:
    return "string";
  default:
    return "unknown";
  }
}

// Callback for Janet `describe`
static void janet_jvalue_tostring(V *p, JanetBuffer *buffer) {
  j_value_t *j_a = (j_value_t *)p;

  C str[32];
  if (j_a->type == JNJPassThrough) {
    sprintf(str, "J bytes");
  } else {
    sprintf(str, "%s [%llu/%llu]", type_tostring(j_a->type), j_a->rank,
            j_a->count);
  }
  janet_buffer_push_cstring(buffer, str);
}

static Janet first_idx(j_value_t *j_a) {
  if (j_a->count <= 0) {
    return janet_wrap_nil();
  } else {
    return janet_wrap_integer(0);
  }
}

// Implement 0-based indexing across `data`
static Janet janet_jvalue_next(V *p, Janet key) {
  j_value_t *j_a = (j_value_t *)p;
  if (janet_type(key) == JANET_NIL) {
    return first_idx(j_a);
  } else {
    int32_t max_count;
    // String arrays have a count field that refers to their number of
    // characters, but will return their value as a Janet String.
    if (j_a->type == JTChar) {
      max_count = 0;
    } else {
      max_count = j_a->count - 1;
    }

    int32_t k = janet_unwrap_integer(key);
    if (k >= max_count) {
      return janet_wrap_nil();
    } else {
      return janet_wrap_integer(k + 1);
    }
  }
}

// Implementation of j_value as Janet Abstract Type
const JanetAbstractType janet_jvalue_type = {"jnj/jvalue",
                                             janet_jvalue_gc,
                                             NULL,
                                             janet_jvalue_get,
                                             NULL,
                                             janet_jvalue_marshal,
                                             janet_jvalue_unmarshal,
                                             janet_jvalue_tostring,
                                             NULL,
                                             NULL,
                                             janet_jvalue_next,
                                             JANET_ATEND_NEXT};

static j_value_t *janet_new_abstract_jvalue() {
  return (j_value_t *)janet_abstract(&janet_jvalue_type, sizeof(j_value_t));
}

static j_value_t *janet_get_abstract_jvalue(Janet *argv, const uint8_t idx) {
  return (j_value_t *)janet_getabstract(argv, 0, &janet_jvalue_type);
}
// Given the pointers from a `getm` call, return a J value representing a J
// value.
static void populate_j_value(j_value_t *j_v, I jt, I jr, I *js, V *jd) {
  // Allocate an integer array to hold the shape. This will be freed in the gc
  // callback on the j_value it's put into.
  I *shape = janet_malloc(jr * sizeof(I));

  I values_count = 1;
  for (unsigned int d = 0; d < jr; d++) {
    values_count *= js[d];
    shape[d] = js[d];
  }

  j_v->type = jt;
  j_v->rank = jr;
  j_v->count = values_count;
  j_v->shape = shape;

  // Malloc the appropriate data; this will be freed when the j_value gets its
  // GC callback.
  switch (jt) {
  case JTBoolean:
  case JTChar: {
    uint8_t *values = janet_malloc(values_count * sizeof(uint8_t));

    for (I i = 0; i < values_count; i++) {
      uint8_t n = ((uint8_t *)jd)[i];
      values[i] = n;
    }
    j_v->data = values;
    break;
  }
  case JTInteger: {
    I *values = janet_malloc(values_count * sizeof(I));

    for (I i = 0; i < values_count; i++) {
      I n = ((I *)jd)[i];
      values[i] = n;
    }
    j_v->data = values;
    break;
  }
  case JTFloat: {
    D *values = janet_malloc(values_count * sizeof(double));

    for (I i = 0; i < values_count; i++) {
      double n = ((D *)jd)[i];
      values[i] = n;
    }
    j_v->data = values;
    break;
  }
  default: {
    janet_panicf("Can't create J Value of type %d", jt);
  }
  };
}

// Get a J variable and return as a J Value struct.
static j_value_t *janet_jgetm(V *j, C *name) {
  I jt, jr, js, jd;
  int r = jgetm(j, name, &jt, &jr, &js, &jd);
  if (r != 0) {
    janet_panicf("Received non-zero status code: %d", r);
  }

  j_value_t *j_v = janet_new_abstract_jvalue();
  populate_j_value(j_v, jt, jr, (I *)js, (V *)jd);
  return j_v;
}

// J Engine
typedef struct {
  V *engine;
} je_t;

// Free memory allocated by J dll
static void je_deinit(je_t *je) { jfree(je->engine); }

// Callback for GC on engine abstract type
static int janet_je_gc(V *p, UNUSED size_t s) {
  je_t *je = (je_t *)p;
  je_deinit(je);
  return 0;
}

// Callback for Janet `describe`
static void janet_je_tostring(V *p, JanetBuffer *buffer) {
  C description[128];
  je_t *je = (je_t *)p;
  sprintf(description, "J Engine %p", je->engine);
  janet_buffer_push_cstring(buffer, description);
}

// Implementation of J Engine struct as Janet Abstract Type
const JanetAbstractType janet_jengine_type = {
    "jnj/je",          janet_je_gc,         NULL, NULL, NULL, NULL, NULL,
    janet_je_tostring, JANET_ATEND_TOSTRING};

static je_t *janet_new_abstract_je() {
  return (je_t *)janet_abstract(&janet_jengine_type, sizeof(je_t));
}

static je_t *janet_get_abstract_je(Janet *argv, const uint8_t idx) {
  return (je_t *)janet_getabstract(argv, idx, &janet_jengine_type);
}

/// cfuns

// Construct an array of numbers from a JanetArray
static D *janet_number_array_to_jd(JanetArray *val) {
  D *jd = janet_malloc(sizeof(D) * val->count);
  for (I i = 0; i < val->count; i++) {
    if (!janet_checktype(val->data[i], JANET_NUMBER)) {
      janet_panicf("Expected a 1-dimensional number array");
    }
    jd[i] = janet_unwrap_number(val->data[i]);
  }
  return jd;
}

// Construct an array of numbers from a JanetTuple
static D *janet_number_tuple_to_jd(JanetTupleHead *val) {
  D *jd = janet_malloc(sizeof(D) * val->length);
  for (I i = 0; i < val->length; i++) {
    if (!janet_checktype(val->data[i], JANET_NUMBER)) {
      janet_panicf("Expected a 1-dimensional number tuple");
    }
    jd[i] = janet_unwrap_number(val->data[i]);
  }
  return jd;
}

// Set a variable on the J engine from a number
static int setm_num(J j, C *name, D val) {
  I jt = JTFloat, jr = 0, *jl = NULL;
  D *jd = &val;

  int r = jsetm(j, name, &jt, &jr, (I *)&jl, (I *)&jd);

  return r;
}

// Set a variable on the J engine from an array
static int setm_array(J j, C *name, JanetArray *val) {
  I jt = JTFloat, jr = 1, *jl = janet_malloc(sizeof(I));
  *jl = val->count;
  D *jd = janet_number_array_to_jd(val);

  int r = jsetm(j, name, &jt, &jr, (I *)&jl, (I *)&jd);

  free(jl);
  free(jd);
  return r;
}

// Set a variable on the J engine from a tuple
static int setm_tuple(J j, C *name, JanetTupleHead *val) {
  I jt = JTFloat, jr = 1, *jl = janet_malloc(sizeof(I));
  *jl = val->length;
  D *jd = janet_number_tuple_to_jd(val);

  int r = jsetm(j, name, &jt, &jr, (I *)&jl, (I *)&jd);

  free(jl);
  free(jd);
  return r;
}

// Set a variable on the J engine from a string
static int setm_string(J j, C *name, const uint8_t *string) {
  I jt = JTChar, jr = 1;
  I length = janet_string_length(string);
  I *jl = &length;

  return jsetm(j, name, &jt, &jr, (I *)&jl, (I *)&string);
}

// Set a variable on the J engine from a boolean
static int setm_bool(J j, C *name, int val) {
  I jt = JTBoolean, jr = 0, *jl = NULL;
  int *jd = &val;

  return jsetm(j, name, &jt, &jr, (I *)&jl, (I *)&jd);
}

// Set a variable on the J engine from an existing j-value
static int setm_jvalue(J j, C *name, j_value_t *val) {
  I jt = val->type, jr = val->rank;
  return jsetm(j, name, &jt, &jr, (I *)&val->shape, (I *)&val->data);
}

// Initialize a new J engine
static Janet cfun_jinit(int32_t argc, UNUSED Janet *argv) {
  janet_fixarity(argc, 0);

  if (!loaddl()) {
    janet_panic("Could not load J library");
  };

  J j = jinit();
  je_t *j_e = janet_new_abstract_je();
  j_e->engine = j;

  return janet_wrap_abstract(j_e);
}

// Execute a sentence on the Janet engine
static Janet cfun_jdo(int32_t argc, Janet *argv) {
  janet_fixarity(argc, 2);

  je_t *j_e = janet_get_abstract_je(argv, 0);
  const uint8_t *sentence = janet_getstring(argv, 1);

  int r = jdo(j_e->engine, (C *)sentence);
  return janet_wrap_number(r);
}

// Get a variable from the Janet engine
static Janet cfun_getm(int32_t argc, Janet *argv) {
  janet_fixarity(argc, 2);

  je_t *j_e = janet_get_abstract_je(argv, 0);
  const uint8_t *name = janet_getstring(argv, 1);

  j_value_t *res = janet_jgetm(j_e->engine, (C *)name);
  return janet_wrap_abstract(res);
}

// Set a variable on the J engine
static Janet cfun_setm(int32_t argc, Janet *argv) {
  janet_fixarity(argc, 3);

  je_t *j_e = janet_get_abstract_je(argv, 0);
  const uint8_t *name = janet_getstring(argv, 1);
  Janet val = argv[2];

  int r;
  JanetAbstract *abs;

  switch (janet_type(val)) {
  case JANET_ABSTRACT:
    abs = janet_unwrap_abstract(val);
    if (!(janet_abstract_type(abs) == &janet_jvalue_type)) {
      janet_panicf("Received unknown abstract type");
    }

    j_value_t *j_v = (j_value_t *)abs;
    r = setm_jvalue(j_e->engine, (C *)name, j_v);

    return janet_wrap_number(r);
  case JANET_BOOLEAN:
    r = setm_bool(j_e->engine, (C *)name, janet_unwrap_boolean(val));
    return janet_wrap_number(r);
  case JANET_NUMBER:
    r = setm_num(j_e->engine, (C *)name, janet_unwrap_number(val));
    return janet_wrap_number(r);
  case JANET_ARRAY:
    r = setm_array(j_e->engine, (C *)name, janet_unwrap_array(val));
    return janet_wrap_number(r);
  case JANET_TUPLE:
    r = setm_tuple(j_e->engine, (C *)name,
                   janet_tuple_head(janet_unwrap_tuple(val)));
    return janet_wrap_number(r);
  case JANET_STRING:
    r = setm_string(j_e->engine, (C *)name, janet_unwrap_string(val));
    return janet_wrap_number(r);
  default:
    janet_panicf("Can't set J variable with value of unknown type");
  }
}

// Get internal type of J Value as keyword
static Janet cfun_type(int32_t argc, Janet *argv) {
  janet_fixarity(argc, 1);

  j_value_t *j_a = janet_get_abstract_jvalue(argv, 0);

  if (j_a->type == JTBoolean) {
    const unsigned char str[] = "boolean";
    const uint8_t *keyw = janet_keyword(str, 7);
    return janet_wrap_keyword(keyw);
  }
  if (j_a->type == JTChar) {
    const unsigned char str[] = "string";
    const uint8_t *keyw = janet_keyword(str, 6);
    return janet_wrap_keyword(keyw);
  }
  if (j_a->type == JTInteger || j_a->type == JTFloat) {
    const unsigned char str[] = "number";
    const uint8_t *keyw = janet_keyword(str, 6);
    return janet_wrap_keyword(keyw);
  }
  return janet_wrap_nil();
}

// Get rank of J Value
static Janet cfun_rank(int32_t argc, Janet *argv) {
  janet_fixarity(argc, 1);

  j_value_t *j_a = janet_get_abstract_jvalue(argv, 0);
  return janet_wrap_number(j_a->rank);
}

// Get shape of J Value
static Janet cfun_shape(int32_t argc, Janet *argv) {
  janet_fixarity(argc, 1);

  j_value_t *j_a = janet_get_abstract_jvalue(argv, 0);
  JanetArray *shape = janet_array(j_a->rank);

  for (int i = 0; i < j_a->rank; i++) {
    Janet n = janet_wrap_number(j_a->shape[i]);
    janet_array_push(shape, n);
  }
  return janet_wrap_array(shape);
}

// Get J Value values as a single Janet array
static Janet cfun_values(int32_t argc, Janet *argv) {
  janet_fixarity(argc, 1);

  j_value_t *j_a = janet_get_abstract_jvalue(argv, 0);
  JanetArray *items = janet_array(j_a->count);

  if (j_a->rank == 0) {
    janet_panicf("Can't convert scalar result to an Array");
  }

  switch (j_a->type) {
  case JTBoolean: {
    const uint8_t *j_data;
    j_data = (uint8_t *)j_a->data;
    for (I i = 0; i < j_a->count; i++) {
      Janet n = janet_wrap_integer(j_data[i]);
      janet_array_push(items, n);
    }
    break;
  }
  case JTInteger: {
    I *j_data;
    j_data = (I *)j_a->data;
    for (I i = 0; i < j_a->count; i++) {
      Janet n = janet_wrap_integer(j_data[i]);
      janet_array_push(items, n);
    }
    break;
  }
  case JTFloat: {
    D *j_data;
    j_data = (D *)j_a->data;
    for (I i = 0; i < j_a->count; i++) {
      Janet n = janet_wrap_number(j_data[i]);
      janet_array_push(items, n);
    }
    break;
  }
  case JTChar: {
    const uint8_t *j_data;
    j_data = (const uint8_t *)j_a->data;
    return janet_wrap_string(janet_string(j_data, j_a->count));
  }
  default:
    janet_panicf("Can't get values from J Value of type %n", j_a->type);
  }

  return janet_wrap_array(items);
}

// Convert a Janet array along with shape data into a full J Value
static Janet cfun_tojvalue(int32_t argc, Janet *argv) {
  janet_fixarity(argc, 4);

  I jr = (I)janet_getnumber(argv, 0);
  I values_count = (I)janet_getnumber(argv, 1);

  j_value_t *j_a = janet_new_abstract_jvalue();
  j_a->type = JTFloat;
  j_a->rank = jr;
  j_a->count = values_count;

  JanetArray *shape_array = janet_getarray(argv, 2);
  I *shape = janet_malloc(jr * sizeof(I));
  for (unsigned int d = 0; d < jr; d++) {
    I shape_val = (I)janet_unwrap_number(shape_array->data[d]);
    shape[d] = shape_val;
  }
  j_a->shape = shape;

  j_a->data = janet_number_array_to_jd(janet_getarray(argv, 3));

  return janet_wrap_abstract(j_a);
}

///

static const JanetReg cfuns[] = {
    {"init", cfun_jinit, "(jnj/init)\n\nCreate a new J Engine instance."},
    {"do", cfun_jdo, "(jnj/do j-engine sentence)\n\nExecute a J sentence."},
    {"getm", cfun_getm, "(jnj/getm j-engine name)\n\nLookup a variable."},
    {"setm", cfun_setm, "(jnj/setm j-engine name)\n\nSet a variable."},
    {"type", cfun_type, "(jnj/type j-value)\n\nGet the type of a J value."},
    {"rank", cfun_rank, "(jnj/rank j-value)\n\nGet the rank of a J value."},
    {"shape", cfun_shape, "(jnj/shape j-value)\n\nGet the shape of a J value."},
    {"values", cfun_values,
     "(jnj/values j-value)\n\nGet all the items of a J value as a Janet "
     "array."},
    {"to-j-value", cfun_tojvalue,
     "(jnj/to-j-value rank count shape data)\n\nTurn the components of a Janet "
     "Array into J Value abstract type."},
    {NULL, NULL, NULL}};

JANET_MODULE_ENTRY(JanetTable *env) {
  janet_register_abstract_type(&janet_jvalue_type);
  janet_register_abstract_type(&janet_jengine_type);
  janet_cfuns(env, "jnj", cfuns);
}