/*
 * Copyright 2026 Behdad Esfahbod. All Rights Reserved.
 */

struct Uniforms {
  float4x4 matViewProjection;
  float2 viewport;
  float gamma;
  float stem_darkening;
  float4 foreground;
  float debug;
};

struct VertexIn {
  float2 position [[attribute(0)]];
  float2 texcoord [[attribute(1)]];
  float2 normal   [[attribute(2)]];
  float  emPerPos [[attribute(3)]];
  uint   glyphLoc [[attribute(4)]];
};

struct VertexOut {
  float4 position [[position]];
  float2 texcoord;
  uint   glyphLoc [[flat]];
};

vertex VertexOut vertex_main(VertexIn in [[stage_in]],
                             constant Uniforms& uniforms [[buffer(1)]]) {
  float2 pos = in.position;
  float2 tex = in.texcoord;
  float4 jac = float4(in.emPerPos, 0.0, 0.0, -in.emPerPos);

  hb_gpu_dilate(pos, tex, in.normal, jac,
                uniforms.matViewProjection, uniforms.viewport);

  VertexOut out;
  out.position = uniforms.matViewProjection * float4(pos, 0.0, 1.0);
  out.texcoord = tex;
  out.glyphLoc = in.glyphLoc;
  return out;
}

fragment float4 fragment_main(VertexOut in [[stage_in]],
                              constant Uniforms& uniforms [[buffer(1)]],
                              device const short4* atlas [[buffer(0)]]) {
#ifdef HB_GPU_DEMO_DRAW
  float cov = hb_gpu_draw(in.texcoord, in.glyphLoc, atlas);
  float4 c = float4(uniforms.foreground.rgb * uniforms.foreground.a,
                    uniforms.foreground.a) * cov;
#else
  float cov;
  float4 c = hb_gpu_paint(in.texcoord, in.glyphLoc, uniforms.foreground,
                          atlas, cov);

  if (cov > 0.0 && cov < 1.0) {
    float adj = cov;
    if (uniforms.stem_darkening > 0.0) {
      float brightness = c.a > 0.0
        ? dot(c.rgb, float3(1.0 / 3.0)) / c.a : 0.0;
      float2 fw = fwidth(in.texcoord);
      adj = hb_gpu_stem_darken(adj, brightness,
        1.0 / max(fw.x, fw.y));
    }
    if (uniforms.gamma != 1.0)
      adj = pow(adj, uniforms.gamma);
    c *= adj / cov;
  }
#endif

  if (uniforms.debug > 0.0) {
    int2 counts = _hb_gpu_curve_counts(in.texcoord, in.glyphLoc, atlas);
    float r = clamp(float(counts.x) / 8.0, 0.0, 1.0);
    float g = clamp(float(counts.y) / 8.0, 0.0, 1.0);
    return float4(r, g, c.a, max(max(r, g), c.a));
  }

  return c;
}
