// MSDF/MTSDF generation shader. One quad per glyph is drawn into its atlas region.

struct SegmentGpu {
    p_from: vec2<f32>,
    p_to: vec2<f32>,
    p_mid: vec2<f32>,
    tangent: vec2<f32>,
    tangent_unit: vec2<f32>,
    inv_tangent_len_sq: f32,
    color_mask: u32,
}

@group(0) @binding(0) var<storage, read> segments: array<SegmentGpu>;

struct VertexInput {
    @location(0) clip_pos: vec2<f32>,
    @location(1) bound_pos: vec2<f32>,
    @location(2) seg_from: u32,
    @location(3) seg_to: u32,
    @location(4) units_per_em: f32,
}

struct VertexOutput {
    @builtin(position) clip_pos: vec4<f32>,
    @location(0) bound_pos: vec2<f32>,
    @location(1) @interpolate(flat) seg_from: u32,
    @location(2) @interpolate(flat) seg_to: u32,
    @location(3) @interpolate(flat) units_per_em: f32,
}

@vertex
fn vs_main(in: VertexInput) -> VertexOutput {
    var out: VertexOutput;
    out.clip_pos = vec4<f32>(in.clip_pos, 0.0, 1.0);
    out.bound_pos = in.bound_pos;
    out.seg_from = in.seg_from;
    out.seg_to = in.seg_to;
    out.units_per_em = in.units_per_em;
    return out;
}

fn cross2(a: vec2<f32>, b: vec2<f32>) -> f32 {
    return a.x * b.y - a.y * b.x;
}

const DIST_TOLERANCE: f32 = 0.1;
const EPS: f32 = 1e-6;

struct ClosestSet {
    seg: array<u32, 3>,
    has: array<u32, 3>,
}

fn find_closest_seg(pos: vec2<f32>, seg_from: u32, seg_to: u32) -> ClosestSet {
    var best_dist = array<f32, 3>(3.4e38, 3.4e38, 3.4e38);
    var best_ortho_sq = array<f32, 3>(-3.4e38, -3.4e38, -3.4e38);
    var best_seg = array<u32, 3>(0u, 0u, 0u);
    var has = array<u32, 3>(0u, 0u, 0u);

    for (var i = seg_from; i < seg_to; i = i + 1u) {
        let seg = segments[i];
        let pf = pos - seg.p_from;
        let t = clamp(dot(pf, seg.tangent) * seg.inv_tangent_len_sq, 0.0, 1.0);
        let perp = seg.p_from + seg.tangent * t - pos;
        let dist = dot(perp, perp);

        let cmask = seg.color_mask;
        for (var col = 0u; col < 3u; col = col + 1u) {
            if (((cmask >> col) & 1u) == 0u) {
                continue;
            }
            let diff = dist - best_dist[col];
            let smaller = diff < -DIST_TOLERANCE;
            let equal = abs(diff) < DIST_TOLERANCE;
            if (smaller) {
                best_dist[col] = dist;
                best_seg[col] = i;
                has[col] = 1u;
                let c = cross2(seg.tangent_unit, perp);
                best_ortho_sq[col] = select(0.0, c * c / dist, dist > 0.0);
            } else if (equal) {
                let c = cross2(seg.tangent_unit, perp);
                let ortho_sq = select(0.0, c * c / dist, dist > 0.0);
                if (ortho_sq > best_ortho_sq[col]) {
                    best_dist[col] = dist;
                    best_seg[col] = i;
                    has[col] = 1u;
                    best_ortho_sq[col] = ortho_sq;
                }
            }
        }
    }

    var out: ClosestSet;
    out.seg = best_seg;
    out.has = has;
    return out;
}

// Cardano's method on f32. WGSL has no f64.
struct CubicRoots {
    n: u32,
    t: array<f32, 3>,
}

fn solve_cubic(a: f32, b: f32, c: f32, d: f32) -> CubicRoots {
    var out: CubicRoots;
    out.n = 0u;
    let aa = b / a;
    let bb = c / a;
    let cc = d / a;

    let a2 = aa * aa;
    let q = (a2 - 3.0 * bb) / 9.0;
    let r = (aa * (2.0 * a2 - 9.0 * bb) + 27.0 * cc) / 54.0;
    let r2 = r * r;
    let q3 = q * q * q;
    let a_third = aa / 3.0;

    if (r2 < q3) {
        let t = clamp(r / sqrt(q3), -1.0, 1.0);
        let t_acos = acos(t);
        let q_factor = -2.0 * sqrt(q);
        out.t[0] = q_factor * cos(t_acos / 3.0) - a_third;
        out.t[1] = q_factor * cos((t_acos + 2.0 * 3.14159265358979) / 3.0) - a_third;
        out.t[2] = q_factor * cos((t_acos - 2.0 * 3.14159265358979) / 3.0) - a_third;
        out.n = 3u;
    } else {
        let abs_r = abs(r);
        let sign_r = select(-1.0, 1.0, r < 0.0);
        let u = pow(abs_r + sqrt(r2 - q3), 1.0 / 3.0) * sign_r;
        let v = select(q / u, 0.0, u == 0.0);
        let x0 = u + v - a_third;
        out.t[0] = x0;
        if (u == v || abs(u - v) < abs(u + v) * 1e-12) {
            out.t[1] = (u + v) * -0.5 - a_third;
            out.n = 2u;
        } else {
            out.n = 1u;
        }
    }
    return out;
}

