#include <metal_stdlib>
using namespace metal;

#define M_PI_F 3.14159265358979323846f

struct VertexOut {
  float4 position [[position]];
  float2 uv;
  float4 color;
};

struct ShapeParams {
  float4 fill_color;
  float4 stroke_color;
  float2 size;
  float radius;
  float feather;
  float stroke_thickness;
  float stroke_align;
  int shape_type;      // 0=rect (rounded via radius), 1=circle, 2=ellipse, 3=pill
  int corner_style;
  float dash_len;
  float gap_len;
  float dash_offset;
  int cap_type;        // 0=butt, 1=square-ish, 2=round-ish
  float4 uv_transform;
  float uv_rotate;
  int has_texture;
  float opacity;
  float _pad;
};

// -------- SDFs --------

float sdf_rounded_rect(float2 p, float2 b, float r) {
  float2 q = abs(p) - b + r;
  return min(max(q.x, q.y), 0.0f) + length(max(q, 0.0f)) - r;
}

float sdf_circle(float2 p, float r) {
  return length(p) - r;
}

float sdf_ellipse(float2 p, float2 ab) {
  float2 pn = p / ab;
  float d = length(pn) - 1.0f;
  return d * min(ab.x, ab.y);
}

float sdf_pill(float2 p, float2 b) {
  float r = min(b.x, b.y);
  float2 q = abs(p) - float2(b.x - r, b.y - r);
  return length(max(q, 0.0f)) + min(max(q.x, q.y), 0.0f) - r;
}

float compute_sdf(float2 p, float2 half_size, int shape_type, float radius) {
  if (shape_type == 0) {
    float r = min(radius, min(half_size.x, half_size.y));
    return sdf_rounded_rect(p, half_size, r);
  }
  else if (shape_type == 1) {
    float r = min(half_size.x, half_size.y);
    return sdf_circle(p, r);
  }
  else if (shape_type == 2) {
    return sdf_ellipse(p, half_size);
  }
  else {
    return sdf_pill(p, half_size);
  }
}

// -------- Helpers --------

float wrap_angle_0_2pi(float a) {
  float w = fmod(a, 2.0f * M_PI_F);
  if (w < 0.0f) w += 2.0f * M_PI_F;
  return w;
}

// Project arbitrary p to ellipse boundary (stable and cheap)
// ab are ellipse radii (half_size.x, half_size.y)
float2 ellipse_project_boundary(float2 p, float2 ab) {
  float2 safe_ab = max(ab, float2(1e-6f));
  float2 pn = p / safe_ab;
  float len_pn = length(pn);
  if (len_pn < 1e-6f) return float2(ab.x, 0.0f);
  float2 dir = pn / len_pn;
  return dir * ab;
}

// Numeric arc length from 0..theta for ellipse with radii a,b using midpoint rule.
float ellipse_arclen_0_theta(float a, float b, float theta) {
  const int N = 12;
  float t = 0.0f;
  float dt = theta / float(N);
  float sum = 0.0f;
  for (int i = 0; i < N; i++) {
    float tm = t + 0.5f * dt;
    float s = sin(tm);
    float c = cos(tm);
    float ds = sqrt(a * a * s * s + b * b * c * c);
    sum += ds * dt;
    t += dt;
  }
  return sum;
}

float ellipse_perimeter_ramanujan(float a, float b) {
  float apb = a + b;
  float h = (a - b) * (a - b) / max(apb * apb, 1e-6f);
  return M_PI_F * apb * (1.0f + (3.0f * h) / (10.0f + sqrt(max(4.0f - 3.0f * h, 1e-6f))));
}

// Rounded-rect boundary projection (for stable dash phase across stroke width)
float2 rounded_rect_project_boundary(float2 p, float2 half_size, float r) {
  float2 inner = max(half_size - r, float2(0.0f));
  float2 ap = abs(p);

  if (ap.x > inner.x && ap.y > inner.y && r > 0.0f) {
    float sx = (p.x < 0.0f) ? -1.0f : 1.0f;
    float sy = (p.y < 0.0f) ? -1.0f : 1.0f;
    float2 c = float2(inner.x * sx, inner.y * sy);
    float2 v = p - c;
    float lv = length(v);
    if (lv < 1e-6f) return c + float2(r, 0.0f);
    return c + v * (r / lv);
  }

  float dx = half_size.x - ap.x;
  float dy = half_size.y - ap.y;

  if (dx < dy) {
    float sx = (p.x < 0.0f) ? -1.0f : 1.0f;
    return float2(sx * half_size.x, clamp(p.y, -inner.y, inner.y));
  }
  float sy = (p.y < 0.0f) ? -1.0f : 1.0f;
  return float2(clamp(p.x, -inner.x, inner.x), sy * half_size.y);
}

