// studio.cpp — colour static props at the StudioRender level.
//
// Dynamic props are tinted through their client entity's m_clrRender (see
// esp.cpp). Map-static props have no entity: the engine's static-prop manager
// calls IStudioRender::DrawModelStaticProp, and the colour modulation used for
// the draw comes from IStudioRender::SetColorModulation. Hook both on the
// interface object so the prop colour is substituted only while a prop model
// is being drawn.
//
// VStudioRender025 vtable (verified against studiorender.so's
// CStudioRenderContext vtable at .data.rel.ro):
//   27 = SetColorModulation(float const*)
//   28 = SetAlphaModulation(float)
//   30 = DrawModelStaticProp(const DrawModelInfo_t&, const matrix3x4_t&, int)
// (The 2013 SDK header lists these five slots earlier; this build's layout
// has additional methods before them.)

#include "hl2h.h"
#include <cstdio>
#include <cstring>
#include <dlfcn.h>

namespace {

using SetColorFn = void (*)(void *, const float *);
using SetAlphaFn = void (*)(void *, float);
using DrawStaticFn = void (*)(void *, void *, void *, int);
// DrawModel takes (this, results, info, boneToWorld, flexWeights,
// flexDelayedWeights, modelOrigin, flags). The modelOrigin argument is a
// `const Vector &`, so it arrives as a pointer; omitting it shifts flags into
// modelOrigin and leaves flags as caller stack garbage.
using DrawModelFn = void (*)(void *, void *, void *, void *, void *, void *,
                             const void *, int);
using DrawArrayFn = void (*)(void *, void *, int, void *, int, int);

SetColorFn g_orig_set_color = nullptr;
SetAlphaFn g_orig_set_alpha = nullptr;
DrawStaticFn g_orig_draw_static = nullptr;
DrawModelFn g_orig_draw_model = nullptr;
DrawArrayFn g_orig_draw_array = nullptr;

bool g_install_tried = false;
bool g_installed = false;
int g_prop_depth = 0;
int g_static_calls = 0;
int g_static_props = 0;

bool prop_tint_on() {
  return feature_active(FEAT_WORLD_TINT) && g_tint_prop_on;
}

void prop_color(float out[4]) {
  out[0] = (float)g_tint_prop_rgb[0] / 255.0f;
  out[1] = (float)g_tint_prop_rgb[1] / 255.0f;
  out[2] = (float)g_tint_prop_rgb[2] / 255.0f;
  out[3] = 1.0f;
}

// DrawModelInfo_t starts with studiohdr_t*; studiohdr_t::name is at +12
// (id, version, checksum, then char name[64]).
const char *info_model_name(void *info) {
  if (!info)
    return nullptr;
  void *hdr = *(void **)info;
  if (!hdr)
    return nullptr;
  const char *name = (const char *)hdr + 12;
  size_t n = strnlen(name, 128);
  if (n < 7 || n >= 128)
    return nullptr;
  for (size_t i = 0; i < n; i++) {
    unsigned char c = (unsigned char)name[i];
    if (c < 0x20 || c > 0x7e)
      return nullptr;
  }
  return name;
}

bool is_prop_model(const char *name) {
  if (!name)
    return false;
  // studiohdr_t::name is the short model name with Windows separators, e.g.
  // "props_junk\wood_crate002a.mdl" (the "models/" prefix is stripped).
  if (strncmp(name, "models/", 7) == 0 || strncmp(name, "models\\", 7) == 0)
    name += 7;
  if (strncmp(name, "props_skybox", 12) == 0 ||
      strncmp(name, "props\\skybox", 13) == 0)
    return false; // distant 3D-skybox props
  return strncmp(name, "props/", 6) == 0 ||
         strncmp(name, "props\\", 6) == 0 ||
         strncmp(name, "props_", 6) == 0;
}

void hook_set_color(void *self, const float *color) {
  if (g_prop_depth > 0 && prop_tint_on()) {
    float c[4];
    prop_color(c);
    g_orig_set_color(self, c);
    return;
  }
  g_orig_set_color(self, color);
}

void hook_set_alpha(void *self, float alpha) {
  if (g_prop_depth > 0 && prop_tint_on()) {
    g_orig_set_alpha(self, 1.0f);
    return;
  }
  g_orig_set_alpha(self, alpha);
}


void hook_draw_model(void *self, void *results, void *info, void *bone,
                     void *flex, void *flex_delayed, const void *model_origin,
                     int flags) {
  const char *model = info_model_name(info);
  bool prop = prop_tint_on() && is_prop_model(model);
  if (prop) {
    float c[4];
    prop_color(c);
    g_orig_set_color(self, c);
    g_prop_depth++;
  }
  g_orig_draw_model(self, results, info, bone, flex, flex_delayed,
                    model_origin, flags);
  if (prop)
    g_prop_depth--;
}

void hook_draw_array(void *self, void *info, int count, void *instances,
                     int stride, int flags) {
  const char *model = info_model_name(info);
  bool prop = prop_tint_on() && is_prop_model(model);
  if (prop)
    g_prop_depth++;
  g_orig_draw_array(self, info, count, instances, stride, flags);
  if (prop)
    g_prop_depth--;
}

void hook_draw_static_prop(void *self, void *info, void *model_to_world,
                           int flags) {
  const char *model = info_model_name(info);
  bool prop = prop_tint_on() && is_prop_model(model);
  g_static_calls++;
  if (prop)
    g_static_props++;
  if (g_static_calls <= 3)
    logf("studio: static draw model=\"%s\" prop=%d", model ? model : "?",
         (int)prop);
  if (g_static_calls == 500)
    logf("studio: static draws=%d prop_draws=%d", g_static_calls,
         g_static_props);
  if (prop)
    g_prop_depth++;
  g_orig_draw_static(self, info, model_to_world, flags);
  if (prop)
    g_prop_depth--;
}

std::string loaded_module_path(const char *needle) {
  FILE *maps = fopen("/proc/self/maps", "r");
  if (!maps)
    return {};
  char line[1024];
  std::string found;
  while (fgets(line, sizeof(line), maps)) {
    char *p = strchr(line, '/');
    if (!p)
      continue;
    size_t pl = strlen(p);
    while (pl && (p[pl - 1] == '\n' || p[pl - 1] == '\r'))
      p[--pl] = 0;
    if (!strstr(p, needle) || !strstr(p, ".so"))
      continue;
    found = p;
    if (!strstr(p, "(deleted)"))
      break;
  }
  fclose(maps);
  return found;
}

// ---------------------------------------------------------------------------
// chams: force a flat material on NPC draws
// ---------------------------------------------------------------------------
//
// IVModelRender::DrawModelExecute is called once per model. For NPC entities
// we force models/debug/debugwhite, optionally with MATERIAL_VAR_IGNOREZ for
// the through-wall pass, and draw twice. IVModelRender::ForcedMaterialOverride
// is slot 1; DrawModelExecute is slot 19 in VEngineModel016.

using ModelExecFn = void (*)(void *, const void *, const void *, void *);
using ForceOverrideFn = void (*)(void *, void *, int);

void *g_modelrender = nullptr;
ModelExecFn g_orig_model_exec = nullptr;
ForceOverrideFn g_force_override = nullptr;
void *g_chams_mat = nullptr;
int g_chams_mat_idx = -1;

const char *chams_material_name(int idx) {
  switch (idx) {
  case 1: return "debug/debugwireframevertexcolorignorez";
  case 2: return "vgui/white_additive";
  case 3: return "debug/debugtranslucentsinglecolor";
  case 4: return "vgui/white";
  default: return "models/debug/debugwhite";
  }
}

void *chams_material() {
  int idx = g_chams_mat_index;
  if (idx < 0) idx = 0;
  if (idx > 4) idx = 4;
  if (g_chams_mat && g_chams_mat_idx == idx)
    return g_chams_mat;
  g_chams_mat = world_find_material(chams_material_name(idx));
  if (!g_chams_mat)
    g_chams_mat = world_find_material("models/debug/debugwhite");
  g_chams_mat_idx = idx;
  logf("chams: style %d -> %p", idx, g_chams_mat);
  return g_chams_mat;
}

void chams_color(void *mat, float r, float g, float b) {
  using ColorFn = void (*)(void *, float, float, float);
  using FindVarFn = void *(*)(void *, const char *, bool *, bool);
  using SetVec3Fn = void (*)(void *, float, float, float);
  bool found = false;
  void *var = vfunc<FindVarFn>(mat, 11)(mat, "$color", &found, false);
  if (var)
    vfunc<SetVec3Fn>(var, 11)(var, r, g, b);
  found = false;
  var = vfunc<FindVarFn>(mat, 11)(mat, "$color2", &found, false);
  if (var)
    vfunc<SetVec3Fn>(var, 11)(var, r, g, b);
  vfunc<ColorFn>(mat, 28)(mat, r, g, b);
}

void chams_apply(void *mat, bool ignorez) {
  using FlagFn = void (*)(void *, int, bool);
  chams_color(mat, (float)g_chams_rgb[0] / 255.0f,
              (float)g_chams_rgb[1] / 255.0f,
              (float)g_chams_rgb[2] / 255.0f);
  vfunc<FlagFn>(mat, 29)(mat, 1 << 15 /* MATERIAL_VAR_IGNOREZ */, ignorez);
}

void hkdModelExecute(void *self, const void *state, const void *info,
                     void *bone) {
  if (!g_orig_model_exec)
    return;
  if (feature_active(FEAT_CHAMS) && info && g_force_override) {
    // ModelRenderInfo_t::entity_index is at 0x30 on 32-bit.
    int idx = *(const int *)((const char *)info + 0x30);
    if (idx > 0 && idx <= 8192) {
      bool is_npc = esp_is_npc_index(idx);
      bool is_vm = g_chams_viewmodel && esp_is_viewmodel_index(idx);
      if (is_npc || is_vm) {
        void *mat = chams_material();
        if (mat) {
          if (g_chams_walls && !is_vm) {
            chams_apply(mat, true);
            g_force_override(self, mat, 0 /* OVERRIDE_NORMAL */);
            g_orig_model_exec(self, state, info, bone);
          }
          chams_apply(mat, false);
          g_force_override(self, mat, 0);
          g_orig_model_exec(self, state, info, bone);
          g_force_override(self, nullptr, 0);
          // debugwhite is shared: restore colour and depth behaviour.
          using FlagFn = void (*)(void *, int, bool);
          vfunc<FlagFn>(mat, 29)(mat, 1 << 15, false);
          chams_color(mat, 1.0f, 1.0f, 1.0f);
          return;
        }
      }
    }
  }
  g_orig_model_exec(self, state, info, bone);
}

} // namespace

