#include "ui_lua.h"
#include "gpu.h"
#include "hsv.h"

#include "GLFW/glfw3.h"

#include <stddef.h>
#include <stdio.h>

#define UI_LUA_CONTAINER "ui_container"
#define UI_LUA_CONTEXT   "ui_context"
#define UI_LUA_GPU       "ui_gpu"

#define UI_LUA_ELEMENT_META "ui_element"

// Opaque element handle passed to scripts. container_id + generation let the
// binding detect stale handles (container unloaded / slot recycled) and raise
// a Lua error instead of touching a reused slot.
typedef struct UIElementHandleStruct {
  uint32_t slot;
  uint32_t generation;
  uint32_t is_string;
  uint32_t container_id;
} UIElementHandle;

static void* registry_pointer(lua_State* L, const char* key) {
  lua_getfield(L, LUA_REGISTRYINDEX, key);
  void* ptr = lua_touserdata(L, -1);
  lua_pop(L, 1);
  return ptr;
}

static void set_registry_pointers(lua_State* L, Container* c, UIContext* ui, RenderContext* gpu) {
  lua_pushlightuserdata(L, c);
  lua_setfield(L, LUA_REGISTRYINDEX, UI_LUA_CONTAINER);
  lua_pushlightuserdata(L, ui);
  lua_setfield(L, LUA_REGISTRYINDEX, UI_LUA_CONTEXT);
  lua_pushlightuserdata(L, gpu);
  lua_setfield(L, LUA_REGISTRYINDEX, UI_LUA_GPU);
}

// ---------- handle plumbing ----------

static UIElementHandle* check_handle(lua_State* L, int idx) {
  return luaL_checkudata(L, idx, UI_LUA_ELEMENT_META);
}

static Container* resolve_handle(lua_State* L, UIElementHandle* h) {
  UIContext* ui = registry_pointer(L, UI_LUA_CONTEXT);
  Container* c = context_container(h->container_id, ui);
  if(c == NULL) {
    luaL_error(L, "element handle's container %d is unloaded", h->container_id);
  }
  if(h->is_string) {
    if(h->slot >= c->num_strings
    || c->string_live[h->slot] == 0
    || c->string_generation[h->slot] != h->generation) {
      luaL_error(L, "stale string handle (slot %d)", h->slot);
    }
  } else {
    if(h->slot >= c->num_drawables
    || c->drawable_live[h->slot] == 0
    || c->drawable_generation[h->slot] != h->generation) {
      luaL_error(L, "stale element handle (slot %d)", h->slot);
    }
  }
  return c;
}

// Pushes the handle userdata and refs it against the slot so dispatch can
// pass the same object back to the script
static void push_new_handle(lua_State* L, Container* c, uint32_t slot, uint32_t is_string) {
  UIElementHandle* h = lua_newuserdata(L, sizeof(UIElementHandle));
  h->slot = slot;
  h->is_string = is_string;
  h->container_id = c->id;
  h->generation = is_string ? c->string_generation[slot] : c->drawable_generation[slot];
  luaL_setmetatable(L, UI_LUA_ELEMENT_META);

  lua_pushvalue(L, -1);
  int ref = luaL_ref(L, LUA_REGISTRYINDEX);
  if(is_string) {
    c->string_ref[slot] = ref;
  } else {
    c->drawable_ref[slot] = ref;
  }
}

static void push_element_handle(lua_State* L, Container* c, uint32_t element) {
  int ref = (element < c->cap_drawables) ? c->drawable_ref[element] : LUA_NOREF;
  if(ref == LUA_NOREF) {
    lua_pushnil(L);
  } else {
    lua_rawgeti(L, LUA_REGISTRYINDEX, ref);
  }
}

// ---------- declaration table parsing ----------

static float opt_number_field(lua_State* L, int table, const char* key, float def) {
  lua_getfield(L, table, key);
  float v = lua_isnil(L, -1) ? def : (float)luaL_checknumber(L, -1);
  lua_pop(L, 1);
  return v;
}

static uint32_t opt_integer_field(lua_State* L, int table, const char* key, uint32_t def) {
  lua_getfield(L, table, key);
  uint32_t v = lua_isnil(L, -1) ? def : (uint32_t)luaL_checkinteger(L, -1);
  lua_pop(L, 1);
  return v;
}

