#include "mesh_overlay.h"

#include <limits.h>
#include <math.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

/* ------------------------------------------------------------------------- */
/* Settings                                                                   */
/* ------------------------------------------------------------------------- */

MeshOverlaySettings mesh_overlay_default_settings(void) {
  static const unsigned char defaults[MESH_OVERLAY_CATEGORY_COUNT][3] = {
      {0x7F, 0x84, 0x9C}, /* ESCAPE     Catppuccin Mocha overlay1 */
      {0xCB, 0xA6, 0xF7}, /* DARK       mauve */
      {0xF9, 0xE2, 0xAF}, /* UNRESOLVED yellow */
      {0xF3, 0x8B, 0xA8}, /* INCOMPLETE red */
      {0x89, 0xB4, 0xFA}, /* UNTRACED   blue */
  };
  MeshOverlaySettings settings;
  memcpy(settings.colors, defaults, sizeof settings.colors);
  settings.opacity = 0.5;
  return settings;
}

static int overlay_hex_nibble(char digit, unsigned char *value) {
  if (digit >= '0' && digit <= '9') {
    *value = (unsigned char)(digit - '0');
    return 0;
  }
  if (digit >= 'a' && digit <= 'f') {
    *value = (unsigned char)(digit - 'a' + 10);
    return 0;
  }
  if (digit >= 'A' && digit <= 'F') {
    *value = (unsigned char)(digit - 'A' + 10);
    return 0;
  }
  return -1;
}

int mesh_overlay_parse_color(const char *text, unsigned char rgb[3]) {
  if (text == NULL || rgb == NULL)
    return -1;
  if (strlen(text) != 7 || text[0] != '#')
    return -1;
  unsigned char parsed[3];
  for (int channel = 0; channel < 3; ++channel) {
    unsigned char high, low;
    if (overlay_hex_nibble(text[1 + 2 * channel], &high) ||
        overlay_hex_nibble(text[2 + 2 * channel], &low))
      return -1;
    parsed[channel] = (unsigned char)((high << 4) | low);
  }
  rgb[0] = parsed[0];
  rgb[1] = parsed[1];
  rgb[2] = parsed[2];
  return 0;
}

/* ------------------------------------------------------------------------- */
/* Edge extraction                                                            */
/* ------------------------------------------------------------------------- */

typedef struct {
  size_t low;
  size_t high;
} OverlayEdge;

static int overlay_edge_compare(const void *lhs, const void *rhs) {
  const OverlayEdge *a = lhs;
  const OverlayEdge *b = rhs;
  if (a->low != b->low)
    return a->low < b->low ? -1 : 1;
  if (a->high != b->high)
    return a->high < b->high ? -1 : 1;
  return 0;
}

static unsigned char overlay_vertex_category(const LensVertex *vertex) {
  if (!vertex->traced)
    return (unsigned char)MESH_OVERLAY_CATEGORY_UNTRACED;
  switch (vertex->outcome) {
  case RAY_OUTCOME_ESCAPED:
    return (unsigned char)MESH_OVERLAY_CATEGORY_ESCAPE;
  case RAY_OUTCOME_DARK:
    return (unsigned char)MESH_OVERLAY_CATEGORY_DARK;
  case RAY_OUTCOME_UNRESOLVED:
    return (unsigned char)MESH_OVERLAY_CATEGORY_UNRESOLVED;
  case RAY_OUTCOME_INCOMPLETE:
  default:
    return (unsigned char)MESH_OVERLAY_CATEGORY_INCOMPLETE;
  }
}