// Rounded-rect perimeter distance along boundary, starting at +X mid, CCW.
float rounded_rect_dist_along(float2 pb, float2 half_size, float r) {
  float2 inner = max(half_size - r, float2(0.0f));
  float ix = inner.x;
  float iy = inner.y;

  float seg_right_up = iy;
  float seg_arc = 0.5f * M_PI_F * r;
  float seg_top = 2.0f * ix;
  float seg_left = 2.0f * iy;
  float seg_bottom = 2.0f * ix;

  float ax = abs(pb.x);
  float ay = abs(pb.y);

  // Right edge
  if (pb.x > 0.0f && ay <= iy + 1e-5f) {
    if (pb.y >= 0.0f) return pb.y;
    float tail = seg_right_up + seg_arc + seg_top + seg_arc + seg_left + seg_arc + seg_bottom + seg_arc;
    return tail + (pb.y + iy);
  }

  // Top edge
  if (pb.y > 0.0f && ax <= ix + 1e-5f) {
    float s0 = seg_right_up + seg_arc;
    return s0 + (ix - pb.x);
  }

  // Left edge
  if (pb.x < 0.0f && ay <= iy + 1e-5f) {
    float s0 = seg_right_up + seg_arc + seg_top + seg_arc;
    return s0 + (iy - pb.y);
  }

  // Bottom edge
  if (pb.y < 0.0f && ax <= ix + 1e-5f) {
    float s0 = seg_right_up + seg_arc + seg_top + seg_arc + seg_left + seg_arc;
    return s0 + (pb.x + ix);
  }

  // Corners
  if (pb.x > 0.0f && pb.y > 0.0f) {
    float2 c = float2(ix, iy);
    float2 v = pb - c;
    float a = wrap_angle_0_2pi(atan2(v.y, v.x));
    float local = clamp(a, 0.0f, 0.5f * M_PI_F);
    return seg_right_up + r * local;
  }
  if (pb.x < 0.0f && pb.y > 0.0f) {
    float2 c = float2(-ix, iy);
    float2 v = pb - c;
    float a = wrap_angle_0_2pi(atan2(v.y, v.x));
    float local = clamp(M_PI_F - a, 0.0f, 0.5f * M_PI_F);
    float s0 = seg_right_up + seg_arc + seg_top;
    return s0 + r * local;
  }
  if (pb.x < 0.0f && pb.y < 0.0f) {
    float2 c = float2(-ix, -iy);
    float2 v = pb - c;
    float a = wrap_angle_0_2pi(atan2(v.y, v.x));
    float local = clamp(a - M_PI_F, 0.0f, 0.5f * M_PI_F);
    float s0 = seg_right_up + seg_arc + seg_top + seg_arc + seg_left;
    return s0 + r * local;
  }

  // Bottom-right
  {
    float2 c = float2(ix, -iy);
    float2 v = pb - c;
    float a = wrap_angle_0_2pi(atan2(v.y, v.x));
    float local = 0.0f;
    if (a >= 1.5f * M_PI_F) local = clamp(a - 1.5f * M_PI_F, 0.0f, 0.5f * M_PI_F);
    float s0 = seg_right_up + seg_arc + seg_top + seg_arc + seg_left + seg_arc + seg_bottom;
    return s0 + r * local;
  }
}

// -------- Pill helpers --------

float pill_perimeter(float2 half_size) {
  float r = min(half_size.x, half_size.y);
  float major = max(half_size.x, half_size.y);
  float straight = max(major - r, 0.0f);
  return 4.0f * straight + 2.0f * M_PI_F * r;
}

float2 pill_project_boundary(float2 p, float2 half_size) {
  float r = min(half_size.x, half_size.y);
  bool horiz = (half_size.x >= half_size.y);
  if (horiz) {
    float cx = max(half_size.x - r, 0.0f);
    float2 c = float2((p.x >= 0.0f ? cx : -cx), 0.0f);
    float2 v = p - c;
    float lv = length(v);
    if (lv < 1e-6f) return c + float2(r, 0.0f);
    return c + v * (r / lv);
  }
  float cy = max(half_size.y - r, 0.0f);
  float2 c = float2(0.0f, (p.y >= 0.0f ? cy : -cy));
  float2 v = p - c;
  float lv = length(v);
  if (lv < 1e-6f) return c + float2(r, 0.0f);
  return c + v * (r / lv);
}