static void opt_vec2_field(lua_State* L, int table, const char* key, vec2 out) {
  out[0] = 0;
  out[1] = 0;
  lua_getfield(L, table, key);
  if(!lua_isnil(L, -1)) {
    luaL_checktype(L, -1, LUA_TTABLE);
    for(int i = 0; i < 2; i++) {
      lua_rawgeti(L, -1, i + 1);
      out[i] = (float)luaL_checknumber(L, -1);
      lua_pop(L, 1);
    }
  }
  lua_pop(L, 1);
}

// Reads {r, g, b, a} from the table at the top of the stack
static void read_vec4(lua_State* L, vec4 out) {
  luaL_checktype(L, -1, LUA_TTABLE);
  for(int i = 0; i < 4; i++) {
    lua_rawgeti(L, -1, i + 1);
    out[i] = (float)luaL_checknumber(L, -1);
    lua_pop(L, 1);
  }
}

// color={r,g,b,a} fills all corners; colors={{...},{...},{...},{...}} sets
// each corner
static void opt_corner_colors(lua_State* L, int table, vec4 corners[4]) {
  memset(corners, 0, sizeof(vec4)*4);

  lua_getfield(L, table, "color");
  if(!lua_isnil(L, -1)) {
    vec4 color;
    read_vec4(L, color);
    for(int i = 0; i < 4; i++) {
      memcpy(corners[i], color, sizeof(vec4));
    }
    lua_pop(L, 1);
    return;
  }
  lua_pop(L, 1);

  lua_getfield(L, table, "colors");
  if(!lua_isnil(L, -1)) {
    luaL_checktype(L, -1, LUA_TTABLE);
    for(int i = 0; i < 4; i++) {
      lua_rawgeti(L, -1, i + 1);
      read_vec4(L, corners[i]);
      lua_pop(L, 1);
    }
  }
  lua_pop(L, 1);
}

// ---------- element constructors ----------

// ui.rect{pos=, size=, z=, type=, var=, events=, color=|colors=} -> handle
static int lua_ui_rect(lua_State* L) {
  Container* container = registry_pointer(L, UI_LUA_CONTAINER);
  RenderContext* gpu = registry_pointer(L, UI_LUA_GPU);
  luaL_checktype(L, 1, LUA_TTABLE);

  GPUDrawable drawable;
  memset(&drawable, 0, sizeof(GPUDrawable));
  opt_vec2_field(L, 1, "pos", drawable.pos);
  opt_vec2_field(L, 1, "size", drawable.size);
  opt_corner_colors(L, 1, drawable.color);
  drawable.type = opt_integer_field(L, 1, "type", DRAWABLE_TYPE_RECT);
  drawable.var = opt_integer_field(L, 1, "var", 0);
  drawable.events = opt_integer_field(L, 1, "events", 0);
  drawable.z = opt_number_field(L, 1, "z", 0);

  uint32_t slot;
  if(ui_create_drawable(container, &drawable, gpu, &slot) != VK_SUCCESS) {
    return luaL_error(L, "ui.rect: failed to create drawable");
  }

  push_new_handle(L, container, slot, 0);
  return 1;
}

// ui.text{pos=, z=, size=, color=, font=, max_length=, text=} -> handle
static int lua_ui_text(lua_State* L) {
  Container* container = registry_pointer(L, UI_LUA_CONTAINER);
  UIContext* ui = registry_pointer(L, UI_LUA_CONTEXT);
  RenderContext* gpu = registry_pointer(L, UI_LUA_GPU);
  luaL_checktype(L, 1, LUA_TTABLE);

  GPUString string;
  memset(&string, 0, sizeof(GPUString));
  opt_vec2_field(L, 1, "pos", string.pos);
  string.size = opt_number_field(L, 1, "size", 16);
  string.font = opt_integer_field(L, 1, "font", 0);
  string.max_length = opt_integer_field(L, 1, "max_length", 0);
  string.z = opt_number_field(L, 1, "z", 0);

  string.color[0] = 1;
  string.color[1] = 1;
  string.color[2] = 1;
  string.color[3] = 1;
  lua_getfield(L, 1, "color");
  if(!lua_isnil(L, -1)) {
    read_vec4(L, string.color);
  }
  lua_pop(L, 1);

  uint32_t slot;
  if(ui_create_string(container, &string, gpu, &slot) != VK_SUCCESS) {
    return luaL_error(L, "ui.text: failed to create string");
  }

  lua_getfield(L, 1, "text");
  if(!lua_isnil(L, -1)) {
    update_ui_string(luaL_checkstring(L, -1), container, slot, ui, gpu);
  }
  lua_pop(L, 1);

  push_new_handle(L, container, slot, 1);
  return 1;
}