int mesh_overlay_prepare(const FrameLensMesh *mesh, MeshOverlayLines *lines) {
  if (lines == NULL)
    return -1;
  lines->lines = NULL;
  lines->count = 0;
  if (mesh == NULL)
    return -1;
  /* An overflowing triangle count is rejected before any pointer is
   * dereferenced so a corrupt mesh cannot drive an out-of-bounds read. */
  if (mesh->triangle_count > SIZE_MAX / 3)
    return -1;
  if (mesh->triangle_count == 0)
    return 0;
  const size_t raw_count = mesh->triangle_count * 3;
  if (raw_count > SIZE_MAX / sizeof(OverlayEdge))
    return -1;
  if (mesh->triangles == NULL || mesh->vertices == NULL)
    return -1;

  OverlayEdge *raw = malloc(raw_count * sizeof *raw);
  if (raw == NULL)
    return -1;
  for (size_t triangle = 0; triangle < mesh->triangle_count; ++triangle) {
    const LensTriangle *leaf = &mesh->triangles[triangle];
    for (int edge = 0; edge < 3; ++edge) {
      const size_t from = leaf->vertex[edge];
      const size_t to = leaf->vertex[(edge + 1) % 3];
      if (from >= mesh->vertex_count || to >= mesh->vertex_count) {
        free(raw);
        return -1;
      }
      const LensVertex *a = &mesh->vertices[from];
      const LensVertex *b = &mesh->vertices[to];
      if (!isfinite(a->image_x) || !isfinite(a->image_y) ||
          !isfinite(b->image_x) || !isfinite(b->image_y)) {
        free(raw);
        return -1;
      }
      OverlayEdge *slot = &raw[3 * triangle + (size_t)edge];
      slot->low = from < to ? from : to;
      slot->high = from < to ? to : from;
    }
  }

  qsort(raw, raw_count, sizeof *raw, overlay_edge_compare);
  size_t unique = 0;
  for (size_t i = 0; i < raw_count; ++i) {
    if (unique == 0 || raw[unique - 1].low != raw[i].low ||
        raw[unique - 1].high != raw[i].high)
      raw[unique++] = raw[i];
  }
  if (unique > SIZE_MAX / sizeof(MeshOverlayLine)) {
    free(raw);
    return -1;
  }
  MeshOverlayLine *out = NULL;
  if (unique != 0) {
    out = malloc(unique * sizeof *out);
    if (out == NULL) {
      free(raw);
      return -1;
    }
  }
  for (size_t i = 0; i < unique; ++i) {
    const LensVertex *a = &mesh->vertices[raw[i].low];
    const LensVertex *b = &mesh->vertices[raw[i].high];
    out[i].x0 = a->image_x;
    out[i].y0 = a->image_y;
    out[i].x1 = b->image_x;
    out[i].y1 = b->image_y;
    out[i].category0 = overlay_vertex_category(a);
    out[i].category1 = overlay_vertex_category(b);
  }
  free(raw);
  lines->lines = out;
  lines->count = unique;
  return 0;
}

void mesh_overlay_lines_destroy(MeshOverlayLines *lines) {
  if (lines == NULL)
    return;
  free(lines->lines);
  lines->lines = NULL;
  lines->count = 0;
}

/* ------------------------------------------------------------------------- */
/* Rasterization                                                              */
/* ------------------------------------------------------------------------- */

static double overlay_fractional_part(double value) { return value - floor(value); }

static void blend_overlay(unsigned char *pixels, int width, int height, int x,
                           int y, const unsigned char rgb[3], double alpha,
                           int rgba) {
  if (alpha <= 0.0 || x < 0 || x >= width || y < 0 || y >= height)
    return;
  if (alpha > 1.0)
    alpha = 1.0;
  const int stride = rgba ? 4 : 3;
  unsigned char *pixel = pixels + stride * ((size_t)y * (size_t)width + (size_t)x);
  for (int channel = 0; channel < stride; ++channel) {
    const double source = channel == 3 ? 255.0 : (double)rgb[channel];
    const double mixed =
        (double)pixel[channel] * (1.0 - alpha) + source * alpha;
    long value = lround(mixed);
    if (value < 0)
      value = 0;
    if (value > 255)
      value = 255;
    pixel[channel] = (unsigned char)value;
  }
}

/* Liang-Barsky clip of the (major, minor) segment to the inclusive box.  Keeps
 * every subsequent cast and loop bounded even for huge finite coordinates.
 * Returns 1 when a nonempty clipped segment remains, 0 when fully outside.  The
 * clipped outputs are guaranteed finite and inside the box before the caller
 * casts them: a nonfinite interpolation result (cancellation) is skipped, and a
 * finite roundoff overshoot is clamped back into the box.  The segment direction
 * keeps x0 <= x1 and the box clamp is monotone, so the order is preserved. */