void chams_init() {
  if (g_modelrender || !g_iface.modelrender) {
    if (!g_iface.modelrender)
      logf("chams: VEngineModel016 unavailable");
    return;
  }
  void **vt = *(void ***)g_iface.modelrender;
  if (!vt) {
    logf("chams: model render has no vtable");
    return;
  }
  g_force_override = (ForceOverrideFn)vt[1];
  g_orig_model_exec = (ModelExecFn)vt[19];
  if (!g_engine.contains((uintptr_t)g_force_override) ||
      !g_engine.contains((uintptr_t)g_orig_model_exec)) {
    logf("chams: implausible vtable entries override=%p draw=%p",
         (void *)g_force_override, (void *)g_orig_model_exec);
    return;
  }
  static void *nvt[32];
  memcpy(nvt, vt, sizeof(nvt));
  nvt[19] = (void *)&hkdModelExecute;
  g_modelrender = g_iface.modelrender;
  *(void ***)g_modelrender = nvt;
  logf("chams: IVModelRender %p hooked (draw=%p override=%p)", g_modelrender,
       (void *)g_orig_model_exec, (void *)g_force_override);
}

void studio_props_init() {
  if (g_install_tried)
    return;
  g_install_tried = true;

  std::string path = loaded_module_path("studiorender");
  void *h = nullptr;
  if (!path.empty())
    h = dlopen(path.c_str(), RTLD_NOLOAD | RTLD_LAZY | RTLD_GLOBAL);
  if (!h)
    h = dlopen("studiorender.so", RTLD_NOLOAD | RTLD_LAZY | RTLD_GLOBAL);
  if (!h) {
    logf("studio: studiorender.so not loaded");
    return;
  }
  auto cf = (CreateInterfaceFn)dlsym(h, "CreateInterface");
  if (!cf) {
    logf("studio: studiorender CreateInterface missing");
    return;
  }
  void *sr = cf("VStudioRender025", nullptr);
  if (!sr) {
    logf("studio: VStudioRender025 not exposed");
    return;
  }

  void **old_vt = *(void ***)sr;
  if (!old_vt) {
    logf("studio: interface has no vtable");
    return;
  }
  static void *vt[49];
  memcpy(vt, old_vt, sizeof(vt));

  g_orig_set_color = (SetColorFn)vt[27];
  g_orig_set_alpha = (SetAlphaFn)vt[28];
  g_orig_draw_model = (DrawModelFn)vt[29];
  g_orig_draw_static = (DrawStaticFn)vt[30];
  g_orig_draw_array = (DrawArrayFn)vt[46];
  vt[27] = (void *)&hook_set_color;
  vt[28] = (void *)&hook_set_alpha;
  vt[29] = (void *)&hook_draw_model;
  vt[30] = (void *)&hook_draw_static_prop;
  vt[46] = (void *)&hook_draw_array;

  *(void ***)sr = vt;
  g_installed = true;
  logf("studio: IStudioRender %p hooked (color=%p draw=%p static=%p "
       "array=%p)", sr, (void *)g_orig_set_color, (void *)g_orig_draw_model,
       (void *)g_orig_draw_static, (void *)g_orig_draw_array);
}