fn quad_bezier(p0: vec2<f32>, p1: vec2<f32>, p2: vec2<f32>, t: f32) -> vec2<f32> {
    return p0 + 2.0 * t * (p1 - p0) + t * t * (p2 - 2.0 * p1 + p0);
}

fn d_quad_bezier(p0: vec2<f32>, p1: vec2<f32>, p2: vec2<f32>, t: f32) -> vec2<f32> {
    return 2.0 * (p1 - p0) + 2.0 * t * (p2 - 2.0 * p1 + p0);
}

fn midpoint(a: vec2<f32>, b: vec2<f32>) -> vec2<f32> {
    return (a + b) * 0.5;
}

fn signed_dist_line(pos: vec2<f32>, pa: vec2<f32>, pb: vec2<f32>) -> f32 {
    let tangent = pb - pa;
    let tlen_sq = dot(tangent, tangent);
    if (tlen_sq <= 1e-6) {
        let perpendicular = pos - midpoint(pa, pb);
        return sign(cross2(tangent, perpendicular)) * length(pos - pa);
    }
    let t = dot(pos - pa, tangent) / tlen_sq;
    let t_clamp = clamp(t, 0.0, 1.0);
    let projected = pa + tangent * t;
    let projected_clamp = pa + tangent * t_clamp;
    let perpendicular = projected - pos;
    let shortest = projected_clamp - pos;
    return sign(cross2(tangent, perpendicular)) * length(shortest);
}

fn pseudo_signed_dist_line(pos: vec2<f32>, pa: vec2<f32>, pb: vec2<f32>) -> f32 {
    let tangent = pb - pa;
    let tlen_sq = dot(tangent, tangent);
    if (tlen_sq <= 1e-6) {
        let perpendicular = pos - midpoint(pa, pb);
        return sign(cross2(tangent, perpendicular)) * length(pos - pa);
    }
    let t = dot(pos - pa, tangent) / tlen_sq;
    let projected = pa + tangent * t;
    let perpendicular = projected - pos;
    return sign(cross2(tangent, perpendicular)) * length(perpendicular);
}

fn signed_dist_seg(pos: vec2<f32>, pa: vec2<f32>, pm: vec2<f32>, pb: vec2<f32>) -> f32 {
    let p = pos - pa;
    let p1 = pm - pa;
    let p2 = pb - 2.0 * pm + pa;
    if (dot(p2, p2) <= EPS) {
        return signed_dist_line(pos, pa, pb);
    }
    let roots = solve_cubic(dot(p2, p2), 3.0 * dot(p1, p2), 2.0 * dot(p1, p1) - dot(p2, p), -dot(p1, p));
    var min_dist = 3.4e38;
    // candidate ts: roots[0..n], plus 0 and 1
    var ts = array<f32, 5>(0.0, 1.0, 0.0, 0.0, 0.0);
    var nt = 2u;
    if (roots.n >= 1u) { ts[nt] = roots.t[0]; nt = nt + 1u; }
    if (roots.n >= 2u) { ts[nt] = roots.t[1]; nt = nt + 1u; }
    if (roots.n >= 3u) { ts[nt] = roots.t[2]; nt = nt + 1u; }
    for (var k = 0u; k < nt; k = k + 1u) {
        let t = ts[k];
        if (t < 0.0 || t > 1.0) { continue; }
        let perpendicular = quad_bezier(pa, pm, pb, t) - pos;
        let dist = length(perpendicular);
        if (dist < abs(min_dist)) {
            min_dist = dist * sign(cross2(d_quad_bezier(pa, pm, pb, t), perpendicular));
        }
    }
    return min_dist;
}