// ---------- handle methods ----------

// h:set_pos(x, y)
static int lua_element_set_pos(lua_State* L) {
  UIElementHandle* h = check_handle(L, 1);
  Container* c = resolve_handle(L, h);
  RenderContext* gpu = registry_pointer(L, UI_LUA_GPU);

  if(h->is_string) {
    GPUString* s = &c->strings_buffer[h->slot];
    s->pos[0] = luaL_checknumber(L, 2);
    s->pos[1] = luaL_checknumber(L, 3);
    add_transfers(
        &s->pos,
        c->strings,
        h->slot*sizeof(GPUString) + offsetof(GPUString, pos),
        sizeof(vec2),
        gpu);
  } else {
    GPUDrawable* d = &c->drawables_buffer[h->slot];
    d->pos[0] = luaL_checknumber(L, 2);
    d->pos[1] = luaL_checknumber(L, 3);
    add_transfers(
        &d->pos,
        c->drawables,
        h->slot*sizeof(GPUDrawable) + offsetof(GPUDrawable, pos),
        sizeof(vec2),
        gpu);
  }
  return 0;
}

// h:set_corner_color(corner, r, g, b, a)
static int lua_element_set_corner_color(lua_State* L) {
  UIElementHandle* h = check_handle(L, 1);
  Container* c = resolve_handle(L, h);
  RenderContext* gpu = registry_pointer(L, UI_LUA_GPU);
  if(h->is_string) {
    return luaL_error(L, "set_corner_color is not valid on text elements");
  }
  uint32_t corner = luaL_checkinteger(L, 2);
  if(corner > 3) {
    return luaL_error(L, "corner %d out of range", corner);
  }

  GPUDrawable* d = &c->drawables_buffer[h->slot];
  d->color[corner][0] = luaL_checknumber(L, 3);
  d->color[corner][1] = luaL_checknumber(L, 4);
  d->color[corner][2] = luaL_checknumber(L, 5);
  d->color[corner][3] = luaL_checknumber(L, 6);

  add_transfers(
      &d->color[corner],
      c->drawables,
      h->slot*sizeof(GPUDrawable) + offsetof(GPUDrawable, color) + corner*sizeof(vec4),
      sizeof(vec4),
      gpu);
  return 0;
}

// h:corner_color(corner) -> r, g, b, a
static int lua_element_corner_color(lua_State* L) {
  UIElementHandle* h = check_handle(L, 1);
  Container* c = resolve_handle(L, h);
  if(h->is_string) {
    return luaL_error(L, "corner_color is not valid on text elements");
  }
  uint32_t corner = luaL_checkinteger(L, 2);
  if(corner > 3) {
    return luaL_error(L, "corner %d out of range", corner);
  }

  GPUDrawable* d = &c->drawables_buffer[h->slot];
  lua_pushnumber(L, d->color[corner][0]);
  lua_pushnumber(L, d->color[corner][1]);
  lua_pushnumber(L, d->color[corner][2]);
  lua_pushnumber(L, d->color[corner][3]);
  return 4;
}

// h:set_text(text)
static int lua_element_set_text(lua_State* L) {
  UIElementHandle* h = check_handle(L, 1);
  Container* c = resolve_handle(L, h);
  UIContext* ui = registry_pointer(L, UI_LUA_CONTEXT);
  RenderContext* gpu = registry_pointer(L, UI_LUA_GPU);
  if(!h->is_string) {
    return luaL_error(L, "set_text is not valid on rect elements");
  }

  update_ui_string(luaL_checkstring(L, 2), c, h->slot, ui, gpu);
  return 0;
}

// h:set_z(z)
static int lua_element_set_z(lua_State* L) {
  UIElementHandle* h = check_handle(L, 1);
  Container* c = resolve_handle(L, h);
  RenderContext* gpu = registry_pointer(L, UI_LUA_GPU);
  float z = luaL_checknumber(L, 2);

  if(h->is_string) {
    ui_set_string_z(c, h->slot, z, gpu);
  } else {
    ui_set_drawable_z(c, h->slot, z, gpu);
  }
  return 0;
}

