
#ifdef WITH_RIVE_SCRIPTING
#include "rive/lua/rive_lua_libs.hpp"
#include "lualib.h"
#include <cmath>
#include <cstring>

using namespace rive;

static int vector_index(lua_State* L)
{
  // Key first: converting a numeric key can move the stack.
  size_t namelen = 0;
  const char* name = luaL_checklstring(L, 2, &namelen);

  const float* vec = luaL_checkvector(L, 1);

  if (namelen == 1)
  {
    switch (name[0])
    {
      case 'x':
      case '1':
        lua_pushnumber(L, vec[0]);
        return 1;
      case 'y':
      case '2':
        lua_pushnumber(L, vec[1]);
        return 1;
      case 'z':
      case '3':
        // Luau's vector type stores 3 components; .z is reachable
        // intrinsically through the VM's named-axis fastpath, but
        // numeric `[3]` indexing routes through this metamethod.
        lua_pushnumber(L, vec[2]);
        return 1;
      default:
        break;
    }
  }

  luaL_error(L, "'%s' is not a valid index of Vector", name);
  return 0;
}

// All magnitude and interpolation ops read all 3 components. For genuine 2D
// vectors z is 0, so results are identical to the old 2D math.
static inline float dot3(const float* a, const float* b)
{
  return a[0] * b[0] + a[1] * b[1] + a[2] * b[2];
}

static int vector_length(lua_State* L)
{
  const float* vec = luaL_checkvector(L, 1);
  lua_pushnumber(L, std::sqrt(dot3(vec, vec)));
  return 1;
}

static int vector_normalized(lua_State* L)
{
  const float* vec = luaL_checkvector(L, 1);
  float len2 = dot3(vec, vec);
  float scale = len2 > 0 ? 1.f / std::sqrt(len2) : 1.f;
  lua_pushvector(L, vec[0] * scale, vec[1] * scale, vec[2] * scale);
  return 1;
}

static int vector_lengthSquared(lua_State* L)
{
  const float* vec = luaL_checkvector(L, 1);
  lua_pushnumber(L, dot3(vec, vec));
  return 1;
}

static int vector_distanceSquared(lua_State* L)
{
  const float* lhs = luaL_checkvector(L, 1);
  const float* rhs = luaL_checkvector(L, 2);
  float d[3] = {rhs[0] - lhs[0], rhs[1] - lhs[1], rhs[2] - lhs[2]};
  lua_pushnumber(L, dot3(d, d));
  return 1;
}

static int vector_distance(lua_State* L)
{
  const float* lhs = luaL_checkvector(L, 1);
  const float* rhs = luaL_checkvector(L, 2);
  float d[3] = {rhs[0] - lhs[0], rhs[1] - lhs[1], rhs[2] - lhs[2]};
  lua_pushnumber(L, std::sqrt(dot3(d, d)));
  return 1;
}

static int vector_dot(lua_State* L)
{
  const float* lhs = luaL_checkvector(L, 1);
  const float* rhs = luaL_checkvector(L, 2);
  lua_pushnumber(L, dot3(lhs, rhs));
  return 1;
}

static int vector_cross(lua_State* L)
{
  auto lhs = lua_checkvec2d(L, 1);
  auto rhs = lua_checkvec2d(L, 2);
  lua_pushnumber(L, Vec2D::cross(*lhs, *rhs));
  return 1;
}

// 3D cross product. Distinct from `cross`, which returns the scalar 2D
// perp-dot.
static int vector_cross3(lua_State* L)
{
  const float* a = luaL_checkvector(L, 1);
  const float* b = luaL_checkvector(L, 2);
  lua_pushvector(
      L,
      a[1] * b[2] - a[2] * b[1],
      a[2] * b[0] - a[0] * b[2],
      a[0] * b[1] - a[1] * b[0]);
  return 1;
}

static int vector_scaleAndAdd(lua_State* L)
{
  const float* a = luaL_checkvector(L, 1);
  const float* b = luaL_checkvector(L, 2);
  float scale = float(luaL_checknumber(L, 3));
  lua_pushvector(
      L,
      a[0] + b[0] * scale,
      a[1] + b[1] * scale,
      a[2] + b[2] * scale);
  return 1;
}

static int vector_scaleAndSub(lua_State* L)
{
  const float* a = luaL_checkvector(L, 1);
  const float* b = luaL_checkvector(L, 2);
  float scale = float(luaL_checknumber(L, 3));
  lua_pushvector(
      L,
      a[0] - b[0] * scale,
      a[1] - b[1] * scale,
      a[2] - b[2] * scale);
  return 1;
}

// Exact at t=1, mirroring the VM fastcall's luai_lerpf so both dispatch
// paths return bit-identical results.
static inline float lerpf(float a, float b, float t)
{
  return t == 1.f ? b : a + (b - a) * t;
}

