#include "cell.h"
#include "prosperon.h"
#include "stb_ds.h"
#include <string.h>
#include <stdio.h>
#include "HandmadeMath.h"

#define VEC2_FMT "[%g,%g]"
#define VEC2_MEMS(s) (s).x, (s).y

#define TR_FMT "(pos) " HMMFMT_VEC3 " scale " HMMFMT_VEC3
#define TR_PRINT(tr) HMMPRINT_VEC3(tr->pos), HMMPRINT_VEC3(tr->scale)

typedef struct transform {
  HMM_Vec3 pos;
  HMM_Vec3 scale;
  HMM_Quat rotation;
  HMM_Mat4 cache;
  HMM_Mat4 gcache;
  int dirty;
  struct transform *parent;
  struct transform **children;
  
  JSValue self;  
  JSValue jsparent;  
  JSValue *jschildren;
  JSValue change_hook;
} transform;


static transform **dirties;

static transform model = {
  .pos = {0,0,0},
  .scale = {1,1,1},
  .rotation = {0,0,0,1},
  .cache = {1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1},
  .gcache = {1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1},
  .dirty = 0,
};

transform *make_transform()
{
  transform *t = malloc(sizeof(transform));
  *t = model;
  t->jsparent = JS_NULL;
  t->change_hook = JS_NULL;
  return t;
}

void transform_free(JSRuntime *rt, transform *t) {
//  JS_FreeValueRT(rt,t->self);
  JS_FreeValueRT(rt,t->change_hook);
  JS_FreeValueRT(rt,t->jsparent);
  for (int i = 0; i < arrlen(t->jschildren); i++)
    JS_FreeValueRT(rt,t->jschildren[i]);

  arrfree(t->jschildren);
  arrfree(t->children);
  free(t);

  for (int i = arrlen(dirties)-1; i >= 0; i--)
    if (dirties[i] == t) arrdelswap(dirties, i);
}

void transform_clean(JSContext *js, transform *t)
{
  if (!t->dirty) return;
  t->dirty = 0;
  t->cache = HMM_M4TRS(t->pos,t->rotation,t->scale);
  if (t->parent)
    t->gcache = HMM_MulM4(t->parent->gcache, t->cache);
  else
    t->gcache = t->cache;

  for (int i = 0; i < arrlen(t->children); i++) {
    t->children[i]->dirty = 1;
    transform_clean(js, t->children[i]);
  }

  if (!JS_IsNull(t->change_hook)) {
    JSValue ret = JS_Call(js, t->change_hook, JS_DupValue(js,t->self), 0, NULL);
    JS_FreeValue(js,ret);
  }
}

void clean_all(JSContext *js)
{
  for (int i = 0; i < arrlen(dirties); i++)
    transform_clean(js, dirties[i]);

  arrsetlen(dirties,0);
}

void transform_apply(transform *t)
{
  t->dirty = 1;
  arrput(dirties,t);
}

void transform_move(transform *t, HMM_Vec3 v)
{
  t->pos = HMM_AddV3(t->pos, v);
  transform_apply(t);
}

HMM_Vec3 transform_direction(transform *t, HMM_Vec3 dir)
{
  return HMM_QVRot(dir, t->rotation);
}

HMM_Vec3 trans_forward(const transform *const trans) { return HMM_QVRot(vFWD, trans->rotation); }
HMM_Vec3 trans_back(const transform *trans) { return HMM_QVRot(vBKWD, trans->rotation); }
HMM_Vec3 trans_up(const transform *trans) { return HMM_QVRot(vUP, trans->rotation); }
HMM_Vec3 trans_down(const transform *trans) { return HMM_QVRot(vDOWN, trans->rotation); }
HMM_Vec3 trans_right(const transform *trans) { return HMM_QVRot(vRIGHT, trans->rotation); }
HMM_Vec3 trans_left(const transform *trans) { return HMM_QVRot(vLEFT, trans->rotation); }

HMM_Vec2 mat_t_pos(HMM_Mat3 m, HMM_Vec2 pos) { return HMM_MulM3V3(m, (HMM_Vec3){pos.x, pos.y, 1}).xy; }

HMM_Vec2 mat_t_dir(HMM_Mat3 m, HMM_Vec2 dir)
{
  m.Columns[2] = (HMM_Vec3){0,0,1};
  return HMM_MulM3V3(m, (HMM_Vec3){dir.x, dir.y, 1}).XY;
}

HMM_Vec2 mat_up(HMM_Mat3 m) { return HMM_NormV2(m.Columns[1].XY); }
HMM_Vec2 mat_right(HMM_Mat3 m) { return HMM_NormV2(m.Columns[0].XY); }
float vec_angle(HMM_Vec2 a, HMM_Vec2 b) { return acos(HMM_DotV2(a,b)/(HMM_LenV2(a)*HMM_LenV2(b))); }
float vec_dirangle(HMM_Vec2 a, HMM_Vec2 b) { return atan2(b.x, b.y) - atan2(a.x, a.y); }