// h:focus()
static int lua_element_focus(lua_State* L) {
  UIElementHandle* h = check_handle(L, 1);
  Container* c = resolve_handle(L, h);
  UIContext* ui = registry_pointer(L, UI_LUA_CONTEXT);
  if(h->is_string) {
    return luaL_error(L, "focus is not valid on text elements");
  }

  set_active_element(c->id, h->slot, ui);
  return 0;
}

// h:is_focused() -> bool
static int lua_element_is_focused(lua_State* L) {
  UIElementHandle* h = check_handle(L, 1);
  Container* c = resolve_handle(L, h);
  UIContext* ui = registry_pointer(L, UI_LUA_CONTEXT);

  lua_pushboolean(L,
      h->is_string == 0
      && ui->active_container == c
      && ui->active_element == h->slot);
  return 1;
}

// h:remove()
static int lua_element_remove(lua_State* L) {
  UIElementHandle* h = check_handle(L, 1);
  Container* c = resolve_handle(L, h);
  RenderContext* gpu = registry_pointer(L, UI_LUA_GPU);

  if(h->is_string) {
    luaL_unref(L, LUA_REGISTRYINDEX, c->string_ref[h->slot]);
    c->string_ref[h->slot] = LUA_NOREF;
    ui_destroy_string(c, h->slot, gpu);
  } else {
    luaL_unref(L, LUA_REGISTRYINDEX, c->drawable_ref[h->slot]);
    c->drawable_ref[h->slot] = LUA_NOREF;
    ui_destroy_drawable(c, h->slot, gpu);
  }
  return 0;
}

// ---------- ui table functions ----------

// ui.blur()
static int lua_ui_blur(lua_State* L) {
  UIContext* ui = registry_pointer(L, UI_LUA_CONTEXT);
  RenderContext* gpu = registry_pointer(L, UI_LUA_GPU);

  clear_active_element(ui, gpu);
  return 0;
}

// ui.raise() — bring the calling script's container to the front
static int lua_ui_raise(lua_State* L) {
  Container* container = registry_pointer(L, UI_LUA_CONTAINER);
  UIContext* ui = registry_pointer(L, UI_LUA_CONTEXT);

  ui_container_to_front(container->id, ui);
  return 0;
}

// ui.hsv_to_rgb(h, s, v) -> r, g, b
static int lua_ui_hsv_to_rgb(lua_State* L) {
  double hsv[3] = {
    luaL_checknumber(L, 1),
    luaL_checknumber(L, 2),
    luaL_checknumber(L, 3),
  };
  vec3 rgb;
  hsv_to_rgb(hsv, rgb);
  lua_pushnumber(L, rgb[0]);
  lua_pushnumber(L, rgb[1]);
  lua_pushnumber(L, rgb[2]);
  return 3;
}

// ui.rgb_to_hsv(r, g, b) -> h, s, v
static int lua_ui_rgb_to_hsv(lua_State* L) {
  vec3 rgb = {
    luaL_checknumber(L, 1),
    luaL_checknumber(L, 2),
    luaL_checknumber(L, 3),
  };
  double hsv[3];
  rgb_to_hsv(rgb, hsv);
  lua_pushnumber(L, hsv[0]);
  lua_pushnumber(L, hsv[1]);
  lua_pushnumber(L, hsv[2]);
  return 3;
}