static int clip_overlay_segment(double *x0, double *y0, double *x1, double *y1,
                                double xmin, double xmax, double ymin,
                                double ymax) {
  const double dx = *x1 - *x0;
  const double dy = *y1 - *y0;
  if (!isfinite(dx) || !isfinite(dy))
    return 0;
  double t0 = 0.0, t1 = 1.0;
  const double p[4] = {-dx, dx, -dy, dy};
  const double q[4] = {*x0 - xmin, xmax - *x0, *y0 - ymin, ymax - *y0};
  for (int i = 0; i < 4; ++i) {
    if (p[i] == 0.0) {
      if (q[i] < 0.0)
        return 0;
    } else {
      const double r = q[i] / p[i];
      if (p[i] < 0.0) {
        if (r > t1)
          return 0;
        if (r > t0)
          t0 = r;
      } else {
        if (r < t0)
          return 0;
        if (r < t1)
          t1 = r;
      }
    }
  }
  double nx0 = *x0 + t0 * dx;
  double ny0 = *y0 + t0 * dy;
  double nx1 = *x0 + t1 * dx;
  double ny1 = *y0 + t1 * dy;
  if (!isfinite(nx0) || !isfinite(ny0) || !isfinite(nx1) || !isfinite(ny1))
    return 0;
  if (nx0 < xmin)
    nx0 = xmin;
  if (nx0 > xmax)
    nx0 = xmax;
  if (ny0 < ymin)
    ny0 = ymin;
  if (ny0 > ymax)
    ny0 = ymax;
  if (nx1 < xmin)
    nx1 = xmin;
  if (nx1 > xmax)
    nx1 = xmax;
  if (ny1 < ymin)
    ny1 = ymin;
  if (ny1 > ymax)
    ny1 = ymax;
  *x0 = nx0;
  *y0 = ny0;
  *x1 = nx1;
  *y1 = ny1;
  return 1;
}

static void plot_overlay_aa(unsigned char *pixels, int width, int height,
                             int steep, int x, int y, double coverage,
                             const unsigned char rgb[3], double opacity,
                             int rgba) {
  if (coverage > 0.0)
    blend_overlay(pixels, width, height, steep ? y : x, steep ? x : y, rgb,
                   coverage * opacity, rgba);
}

/* Xiaolin Wu line rasterization, one pixel wide, with a color that switches to
 * the second endpoint category at the major-axis midpoint.  A single pass
 * colors the whole edge, so a midpoint pixel is never blended from both halves. */