HMM_Vec3 mat3_t_pos(HMM_Mat4 m, HMM_Vec3 pos) { return HMM_MulM4V4(m, (HMM_Vec4){pos.X, pos.Y, pos.Z, 1}).XYZ; }

HMM_Vec3 mat3_t_dir(HMM_Mat4 m, HMM_Vec3 dir)
{
  m.Columns[3] = (HMM_Vec4){0,0,0,1};
  return mat3_t_pos(m, dir);
}

HMM_Mat4 transform2mat(transform *t)
{
  return t->cache;
}

HMM_Mat3 transform2mat3(transform *t)
{
  return HMM_M3TRS(t->pos.xy, t->rotation.x, t->scale.xy);
}

HMM_Mat4 transform2mat4_global(transform *t)
{
  return t->gcache;
}

rect transform2rect(transform *t)
{
  HMM_Mat4 m = t->gcache;
  rect r = {0};
  r.x = m.Columns[3].x;
  r.y = m.Columns[3].y;
  r.w = m.Columns[0].x;
  r.h = m.Columns[1].y;
  return r;
}

HMM_Quat angle2rotation(float angle)
{
  return HMM_QFromAxisAngle_RH(vBKWD, angle);
}

transform mat2transform(HMM_Mat4 m)
{
  transform t;
  t.pos = m.Columns[3].xyz;
  for (int i = 0; i < 2; i++)
    t.scale.Elements[i] = HMM_LenV3(m.Columns[i].xyz);
//  for (int i = 0; i < 2; i++)
//    m.Columns[i].xyz = HMM_MulV3(m.Columns[i].xyz, t.scale.Elements[i]);
  t.rotation = HMM_M4ToQ_RH(m);
  return t;
}

// Transform class definitions
JSClassID js_transform_id;
static void js_transform_mark(JSRuntime *rt, JSValueConst val, JS_MarkFunc *mark_func) {
  transform *t = JS_GetOpaque(val, js_transform_id);
  if (!t) return;
  // Mark the JSValue references stored in your struct
  JS_MarkValue(rt, t->change_hook, mark_func);
  JS_MarkValue(rt, t->jsparent, mark_func);
  // Mark the array elements
  for (int i = 0; i < arrlen(t->jschildren); i++)
    JS_MarkValue(rt, t->jschildren[i], mark_func);
}

QJSCLASSMARK_EXTERN(transform,)

static JSValue js_transform_get_change_hook(JSContext *js, JSValueConst self)
{
  transform *t = js2transform(js,self);
  return JS_DupValue(js,t->change_hook);
}

static JSValue js_transform_set_change_hook(JSContext *js, JSValueConst self, JSValue v)
{
  transform *t = js2transform(js,self);
  if (!JS_IsNull(v) && !JS_IsFunction(v)) return JS_ThrowReferenceError(js, "Hook must be a function.");
  JS_FreeValue(js,t->change_hook);
  t->change_hook = JS_DupValue(js,v);
  return JS_NULL;
}

static JSValue js_transform_get_parent(JSContext *js, JSValueConst self)
{
  transform *t = js2transform(js,self);
  if (t->parent) return JS_DupValue(js,t->jsparent);
  return JS_NULL;
}

static JSValue js_transform_set_parent(JSContext *js, JSValueConst self, JSValue v)
{
  transform *p = js2transform(js,v);
  if (!JS_IsNull(v) && !p)
    return JS_ThrowReferenceError(js,"Parent must be another transform.");

  transform *t = js2transform(js,self);

  if (t == p)
    return JS_ThrowReferenceError(js, "A transform cannot be its own parent.");

  if (t->parent) {
    transform *cur_parent = t->parent;
    JS_FreeValue(js,t->jsparent);
    t->jsparent = JS_NULL;

    for (int i = 0; i < arrlen(cur_parent->children); i++) {
      if (cur_parent->children[i] == t) {
        arrdelswap(cur_parent->children,i);
        break;
      }
    }

    for (int i = 0; i < arrlen(cur_parent->jschildren); i++) {
      if (JS_SameValue(js,cur_parent->jschildren[i],self)) {
        JS_FreeValue(js,cur_parent->jschildren[i]);
        arrdelswap(cur_parent->jschildren,i);
        break;
      }
    }
  }

  t->parent = p;
  t->jsparent = JS_DupValue(js,v);

  if (p) {
    arrput(p->children, t);
    JSValue child = JS_DupValue(js,self);
    arrput(p->jschildren,child);
  }

  transform_apply(t);
    
  return JS_NULL;
}

// TRANSFORM ACTION FUNCTIONS

JSC_CCALL(transform_move,
  transform *t = js2transform(js,self);
  transform_move(t, js2vec3(js,argv[0]));
)

JSC_CCALL(transform_lookat,
  HMM_Vec3 point = js2vec3(js,argv[0]);
  transform *go = js2transform(js,self);
  HMM_Mat4 m = HMM_LookAt_RH(go->pos, point, vUP);
  go->rotation = HMM_M4ToQ_RH(m);
  transform_apply(go);
)