void ui_lua_register(lua_State* L) {
  static const luaL_Reg ui_funcs[] = {
    {"rect",       lua_ui_rect},
    {"text",       lua_ui_text},
    {"blur",       lua_ui_blur},
    {"raise",      lua_ui_raise},
    {"hsv_to_rgb", lua_ui_hsv_to_rgb},
    {"rgb_to_hsv", lua_ui_rgb_to_hsv},
    {NULL, NULL},
  };
  luaL_newlib(L, ui_funcs);
  lua_setglobal(L, "ui");

  static const luaL_Reg element_methods[] = {
    {"set_pos",          lua_element_set_pos},
    {"set_corner_color", lua_element_set_corner_color},
    {"corner_color",     lua_element_corner_color},
    {"set_text",         lua_element_set_text},
    {"set_z",            lua_element_set_z},
    {"focus",            lua_element_focus},
    {"is_focused",       lua_element_is_focused},
    {"remove",           lua_element_remove},
    {NULL, NULL},
  };
  luaL_newmetatable(L, UI_LUA_ELEMENT_META);
  luaL_newlib(L, element_methods);
  lua_setfield(L, -2, "__index");
  lua_pop(L, 1);

  lua_pushinteger(L, GLFW_PRESS);               lua_setglobal(L, "PRESS");
  lua_pushinteger(L, GLFW_RELEASE);             lua_setglobal(L, "RELEASE");
  lua_pushinteger(L, GLFW_MOUSE_BUTTON_LEFT);   lua_setglobal(L, "MOUSE_LEFT");
  lua_pushinteger(L, GLFW_MOUSE_BUTTON_RIGHT);  lua_setglobal(L, "MOUSE_RIGHT");
  lua_pushinteger(L, GLFW_MOUSE_BUTTON_MIDDLE); lua_setglobal(L, "MOUSE_MIDDLE");
  lua_pushinteger(L, GLFW_KEY_ESCAPE);          lua_setglobal(L, "KEY_ESCAPE");
  lua_pushinteger(L, GLFW_KEY_ENTER);           lua_setglobal(L, "KEY_ENTER");
  lua_pushinteger(L, GLFW_KEY_BACKSPACE);       lua_setglobal(L, "KEY_BACKSPACE");

  lua_pushinteger(L, DRAWABLE_TYPE_RECT);       lua_setglobal(L, "RECT");
  lua_pushinteger(L, DRAWABLE_TYPE_RECT_HSV);   lua_setglobal(L, "RECT_HSV");
  lua_pushinteger(L, DRAWABLE_TYPE_IMAGE);      lua_setglobal(L, "IMAGE");
  lua_pushinteger(L, UI_EVENT_BUTTON);          lua_setglobal(L, "EVENT_BUTTON");
  lua_pushinteger(L, UI_EVENT_SCROLL);          lua_setglobal(L, "EVENT_SCROLL");
  lua_pushinteger(L, UI_EVENT_CURSOR);          lua_setglobal(L, "EVENT_CURSOR");
}

VkResult ui_lua_load_container(
    lua_State* L,
    Container* c,
    UIContext* ui,
    RenderContext* gpu,
    const char* path) {

  if(luaL_loadfile(L, path) != LUA_OK) {
    fprintf(stderr, "ui_lua_load_container: %s\n", lua_tostring(L, -1));
    lua_pop(L, 1);
    return VK_ERROR_UNKNOWN;
  }

  // Sandboxed environment table: env = setmetatable({}, {__index = _G}),
  // set as the chunk's _ENV so script globals land in env, not _G
  lua_newtable(L);
  lua_newtable(L);
  lua_getglobal(L, "_G");
  lua_setfield(L, -2, "__index");
  lua_setmetatable(L, -2);

  lua_pushvalue(L, -1);
  lua_setupvalue(L, -3, 1);

  // Ref the env first so the chunk (now at the top) can be called
  c->script_env = luaL_ref(L, LUA_REGISTRYINDEX);

  // Top-level script code declares elements via ui.* during execution
  set_registry_pointers(L, c, ui, gpu);

  if(lua_pcall(L, 0, 0, 0) != LUA_OK) {
    fprintf(stderr, "ui_lua_load_container: %s\n", lua_tostring(L, -1));
    lua_pop(L, 1);
    luaL_unref(L, LUA_REGISTRYINDEX, c->script_env);
    c->script_env = 0;
    return VK_ERROR_UNKNOWN;
  }

  return VK_SUCCESS;
}

// Pushes the container's handler onto the stack, ready for args.
// Returns false if the container has no script or no such handler.
static bool ui_lua_begin_dispatch(
    lua_State* L,
    UIContext* ui,
    RenderContext* gpu,
    Container* c,
    const char* handler) {
  if(c == NULL || c->script_env == 0) return false;

  set_registry_pointers(L, c, ui, gpu);

  lua_rawgeti(L, LUA_REGISTRYINDEX, c->script_env);
  lua_getfield(L, -1, handler);
  lua_remove(L, -2);
  if(!lua_isfunction(L, -1)) {
    lua_pop(L, 1);
    return false;
  }
  return true;
}