static int vector_lerp(lua_State* L)
{
  const float* a = luaL_checkvector(L, 1);
  const float* b = luaL_checkvector(L, 2);
  float t = float(luaL_checknumber(L, 3));
  lua_pushvector(
      L,
      lerpf(a[0], b[0], t),
      lerpf(a[1], b[1], t),
      lerpf(a[2], b[2], t));
  return 1;
}

// vec:writeToBuffer(buf, byteOffset) — 12 bytes (x, y, z as float32).
static int vector_writeToBuffer(lua_State* L)
{
  const float* vec = luaL_checkvector(L, 1);
  size_t bufLen = 0;
  void* buf = luaL_checkbuffer(L, 2, &bufLen);
  int off = int(luaL_checkinteger(L, 3));
  if (off < 0 || size_t(off) + 12 > bufLen)
  {
    luaL_error(L, "Vector:writeToBuffer offset out of range");
    return 0;
  }
  std::memcpy(static_cast<uint8_t*>(buf) + off, vec, 12);
  return 0;
}

// vec:writeVec4(buf, byteOffset, w) — 16 bytes (x, y, z, w as float32),
// matching a vec4 uniform-buffer slot.
static int vector_writeVec4(lua_State* L)
{
  const float* vec = luaL_checkvector(L, 1);
  size_t bufLen = 0;
  void* buf = luaL_checkbuffer(L, 2, &bufLen);
  int off = int(luaL_checkinteger(L, 3));
  float w = float(luaL_checknumber(L, 4));
  if (off < 0 || size_t(off) + 16 > bufLen)
  {
    luaL_error(L, "Vector:writeVec4 offset out of range");
    return 0;
  }
  uint8_t* dst = static_cast<uint8_t*>(buf) + off;
  std::memcpy(dst, vec, 12);
  std::memcpy(dst + 12, &w, 4);
  return 0;
}

static int vector_xy(lua_State* L)
{
  float x = (float)lua_tonumber(L, 1);
  float y = (float)lua_tonumber(L, 2);

  rive_lua_pushvector2(L, x, y);
  return 1;
}

static int vector_xyz(lua_State* L)
{
  float x = (float)lua_tonumber(L, 1);
  float y = (float)lua_tonumber(L, 2);
  float z = (float)lua_tonumber(L, 3);

  lua_pushvector(L, x, y, z);
  return 1;
}

static int vector_origin(lua_State* L)
{
  rive_lua_pushvector2(L, 0.0f, 0.0f);
  return 1;
}

static int vector_namecall(lua_State* L)
{
  int atom;
  const char* str = lua_namecallatom(L, &atom);
  if (str != nullptr)
  {
    switch (atom)
    {
      case (int)LuaAtoms::length:
        return vector_length(L);
      case (int)LuaAtoms::lengthSquared:
        return vector_lengthSquared(L);
      case (int)LuaAtoms::normalized:
        return vector_normalized(L);
      case (int)LuaAtoms::distance:
        return vector_distance(L);
      case (int)LuaAtoms::distanceSquared:
        return vector_distanceSquared(L);
      case (int)LuaAtoms::dot:
        return vector_dot(L);
      case (int)LuaAtoms::lerp:
        return vector_lerp(L);
      case (int)LuaAtoms::writeToBuffer:
        return vector_writeToBuffer(L);
      case (int)LuaAtoms::writeVec4:
        return vector_writeVec4(L);
    }
  }

  luaL_error(L, "%s is not a valid method of Vector", luaL_checkstring(L, 1));
  return 0;
}

static const luaL_Reg vectorStaticMethods[] = {
    {"distance", vector_distance},
    {"distanceSquared", vector_distanceSquared},
    {"dot", vector_dot},
    {"cross", vector_cross},
    {"cross3", vector_cross3},
    {"scaleAndAdd", vector_scaleAndAdd},
    {"scaleAndSub", vector_scaleAndSub},
    {"lerp", vector_lerp},
    {"xy", vector_xy},
    {"xyz", vector_xyz},
    {"origin", vector_origin},
    {"length", vector_length},
    {"lengthSquared", vector_lengthSquared},
    {"normalized", vector_normalized},
    {nullptr, nullptr}};

int luaopen_rive_vector(lua_State* L)
{
  luaL_register(L, "Vector", vectorStaticMethods);
  // create metatable for T
  lua_createtable(L, 0, 1);

  // Dupe the metatable on the stack
  lua_pushvalue(L, -1);
  // Add a vector
  lua_pushvector(L, 0.0f, 0.0f, 0.0f);
  // Move it before the metatable
  lua_insert(L, -2);
  // Set metatable on the vector (globally sets the metatable for all
  // vectors).
  lua_setmetatable(L, -2);
  // pop vector
  lua_pop(L, 1);

  lua_pushcfunction(L, vector_index, nullptr);
  lua_setfield(L, -2, "__index");

  lua_pushcfunction(L, vector_namecall, nullptr);
  lua_setfield(L, -2, "__namecall");

  lua_setreadonly(L, -1, true);
  lua_pop(L, 1); // pop the metatable

  return 1;
}

#endif