static void draw_overlay_line(unsigned char *pixels, int width, int height,
                               const MeshOverlayLine *line,
                               const MeshOverlaySettings *settings, int rgba) {
  double x0 = line->x0, y0 = line->y0;
  double x1 = line->x1, y1 = line->y1;
  const unsigned char *first = settings->colors[line->category0];
  const unsigned char *second = settings->colors[line->category1];
  const int steep = fabs(y1 - y0) > fabs(x1 - x0);
  if (steep) {
    double swap = x0;
    x0 = y0;
    y0 = swap;
    swap = x1;
    x1 = y1;
    y1 = swap;
  }
  if (x0 > x1) {
    double swap = x0;
    x0 = x1;
    x1 = swap;
    swap = y0;
    y0 = y1;
    y1 = swap;
    const unsigned char *color_swap = first;
    first = second;
    second = color_swap;
  }
  /* The switch is fixed to the true midpoint of the unclipped edge.  Halving
   * each endpoint separately cannot overflow for finite same-sign endpoints. */
  const double midpoint = 0.5 * x0 + 0.5 * x1;
  const int major_limit = steep ? height : width;
  const int minor_limit = steep ? width : height;
  double cx0 = x0, cy0 = y0, cx1 = x1, cy1 = y1;
  if (!clip_overlay_segment(&cx0, &cy0, &cx1, &cy1, -1.0, (double)major_limit,
                            -1.0, (double)minor_limit))
    return;
  x0 = cx0;
  y0 = cy0;
  x1 = cx1;
  y1 = cy1;
  const double dx = x1 - x0;
  /* A zero-length edge has no coverage; do not turn it into a vertex dot. */
  if (!(dx > 0.0))
    return;
  const double gradient = (y1 - y0) / dx;
  /* After the steep/orientation normalization |gradient| <= 1, so it is finite
   * for a finite nonzero dx; this guard keeps a pathological subnormal dx from
   * ever reaching a float-to-int cast. */
  if (!isfinite(gradient))
    return;
  const int first_column = (int)round(x0);
  const int last_column = (int)round(x1);
  if (first_column == last_column) {
    /* Wu's two endpoint formulas overlap in the same column for a subpixel
     * segment. Paint its length-weighted coverage once, rather than applying
     * two alpha blends that make tiny edges brighter than full-length ones. */
    const double center_y = 0.5 * y0 + 0.5 * y1;
    const int row = (int)floor(center_y);
    const double fraction = overlay_fractional_part(center_y);
    const unsigned char *color = first_column < midpoint ? first : second;
    plot_overlay_aa(pixels, width, height, steep, first_column, row,
                    dx * (1.0 - fraction), color, settings->opacity, rgba);
    plot_overlay_aa(pixels, width, height, steep, first_column, row + 1,
                    dx * fraction, color, settings->opacity, rgba);
    return;
  }
  double x_end = (double)first_column;
  double y_end = y0 + gradient * (x_end - x0);
  if (!isfinite(y_end))
    return;
  double x_gap = 1.0 - overlay_fractional_part(x0 + 0.5);
  const int x_pixel_start = (int)x_end;
  int y_pixel = (int)floor(y_end);
  const unsigned char *start_color = (x_pixel_start < midpoint) ? first : second;
  plot_overlay_aa(pixels, width, height, steep, x_pixel_start, y_pixel,
                  (1.0 - overlay_fractional_part(y_end)) * x_gap, start_color,
                  settings->opacity, rgba);
  plot_overlay_aa(pixels, width, height, steep, x_pixel_start, y_pixel + 1,
                  overlay_fractional_part(y_end) * x_gap, start_color,
                  settings->opacity, rgba);
  double inter_y = y_end + gradient;
  x_end = (double)last_column;
  y_end = y1 + gradient * (x_end - x1);
  if (!isfinite(inter_y) || !isfinite(y_end))
    return;
  x_gap = overlay_fractional_part(x1 + 0.5);
  const int x_pixel_end = (int)x_end;
  y_pixel = (int)floor(y_end);
  const unsigned char *end_color = (x_pixel_end < midpoint) ? first : second;
  plot_overlay_aa(pixels, width, height, steep, x_pixel_end, y_pixel,
                  (1.0 - overlay_fractional_part(y_end)) * x_gap, end_color,
                  settings->opacity, rgba);
  plot_overlay_aa(pixels, width, height, steep, x_pixel_end, y_pixel + 1,
                  overlay_fractional_part(y_end) * x_gap, end_color,
                  settings->opacity, rgba);
  for (int x = x_pixel_start + 1; x < x_pixel_end; ++x) {
    y_pixel = (int)floor(inter_y);
    const unsigned char *color = (x < midpoint) ? first : second;
    plot_overlay_aa(pixels, width, height, steep, x, y_pixel,
                    1.0 - overlay_fractional_part(inter_y), color,
                    settings->opacity, rgba);
    plot_overlay_aa(pixels, width, height, steep, x, y_pixel + 1,
                    overlay_fractional_part(inter_y), color,
                    settings->opacity, rgba);
    inter_y += gradient;
  }
}