float pill_dist_along(float2 pb, float2 half_size) {
  float r = min(half_size.x, half_size.y);
  bool horiz = (half_size.x >= half_size.y);

  if (horiz) {
    float cx = max(half_size.x - r, 0.0f);

    if (pb.x > cx - 1e-5f) {
      float2 c = float2(cx, 0.0f);
      float2 v = pb - c;
      float a = wrap_angle_0_2pi(atan2(v.y, v.x));
      float local = clamp(a, 0.0f, M_PI_F);
      return r * local;
    }

    if (abs(pb.y - r) <= 1e-3f && abs(pb.x) <= cx + 1e-3f) {
      float s0 = M_PI_F * r;
      return s0 + (cx - pb.x);
    }

    if (pb.x < -cx + 1e-5f) {
      float2 c = float2(-cx, 0.0f);
      float2 v = pb - c;
      float a = wrap_angle_0_2pi(atan2(v.y, v.x));
      float local = clamp(a - M_PI_F, 0.0f, M_PI_F);
      float s0 = M_PI_F * r + 2.0f * cx;
      return s0 + r * local;
    }

    float s0 = M_PI_F * r + 2.0f * cx + M_PI_F * r;
    return s0 + (pb.x + cx);
  }

  float cy = max(half_size.y - r, 0.0f);

  if (pb.y > cy - 1e-5f) {
    float2 c = float2(0.0f, cy);
    float2 v = pb - c;
    float a = wrap_angle_0_2pi(atan2(v.y, v.x));
    float local = clamp(a, 0.0f, M_PI_F);
    return r * local;
  }

  if (abs(pb.x + r) <= 1e-3f && abs(pb.y) <= cy + 1e-3f) {
    float s0 = M_PI_F * r;
    return s0 + (cy - pb.y);
  }

  if (pb.y < -cy + 1e-5f) {
    float2 c = float2(0.0f, -cy);
    float2 v = pb - c;
    float a = wrap_angle_0_2pi(atan2(v.y, v.x));
    float local = clamp(a - M_PI_F, 0.0f, M_PI_F);
    float s0 = M_PI_F * r + 2.0f * cy;
    return s0 + r * local;
  }

  float s0 = M_PI_F * r + 2.0f * cy + M_PI_F * r;
  return s0 + (pb.y + cy);
}

// -------- Perimeter mapping for dashes --------

float perimeter_dist_along(float2 p, float2 half_size, int shape_type, float radius) {
  if (shape_type == 1) {
    float r = min(half_size.x, half_size.y);
    float2 pb = (length(p) < 1e-6f) ? float2(r, 0.0f) : normalize(p) * r;
    float theta = wrap_angle_0_2pi(atan2(pb.y, pb.x));
    return r * theta;
  }

  if (shape_type == 2) {
    float a = max(half_size.x, 1e-6f);
    float b = max(half_size.y, 1e-6f);
    float2 pb = ellipse_project_boundary(p, float2(a, b));
    float theta = wrap_angle_0_2pi(atan2(pb.y / b, pb.x / a));
    return ellipse_arclen_0_theta(a, b, theta);
  }

  if (shape_type == 3) {
    float2 pb = pill_project_boundary(p, half_size);
    return pill_dist_along(pb, half_size);
  }

  float r = min(radius, min(half_size.x, half_size.y));
  float2 pb = rounded_rect_project_boundary(p, half_size, r);
  return rounded_rect_dist_along(pb, half_size, r);
}

float perimeter_total(float2 half_size, int shape_type, float radius) {
  if (shape_type == 1) {
    float r = min(half_size.x, half_size.y);
    return 2.0f * M_PI_F * r;
  }
  if (shape_type == 2) {
    float a = max(half_size.x, 1e-6f);
    float b = max(half_size.y, 1e-6f);
    return ellipse_perimeter_ramanujan(a, b);
  }
  if (shape_type == 3) {
    return pill_perimeter(half_size);
  }
  float r = min(radius, min(half_size.x, half_size.y));
  float2 inner = max(half_size - r, float2(0.0f));
  return 4.0f * (inner.x + inner.y) + 2.0f * M_PI_F * r;
}