JSC_CCALL(transform_rotate,
  HMM_Vec3 axis = js2vec3(js,argv[0]);
  transform *t = js2transform(js,self);
  HMM_Quat rot = HMM_QFromAxisAngle_RH(axis, js2angle(js,argv[1]));
  t->rotation = HMM_MulQ(t->rotation,rot);
  transform_apply(t);
)

JSC_CCALL(transform_angle,
  HMM_Vec3 axis = js2vec3(js,argv[0]);
  transform *t = js2transform(js,self);
  if (axis.x) return angle2js(js,HMM_Q_Roll(t->rotation));
  if (axis.y) return angle2js(js,HMM_Q_Pitch(t->rotation));
  if (axis.z) return angle2js(js,HMM_Q_Yaw(t->rotation));
  return angle2js(js,0);
)

JSC_CCALL(transform_direction,
  transform *t = js2transform(js,self);
  return vec32js(js, HMM_QVRot(js2vec3(js,argv[0]), t->rotation));
)

JSC_CCALL(transform_phys2d,
  transform *t = js2transform(js,self);
  HMM_Vec2 v = js2vec2(js,argv[0]);
  float av = js2number(js,argv[1]);
  float dt = js2number(js,argv[2]);
  transform_move(t, (HMM_Vec3){v.x*dt,v.y*dt,0});
  HMM_Quat rot = HMM_QFromAxisAngle_RH((HMM_Vec3){0,0,1}, av*dt);
  t->rotation = HMM_MulQ(t->rotation, rot);
  transform_apply(t);
)

JSC_CCALL(transform_unit,
  transform *t = js2transform(js,self);
  t->pos = v3zero;
  t->rotation = QUAT1;
  t->scale = v3one;
  transform_apply(t);
)

JSC_CCALL(transform_trs,
  transform *t = js2transform(js,self);
  t->pos = JS_IsNull(argv[0]) ? v3zero : js2vec3(js,argv[0]);
  t->rotation = JS_IsNull(argv[1]) ? QUAT1 : js2quat(js,argv[1]);
  t->scale = JS_IsNull(argv[2]) ? v3one : js2vec3(js,argv[2]);
  transform_apply(t);
)

JSC_CCALL(transform_rect,
  transform *t = js2transform(js,self);
  rect r = js2rect(js,argv[0]);
  t->pos = (HMM_Vec3){r.x,r.y,0};
  t->scale = (HMM_Vec3){r.w,r.h,1};
  t->rotation = QUAT1;
  transform_apply(t);
)

JSC_CCALL(transform_array,
  transform *t = js2transform(js,self);
  HMM_Mat4 m= transform2mat(t);
  ret = JS_NewArray(js);
  for (int i = 0; i < 16; i++)
    JS_SetPropertyUint32(js,ret,i, number2js(js,m.em[i]));
)

JSC_CCALL(transform_torect,
  transform *t = js2transform(js,self);
  return rect2js(js,transform2rect(t));
)

JSC_CCALL(transform_children,
  transform *t = js2transform(js,self);
  ret = JS_NewArray(js);
  for (int i = 0; i < arrlen(t->jschildren); i++)
    JS_SetPropertyUint32(js,ret,i,JS_DupValue(js,t->jschildren[i]));
)

static const JSCFunctionListEntry js_transform_funcs[] = {
//  CGETSET_ADD(transform, pos),
//  CGETSET_ADD(transform, scale),
//  CGETSET_ADD(transform, rotation),
//  CGETSET_ADD(transform, parent),
//  CGETSET_ADD(transform, change_hook),
  MIST_FUNC_DEF(transform, trs, 3),
  MIST_FUNC_DEF(transform, phys2d, 3),
  MIST_FUNC_DEF(transform, move, 1),
  MIST_FUNC_DEF(transform, rotate, 2),
  MIST_FUNC_DEF(transform, angle, 1),
  MIST_FUNC_DEF(transform, lookat, 1),
  MIST_FUNC_DEF(transform, direction, 1),
  MIST_FUNC_DEF(transform, unit, 0),
  MIST_FUNC_DEF(transform, rect, 1),
  MIST_FUNC_DEF(transform, array, 0),
  MIST_FUNC_DEF(transform, torect, 0),
  MIST_FUNC_DEF(transform, children, 0),
};

// Constructor function for transform
static JSValue js_transform_constructor(JSContext *js, JSValueConst new_target, int argc, JSValueConst *argv) {
  transform *t = make_transform();
  if (!t) return JS_ThrowOutOfMemory(js);
  
  JSValue ret = transform2js(js, t);
  t->self = ret;
  
  return ret;
}

JSC_CCALL(transform_clean,
  clean_all(js);
)

static const JSCFunctionListEntry js_transform_ctor_funcs[] = {
  MIST_FUNC_DEF(transform, clean, 0),
};

CELL_USE_INIT(
  QJSCLASSPREP_FUNCS(transform);
  
  JSValue ctor = JS_NewCFunction2(js, js_transform_constructor, "transform", 0, JS_CFUNC_generic, 0);
  
  JS_SetPropertyFunctionList(js, ctor, js_transform_ctor_funcs, countof(js_transform_ctor_funcs));
  
  return ctor;
)