fn pseudo_signed_dist_seg(pos: vec2<f32>, pa: vec2<f32>, pm: vec2<f32>, pb: vec2<f32>) -> f32 {
    let p = pos - pa;
    let p1 = pm - pa;
    let p2 = pb - 2.0 * pm + pa;
    if (dot(p2, p2) <= EPS) {
        return pseudo_signed_dist_line(pos, pa, pb);
    }
    let roots = solve_cubic(dot(p2, p2), 3.0 * dot(p1, p2), 2.0 * dot(p1, p1) - dot(p2, p), -dot(p1, p));
    var min_dist = 3.4e38;
    var ts = array<f32, 3>(0.0, 0.0, 0.0);
    var nt = 0u;
    if (roots.n >= 1u) { ts[nt] = roots.t[0]; nt = nt + 1u; }
    if (roots.n >= 2u) { ts[nt] = roots.t[1]; nt = nt + 1u; }
    if (roots.n >= 3u) { ts[nt] = roots.t[2]; nt = nt + 1u; }
    for (var k = 0u; k < nt; k = k + 1u) {
        let t = ts[k];
        if (t < 0.0) {
            let tangent = pm - pa;
            var perpendicular: vec2<f32>;
            if (dot(tangent, tangent) <= 1e-6) {
                perpendicular = pos - midpoint(pa, pm);
            } else {
                let tt = dot(pos - pa, tangent) / dot(tangent, tangent);
                perpendicular = (pa + tangent * tt) - pos;
            }
            let dist = length(perpendicular);
            if (dist < abs(min_dist)) {
                min_dist = dist * sign(cross2(tangent, perpendicular));
            }
        } else if (t > 1.0) {
            let tangent = pb - pm;
            var perpendicular: vec2<f32>;
            if (dot(tangent, tangent) <= 1e-6) {
                perpendicular = pos - midpoint(pm, pb);
            } else {
                let tt = dot(pos - pm, tangent) / dot(tangent, tangent);
                perpendicular = (pm + tangent * tt) - pos;
            }
            let dist = length(perpendicular);
            if (dist < abs(min_dist)) {
                min_dist = dist * sign(cross2(tangent, perpendicular));
            }
        } else {
            let tangent = d_quad_bezier(pa, pm, pb, t);
            let perpendicular = quad_bezier(pa, pm, pb, t) - pos;
            let dist = length(perpendicular);
            if (dist < abs(min_dist)) {
                min_dist = dist * sign(cross2(tangent, perpendicular));
            }
        }
    }
    return min_dist;
}

fn norm_dist(dist: f32, units_per_em: f32) -> f32 {
    return clamp(dist / units_per_em + 0.5, 0.0, 1.0);
}

// Compute the unclamped MTSDF at `pos`. RGB are pseudo distances, A is the
// true signed distance, all in em units (not yet normalised).
fn mtsdf_at(pos: vec2<f32>, seg_from: u32, seg_to: u32) -> vec4<f32> {
    let closest = find_closest_seg(pos, seg_from, seg_to);
    var pseudo = vec3<f32>(0.0);
    for (var col = 0u; col < 3u; col = col + 1u) {
        if (closest.has[col] != 0u) {
            let s = segments[closest.seg[col]];
            pseudo[col] = pseudo_signed_dist_seg(pos, s.p_from, s.p_mid, s.p_to);
        }
    }

    var closest_true = 3.4e38;
    var max_ortho = 0.0;
    var sign_seg: u32 = 0u;
    var has_true = false;
    for (var col = 0u; col < 3u; col = col + 1u) {
        if (closest.has[col] != 0u) {
            let s = segments[closest.seg[col]];
            let pf = pos - s.p_from;
            let t = clamp(dot(pf, s.tangent) * s.inv_tangent_len_sq, 0.0, 1.0);
            let proj = s.p_from + s.tangent * t;
            let d = dot(proj - pos, proj - pos);
            // Orthogonality tie-break (Chlumsky 2.43): |cross(tangent_unit, perp/|perp|)|.
            // perp here is `pos - projected`, matching `orthogonality_from_soa`.
            let perp_pos = pos - proj;
            let perp_len = length(perp_pos);
            let ortho = select(0.0, abs(cross2(s.tangent_unit, perp_pos / perp_len)), perp_len > 0.0);
            let diff = d - closest_true;
            let smaller = diff <= -1.0;
            let equal = abs(diff) < 1.0;
            if (smaller || (equal && ortho > max_ortho)) {
                closest_true = d;
                sign_seg = closest.seg[col];
                max_ortho = ortho;
                has_true = true;
            }
        }
    }
    var true_signed: f32 = 0.0;
    if (has_true) {
        let s = segments[sign_seg];
        true_signed = signed_dist_seg(pos, s.p_from, s.p_mid, s.p_to);
    }
    return vec4<f32>(pseudo, true_signed);
}

@fragment
fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
    let pos = in.bound_pos;
    let em = in.units_per_em;

    let raw = mtsdf_at(pos, in.seg_from, in.seg_to);
    let rgb = vec3<f32>(
        norm_dist(raw.x, em),
        norm_dist(raw.y, em),
        norm_dist(raw.z, em),
    );
    let alpha = norm_dist(raw.w, em);

    return vec4<f32>(rgb, alpha);
}