static int draw_overlay(const MeshOverlayLines *lines, unsigned char *pixels,
                         int width, int height,
                         const MeshOverlaySettings *settings, int rgba) {
  if (lines == NULL || pixels == NULL || settings == NULL)
    return -1;
  if (width <= 0 || height <= 0)
    return -1;
  /* Keep the raster's y+1 / x+1 and clipped-box endpoint casts strictly inside
   * `int`, and keep the RGB8/RGBA8 byte count inside `size_t`. */
  if (width > INT_MAX - 2 || height > INT_MAX - 2)
    return -1;
  const size_t dim_width = (size_t)width;
  const size_t dim_height = (size_t)height;
  if (dim_width > SIZE_MAX / dim_height)
    return -1;
  const size_t pixel_count = dim_width * dim_height;
  if (pixel_count > SIZE_MAX / (rgba ? 4 : 3))
    return -1;
  if (!isfinite(settings->opacity) || settings->opacity < 0.0 ||
      settings->opacity > 1.0)
    return -1;
  /* A line batch larger than any allocatable MeshOverlayLine array cannot be
   * real; reject it before dereferencing the array. */
  if (lines->count > SIZE_MAX / sizeof(MeshOverlayLine))
    return -1;
  if (lines->count != 0 && lines->lines == NULL)
    return -1;
  /* Validate the whole batch before drawing so an invalid line cannot leave a
   * partially painted image behind.  Endpoint differences too large to
   * represent (e.g. -DBL_MAX..+DBL_MAX) are rejected up front, before the
   * midpoint or clipping math can produce a nonfinite value. */
  for (size_t i = 0; i < lines->count; ++i) {
    const MeshOverlayLine *line = &lines->lines[i];
    if (line->category0 >= MESH_OVERLAY_CATEGORY_COUNT ||
        line->category1 >= MESH_OVERLAY_CATEGORY_COUNT)
      return -1;
    if (!isfinite(line->x0) || !isfinite(line->y0) || !isfinite(line->x1) ||
        !isfinite(line->y1))
      return -1;
    if (!isfinite(line->x1 - line->x0) || !isfinite(line->y1 - line->y0))
      return -1;
  }
  for (size_t i = 0; i < lines->count; ++i)
    draw_overlay_line(pixels, width, height, &lines->lines[i], settings, rgba);
  return 0;
}

int mesh_overlay_draw_rgb8(const MeshOverlayLines *lines, unsigned char *pixels,
                            int width, int height,
                            const MeshOverlaySettings *settings) {
  return draw_overlay(lines, pixels, width, height, settings, 0);
}

void mesh_overlay_layer_destroy(MeshOverlayLayer *layer) {
  if (layer == NULL)
    return;
  free(layer->rgba);
  *layer = (MeshOverlayLayer){0};
}

int mesh_overlay_build_layer(const FrameLensMesh *mesh, int width, int height,
                             const MeshOverlaySettings *settings,
                             MeshOverlayLayer *layer) {
  if (layer == NULL)
    return -1;
  *layer = (MeshOverlayLayer){0};
  if (width <= 0 || height <= 0 || width > INT_MAX - 2 ||
      height > INT_MAX - 2 || (size_t)width > SIZE_MAX / (size_t)height ||
      (size_t)width * height > SIZE_MAX / 4 || settings == NULL ||
      !isfinite(settings->opacity) || settings->opacity < 0.0 ||
      settings->opacity > 1.0)
    return -1;
  MeshOverlayLines lines = {0};
  if (mesh_overlay_prepare(mesh, &lines))
    return -1;
  unsigned char *pixels = calloc((size_t)width * height, 4);
  const int result = pixels == NULL ? -1 :
      draw_overlay(&lines, pixels, width, height, settings, 1);
  mesh_overlay_lines_destroy(&lines);
  if (result) {
    free(pixels);
    return -1;
  }
  layer->rgba = pixels;
  layer->width = width;
  layer->height = height;
  return 0;
}

int mesh_overlay_composite_rgb8(const MeshOverlayLayer *layer,
                                unsigned char *rgb8, int width, int height) {
  if (layer == NULL || layer->rgba == NULL || rgb8 == NULL || width <= 0 ||
      height <= 0 || width != layer->width || height != layer->height ||
      (size_t)width > SIZE_MAX / (size_t)height ||
      (size_t)width * height > SIZE_MAX / 4)
    return -1;
  const size_t count = (size_t)width * height;
  for (size_t pixel = 0; pixel < count; ++pixel) {
    const unsigned char *source = &layer->rgba[4 * pixel];
    const unsigned int alpha = source[3];
    if (alpha == 0)
      continue;
    for (int channel = 0; channel < 3; ++channel) {
      const unsigned int value = source[channel] +
          (rgb8[3 * pixel + channel] * (255 - alpha) + 127) / 255;
      rgb8[3 * pixel + channel] = value > 255 ? 255 : (unsigned char)value;
    }
  }
  return 0;
}