float dash_mask_perimeter(float dist_along, float perimeter,
                          float dash_len, float gap_len, float dash_offset,
                          int cap_type, float aa) {
  if (dash_len <= 0.0f) return 1.0f;

  float pattern_len = max(dash_len + gap_len, 1e-6f);
  float s = dist_along + dash_offset;

  float pwrap = max(perimeter, 1e-6f);
  float s_wrap = fmod(s, pwrap);
  if (s_wrap < 0.0f) s_wrap += pwrap;

  float pos = fmod(s_wrap, pattern_len);
  if (pos < 0.0f) pos += pattern_len;

  if (cap_type == 0) {
    return pos < dash_len ? 1.0f : 0.0f;
  }

  if (cap_type == 1) {
    float edge = dash_len;
    return 1.0f - smoothstep(edge - aa, edge + aa, pos);
  }

  float half_dash = dash_len * 0.5f;
  float center = half_dash;
  float d = abs(pos - center);
  return smoothstep(half_dash + aa, half_dash - aa, d);
}

// -------- Fragment --------

fragment float4 fragment_main(VertexOut in [[stage_in]],
                              constant ShapeParams& params [[buffer(0)]],
                              texture2d<float> fill_texture [[texture(0)]],
                              sampler smp [[sampler(0)]]) {
  float2 uv = in.uv;

  float2 p = (uv - 0.5f) * params.size;
  float2 half_size = params.size * 0.5f;

  float d = compute_sdf(p, half_size, params.shape_type, params.radius);

  float aa = max(params.feather, fwidth(d));

  float fill_alpha = 0.0f;
  float stroke_alpha = 0.0f;

  if (params.stroke_thickness > 0.0f) {
    float inner_offset, outer_offset;

    if (params.stroke_align < 0.25f) {
      inner_offset = -params.stroke_thickness;
      outer_offset = 0.0f;
    }
    else if (params.stroke_align > 0.75f) {
      inner_offset = 0.0f;
      outer_offset = params.stroke_thickness;
    }
    else {
      inner_offset = -params.stroke_thickness * 0.5f;
      outer_offset = params.stroke_thickness * 0.5f;
    }

    float fill_d = d - inner_offset;
    fill_alpha = 1.0f - smoothstep(-aa, aa, fill_d);

    float start = smoothstep(-aa, aa, d - inner_offset);
    float end = 1.0f - smoothstep(-aa, aa, d - outer_offset);
    float in_stroke = start * end;

    float perim = perimeter_total(half_size, params.shape_type, params.radius);
    float dist_along = perimeter_dist_along(p, half_size, params.shape_type, params.radius);
    float dash = dash_mask_perimeter(dist_along, perim,
                                     params.dash_len, params.gap_len, params.dash_offset,
                                     params.cap_type, aa);

    stroke_alpha = in_stroke * dash;
  }
  else {
    fill_alpha = 1.0f - smoothstep(-aa, aa, d);
  }

  float4 fill_col = params.fill_color;
  if (params.has_texture > 0) {
    float2 tex_uv = uv;
    tex_uv = tex_uv * params.uv_transform.xy + params.uv_transform.zw;

    if (abs(params.uv_rotate) > 0.001f) {
      float2 center = float2(0.5f, 0.5f);
      tex_uv -= center;
      float c = cos(params.uv_rotate);
      float s = sin(params.uv_rotate);
      tex_uv = float2(tex_uv.x * c - tex_uv.y * s, tex_uv.x * s + tex_uv.y * c);
      tex_uv += center;
    }

    float4 tex_col = fill_texture.sample(smp, tex_uv);
    fill_col *= tex_col;
  }

  float4 fill_out = fill_col;
  fill_out.a *= fill_alpha;

  float4 stroke_out = params.stroke_color;
  stroke_out.a *= stroke_alpha;

  float sa = stroke_out.a;
  float fa = fill_out.a;

  float3 rgb_premul = stroke_out.rgb * sa + fill_out.rgb * fa * (1.0f - sa);
  float a = sa + fa * (1.0f - sa);

  float4 result = (a > 1e-6f) ? float4(rgb_premul / a, a) : float4(0.0f);

  result *= in.color;
  result.a *= params.opacity;

  return result;
}