// Calls the handler with nargs already on the stack, returns its boolean result
static bool ui_lua_finish_dispatch(lua_State* L, const char* handler, int nargs) {
  if(lua_pcall(L, nargs, 1, 0) != LUA_OK) {
    fprintf(stderr, "%s: %s\n", handler, lua_tostring(L, -1));
    lua_pop(L, 1);
    return false;
  }
  bool consumed = lua_toboolean(L, -1);
  lua_pop(L, 1);
  return consumed;
}

VkResult ui_lua_call(
    UIContext* ui,
    RenderContext* gpu,
    Container* c,
    const char* function,
    const char* argument) {
  lua_State* L = ui->lua;
  if(!ui_lua_begin_dispatch(L, ui, gpu, c, function)) {
    return VK_ERROR_UNKNOWN;
  }
  lua_pushstring(L, argument);
  if(lua_pcall(L, 1, 0, 0) != LUA_OK) {
    fprintf(stderr, "%s: %s\n", function, lua_tostring(L, -1));
    lua_pop(L, 1);
    return VK_ERROR_UNKNOWN;
  }
  return VK_SUCCESS;
}

bool ui_lua_dispatch_button(
    UIContext* ui,
    RenderContext* gpu,
    Container* c,
    uint32_t element,
    float x,
    float y,
    int button,
    int action,
    int mods) {
  lua_State* L = ui->lua;
  if(!ui_lua_begin_dispatch(L, ui, gpu, c, "on_button")) return false;
  push_element_handle(L, c, element);
  lua_pushnumber(L, x);
  lua_pushnumber(L, y);
  lua_pushinteger(L, button);
  lua_pushinteger(L, action);
  lua_pushinteger(L, mods);
  return ui_lua_finish_dispatch(L, "on_button", 6);
}

bool ui_lua_dispatch_cursor(
    UIContext* ui,
    RenderContext* gpu,
    Container* c,
    uint32_t element,
    float x,
    float y) {
  lua_State* L = ui->lua;
  if(!ui_lua_begin_dispatch(L, ui, gpu, c, "on_cursor")) return false;
  push_element_handle(L, c, element);
  lua_pushnumber(L, x);
  lua_pushnumber(L, y);
  return ui_lua_finish_dispatch(L, "on_cursor", 3);
}

bool ui_lua_dispatch_scroll(
    UIContext* ui,
    RenderContext* gpu,
    Container* c,
    uint32_t element,
    double x,
    double y) {
  lua_State* L = ui->lua;
  if(!ui_lua_begin_dispatch(L, ui, gpu, c, "on_scroll")) return false;
  push_element_handle(L, c, element);
  lua_pushnumber(L, x);
  lua_pushnumber(L, y);
  return ui_lua_finish_dispatch(L, "on_scroll", 3);
}

bool ui_lua_dispatch_key(
    UIContext* ui,
    RenderContext* gpu,
    Container* c,
    uint32_t element,
    int key,
    int action,
    int mods) {
  lua_State* L = ui->lua;
  if(!ui_lua_begin_dispatch(L, ui, gpu, c, "on_key")) return false;
  push_element_handle(L, c, element);
  lua_pushinteger(L, key);
  lua_pushinteger(L, action);
  lua_pushinteger(L, mods);
  return ui_lua_finish_dispatch(L, "on_key", 4);
}

bool ui_lua_dispatch_text(
    UIContext* ui,
    RenderContext* gpu,
    Container* c,
    uint32_t element,
    unsigned int codepoint) {
  lua_State* L = ui->lua;
  if(!ui_lua_begin_dispatch(L, ui, gpu, c, "on_text")) return false;
  push_element_handle(L, c, element);
  lua_pushinteger(L, codepoint);
  return ui_lua_finish_dispatch(L, "on_text", 2);
}

bool ui_lua_dispatch_deselect(
    UIContext* ui,
    RenderContext* gpu,
    Container* c,
    uint32_t element) {
  lua_State* L = ui->lua;
  if(!ui_lua_begin_dispatch(L, ui, gpu, c, "on_deselect")) return false;
  push_element_handle(L, c, element);
  return ui_lua_finish_dispatch(L, "on_deselect", 1);
}