// 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>
#include <time.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_style_mats[7] = {};

void *viewmodel_material();

const char *chams_material_name(int idx) {
  switch (idx) {
  case 1: return "debug/debugwireframe";
  case 2:
  case 3:
  case 4:
  case 5:
  case 6: return "debug/debugtranslucentsinglecolor";
  default: return "vgui/white"; // models/debug/debugwhite is EP2-only
  }
}

void *style_material(int idx) {
  if (idx < 0) idx = 0;
  if (idx > 6) idx = 6;
  if (g_style_mats[idx])
    return g_style_mats[idx];
  if (g_style_mats[3])
    return g_style_mats[3];
  return viewmodel_material();
}

// Viewmodels need a model shader; vgui materials do not draw skinned models.
void *viewmodel_material() {
  if (g_style_mats[3])
    return g_style_mats[3];
  static void *mat = nullptr;
  if (!mat)
    mat = world_find_material("debug/debugtranslucentsinglecolor");
  return 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, int style, bool force_white,
                 const int *rgb) {
  using FlagFn = void (*)(void *, int, bool);
  float r = force_white ? 1.0f : (float)rgb[0] / 255.0f;
  float g = force_white ? 1.0f : (float)rgb[1] / 255.0f;
  float b = force_white ? 1.0f : (float)rgb[2] / 255.0f;
  chams_color(mat, r, g, b);
  // 0 Flat, 1 Wireframe, 2 Additive, 3 Solid, 4 White, 5 Fullbright,
  // 6 Glow outline (solid base; the wireframe overlay is a second draw).
  vfunc<FlagFn>(mat, 29)(mat, 1 << 15 /* IGNOREZ */, ignorez);
  vfunc<FlagFn>(mat, 29)(mat, 1 << 28 /* WIREFRAME */, style == 1);
  vfunc<FlagFn>(mat, 29)(mat, 1 << 7 /* ADDITIVE */, style == 2);
  vfunc<FlagFn>(mat, 29)(mat, 1 << 12 /* FLAT */,
                         style == 0 || style == 5);
  vfunc<FlagFn>(mat, 29)(mat, 1 << 6 /* SELFILLUM */, style == 5);
}

void chams_reset(void *mat) {
  using FlagFn = void (*)(void *, int, bool);
  vfunc<FlagFn>(mat, 29)(mat, 1 << 15, false);
  vfunc<FlagFn>(mat, 29)(mat, 1 << 28, false);
  vfunc<FlagFn>(mat, 29)(mat, 1 << 7, false);
  vfunc<FlagFn>(mat, 29)(mat, 1 << 12, false);
  vfunc<FlagFn>(mat, 29)(mat, 1 << 6, false);
  chams_color(mat, 1.0f, 1.0f, 1.0f);
}

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_vm = g_chams_viewmodel && esp_is_viewmodel_index(idx);
      ChamsGroup *grp = nullptr;
      if (!is_vm) {
        int rel = esp_npc_relation(idx);
        if (rel >= 0)
          grp = rel == 0 ? &g_chams_hostile
                         : rel == 1 ? &g_chams_neutral : &g_chams_friendly;
      }
      if (is_vm || grp) {
        auto draw_pass = [&](int style, const int *rgb,
                             bool ignorez) -> bool {
          if (style < 0) style = 0;
          if (style > 6) style = 6;
          void *mat = style_material(style);
          if (!mat)
            return false;
          bool white = style == 4;
          chams_apply(mat, ignorez, style, white, rgb);
          g_force_override(self, mat, 0 /* OVERRIDE_NORMAL */);
          g_orig_model_exec(self, state, info, bone);
          if (style == 6) {
            // Solid base plus a wireframe contour pass for depth.
            void *wire = g_style_mats[1];
            if (wire && wire != mat) {
              chams_apply(wire, ignorez, 1, false, rgb);
              g_force_override(self, wire, 0);
              g_orig_model_exec(self, state, info, bone);
              chams_reset(wire);
            }
          }
          chams_reset(mat);
          return true;
        };
        bool did = false;
        if (is_vm) {
          did = draw_pass(g_chams_vm_mat, g_chams_vm_rgb, false);
        } else {
          if (grp->inv_on)
            did |= draw_pass(grp->inv_mat, grp->inv_rgb, true);
          if (grp->vis_on)
            did |= draw_pass(grp->vis_mat, grp->vis_rgb, false);
        }
        if (did) {
          g_force_override(self, nullptr, 0);
          return;
        }
      }
    }
  }
  g_orig_model_exec(self, state, info, bone);
}

} // namespace

static double chams_now() {
  timespec ts;
  clock_gettime(CLOCK_MONOTONIC, &ts);
  return (double)ts.tv_sec + (double)ts.tv_nsec * 1e-9;
}

// Called from outer_draw (outside the model draw) so FindMaterial can load
// the style VMTs. The render hook only reads the cached pointers.
void chams_preload() {
  static double last = 0.0;
  static bool probed = false;
  double now = chams_now();
  if (now - last < 1.5)
    return;
  last = now;
  if (!probed) {
    probed = true;
    void *probe = world_find_material("vgui/white");
    logf("chams: probe vgui/white = %p", probe);
  }
  for (int i = 0; i < 7; i++) {
    if (g_style_mats[i])
      continue;
    void *m = world_find_material(chams_material_name(i));
    if (m) {
      g_style_mats[i] = m;
      logf("chams: preloaded style %d %s = %p", i, chams_material_name(i), m);
    }
  }
}

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);
}