LlamaCasty/Services/BackgroundEffects/BackgroundEffectProcessor.cs

using ytLive.Helpers;
using ytLive.Models;

namespace ytLive.Services.BackgroundEffects;

/// <summary>
/// Applies a webcam's <see cref="BackgroundEffect"/> to one captured frame (TASK 51).
/// Pure pixel work over the <see cref="VideoFrame"/> seam — WPF-free and thread-agnostic,
/// so the camera's processing worker can call it before the frame reaches the preview bitmap
/// and the compositor.
///
/// <list type="bullet">
/// <item><see cref="BackgroundEffect.None"/> returns the SAME frame instance — a true
/// passthrough, no copy.</item>
/// <item><see cref="BackgroundEffect.Remove"/> returns a copy whose alpha is the person mask,
/// RGB untouched, <c>IsOpaque = false</c> — the compositor's straight-alpha blend then shows
/// the layers below the webcam. Strength = edge softness (0 = crisp cut-out, 1 = the model's
/// raw soft mask).</item>
/// <item><see cref="BackgroundEffect.Blur"/> returns a copy where the background is replaced
/// by a blurred version of itself and the person stays sharp. Strength = blur amount.</item>
/// </list>
///
/// The mask is computed ONCE per camera frame (<see cref="TrySegment"/>) and then rendered for
/// every effect variant in use (<see cref="Render"/>). A segmenter that returns no mask (or
/// throws) never drops the layer: the frame passes through and the failure is logged once.
/// </summary>
public sealed class BackgroundEffectProcessor
{
    public const double DefaultStrength = 0.5;

    /// <summary>The background is blurred at 1/<see cref="BlurDownscale"/> resolution and
    /// upsampled while blending — a strong blur at a fraction of the cost.</summary>
    internal const int BlurDownscale = 4;

    /// <summary>Box-blur radius (in downscaled pixels) at strength 0 and 1. The maximum was halved
    /// from 12 (2026-10-09): the creator found full range too heavy — the old max spread ≈48 full
    /// px. The established tools blur far less: StreamFX defaults its Box blur size to 5 and its
    /// author stops downscale-based blurs around 8 px.</summary>
    internal const int MinBlurRadius = 1;
    internal const int MaxBlurRadius = 6;

    private readonly IPersonSegmenter _segmenter;
    private bool _loggedFailure;

    public BackgroundEffectProcessor(IPersonSegmenter segmenter)
    {
        _segmenter = segmenter;
    }

    public VideoFrame Apply(VideoFrame frame, BackgroundEffect effect, double strength = DefaultStrength)
    {
        if (effect == BackgroundEffect.None) return frame;
        var mask = TrySegment(frame);
        return mask == null ? frame : Render(frame, mask, effect, strength);
    }

    /// <summary>Runs the segmenter once; null (logged once) when it can't produce a usable mask.</summary>
    public SegmentationMask? TrySegment(VideoFrame frame)
    {
        SegmentationMask? mask;
        try
        {
            mask = _segmenter.Segment(frame);
        }
        catch (Exception ex)
        {
            LogFailureOnce($"{ex.GetType().Name}: {ex.Message}");
            return null;
        }

        if (mask == null || mask.Width < 1 || mask.Height < 1 || mask.Values.Length < mask.Width * mask.Height)
        {
            LogFailureOnce("segmenter returned no usable mask");
            return null;
        }
        return mask;
    }

    public static VideoFrame Render(VideoFrame frame, SegmentationMask mask, BackgroundEffect effect, double strength)
    {
        var s = Math.Clamp(strength, 0, 1);
        return effect switch
        {
            BackgroundEffect.Remove => ApplyRemove(frame, mask, s),
            BackgroundEffect.Blur => ApplyBlur(frame, mask, s),
            _ => frame,
        };
    }

    internal static VideoFrame ApplyRemove(VideoFrame frame, SegmentationMask mask, double strength)
    {
        var src = frame.BgraPixels;
        var dst = new byte[src.Length];
        Buffer.BlockCopy(src, 0, dst, 0, src.Length);

        // Edge feather: map the mask's [0.5 - w, 0.5 + w] band to alpha 0..1, as a 256-entry LUT
        // over the quantized mask value (so the per-pixel work is a lookup + one multiply).
        var w = 0.05 + 0.45 * strength;
        var lo = 0.5 - w;
        var span = 2 * w;
        var alphaLut = new int[256];
        for (var i = 0; i < 256; i++)
            alphaLut[i] = (int)Math.Round(Math.Clamp((i / 255.0 - lo) / span, 0, 1) * 256);

        var width = frame.Width;
        var height = frame.Height;
        var stride = frame.Stride;
        var maskRows = UpsampleMaskRows(mask, width);
        var (r0, r1, fy) = AxisLut(mask.Height, height);

        Parallel.For(0, height, RowParallelism, y =>
        {
            var a0 = maskRows[r0[y]];
            var a1 = maskRows[r1[y]];
            var t = fy[y];
            var row = y * stride + 3;
            for (var x = 0; x < width; x++)
            {
                var m = a0[x] + (a1[x] - a0[x]) * t;
                var q = (int)(m * 255f + 0.5f);
                q = q < 0 ? 0 : q > 255 ? 255 : q;
                var p = row + x * 4;
                dst[p] = (byte)((src[p] * alphaLut[q] + 128) >> 8);
            }
        });

        return new VideoFrame(width, height, dst) { Epoch = frame.Epoch, IsOpaque = false };
    }

    internal static VideoFrame ApplyBlur(VideoFrame frame, SegmentationMask mask, double strength)
    {
        var width = frame.Width;
        var height = frame.Height;
        var src = frame.BgraPixels;
        var stride = frame.Stride;

        // 1. Downscale by box-averaging BlurDownscale×BlurDownscale blocks (3 float channels).
        var f = BlurDownscale;
        var sw = (width + f - 1) / f;
        var sh = (height + f - 1) / f;
        var small = new float[sw * sh * 3];
        Parallel.For(0, sh, RowParallelism, sy =>
        {
            var y0 = sy * f;
            var y1 = Math.Min(height, y0 + f);
            for (var sx = 0; sx < sw; sx++)
            {
                var x0 = sx * f;
                var x1 = Math.Min(width, x0 + f);
                float b = 0, g = 0, r = 0;
                var n = 0;
                for (var y = y0; y < y1; y++)
                {
                    var p = y * stride + x0 * 4;
                    for (var x = x0; x < x1; x++, p += 4)
                    {
                        b += src[p]; g += src[p + 1]; r += src[p + 2];
                        n++;
                    }
                }
                var o = (sy * sw + sx) * 3;
                small[o] = b / n; small[o + 1] = g / n; small[o + 2] = r / n;
            }
        });

        // 2. Three separable box-blur passes ≈ a gaussian.
        var radius = BlurRadius(strength);
        var scratch = new float[small.Length];
        for (var pass = 0; pass < 3; pass++)
        {
            BoxBlurHorizontal(small, scratch, sw, sh, radius);
            BoxBlurVertical(scratch, small, sw, sh, radius);
        }

        // 3. Upsample: horizontal interpolation precomputed once per small row (and per mask row),
        //    then each output row lerps two precomputed rows — no per-pixel 2-D bilinear.
        var blurRows = UpsampleColorRows(small, sw, sh, width);
        var (b0, b1, bfy) = AxisLut(sh, height);
        var maskRows = UpsampleMaskRows(mask, width);
        var (m0, m1, mfy) = AxisLut(mask.Height, height);

        // 4. Blend: person (mask) keeps the source pixel, background takes the blur.
        var dst = new byte[src.Length];
        Parallel.For(0, height, RowParallelism, y =>
        {
            var ma = maskRows[m0[y]];
            var mb = maskRows[m1[y]];
            var mt = mfy[y];
            var ba = blurRows[b0[y]];
            var bb = blurRows[b1[y]];
            var bt = bfy[y];
            var row = y * stride;
            for (var x = 0; x < width; x++)
            {
                var p = row + x * 4;
                var m = ma[x] + (mb[x] - ma[x]) * mt;
                if (m >= 0.999f)
                {
                    dst[p] = src[p]; dst[p + 1] = src[p + 1]; dst[p + 2] = src[p + 2]; dst[p + 3] = src[p + 3];
                    continue;
                }
                if (m < 0f) m = 0f;
                var k = 1f - m;
                var c = x * 3;
                var vb = ba[c] + (bb[c] - ba[c]) * bt;
                var vg = ba[c + 1] + (bb[c + 1] - ba[c + 1]) * bt;
                var vr = ba[c + 2] + (bb[c + 2] - ba[c + 2]) * bt;
                dst[p] = ClampByte(src[p] * m + vb * k);
                dst[p + 1] = ClampByte(src[p + 1] * m + vg * k);
                dst[p + 2] = ClampByte(src[p + 2] * m + vr * k);
                dst[p + 3] = src[p + 3];
            }
        });

        return new VideoFrame(width, height, dst) { Epoch = frame.Epoch, IsOpaque = frame.IsOpaque };
    }

    /// <summary>Blur radius in downscaled pixels for a strength. Quadratic (s²), not linear, so the
    /// default 0.5 stays gentle (2 small px ≈ ~8–10 full px of perceived spread, ~3× lighter than
    /// the old linear 6) while only the top of the slider approaches the full range.</summary>
    internal static int BlurRadius(double strength)
    {
        var s = Math.Clamp(strength, 0, 1);
        return (int)Math.Round(MinBlurRadius + (MaxBlurRadius - MinBlurRadius) * s * s);
    }

    /// <summary>Row-parallel pixel loops, capped so the effect worker can't starve the compositor
    /// and encoder threads.</summary>
    private static readonly ParallelOptions RowParallelism = new()
    {
        MaxDegreeOfParallelism = Math.Clamp(Environment.ProcessorCount / 2, 1, 4),
    };

    private static byte ClampByte(float v)
    {
        var i = (int)(v + 0.5f);
        return (byte)(i < 0 ? 0 : i > 255 ? 255 : i);
    }

    /// <summary>For each destination index along one axis: the two source indices and the
    /// fractional weight — pixel-center aligned, edge-clamped (the same mapping a bilinear
    /// sample uses).</summary>
    private static (int[] I0, int[] I1, float[] F) AxisLut(int srcSize, int dstSize)
    {
        var i0 = new int[dstSize];
        var i1 = new int[dstSize];
        var fr = new float[dstSize];
        for (var d = 0; d < dstSize; d++)
        {
            var s = (d + 0.5f) * srcSize / dstSize - 0.5f;
            var a = (int)MathF.Floor(s);
            var t = s - a;
            i0[d] = Math.Clamp(a, 0, srcSize - 1);
            i1[d] = Math.Clamp(a + 1, 0, srcSize - 1);
            fr[d] = a < 0 ? 0f : a + 1 > srcSize - 1 ? 0f : t;
        }
        return (i0, i1, fr);
    }

    /// <summary>Each mask row, horizontally interpolated (and clamped to 0..1) to the frame width.</summary>
    private static float[][] UpsampleMaskRows(SegmentationMask mask, int width)
    {
        var (x0, x1, fx) = AxisLut(mask.Width, width);
        var rows = new float[mask.Height][];
        var v = mask.Values;
        for (var r = 0; r < mask.Height; r++)
        {
            var row = new float[width];
            var o = r * mask.Width;
            for (var x = 0; x < width; x++)
            {
                var a = v[o + x0[x]];
                var m = a + (v[o + x1[x]] - a) * fx[x];
                row[x] = m < 0f ? 0f : m > 1f ? 1f : m;
            }
            rows[r] = row;
        }
        return rows;
    }

    /// <summary>Each row of a 3-channel float image, horizontally interpolated to <paramref name="width"/>.</summary>
    private static float[][] UpsampleColorRows(float[] img, int w, int h, int width)
    {
        var (x0, x1, fx) = AxisLut(w, width);
        var rows = new float[h][];
        Parallel.For(0, h, RowParallelism, r =>
        {
            var row = new float[width * 3];
            var o = r * w;
            for (var x = 0; x < width; x++)
            {
                var a = (o + x0[x]) * 3;
                var b = (o + x1[x]) * 3;
                var t = fx[x];
                row[x * 3] = img[a] + (img[b] - img[a]) * t;
                row[x * 3 + 1] = img[a + 1] + (img[b + 1] - img[a + 1]) * t;
                row[x * 3 + 2] = img[a + 2] + (img[b + 2] - img[a + 2]) * t;
            }
            rows[r] = row;
        });
        return rows;
    }

    private static void BoxBlurHorizontal(float[] input, float[] output, int w, int h, int r)
    {
        var norm = 1f / (2 * r + 1);
        for (var y = 0; y < h; y++)
        {
            var row = y * w;
            for (var c = 0; c < 3; c++)
            {
                float sum = 0;
                for (var i = -r; i <= r; i++)
                    sum += input[(row + Math.Clamp(i, 0, w - 1)) * 3 + c];
                for (var x = 0; x < w; x++)
                {
                    output[(row + x) * 3 + c] = sum * norm;
                    var add = Math.Min(w - 1, x + r + 1);
                    var sub = Math.Max(0, x - r);
                    sum += input[(row + add) * 3 + c] - input[(row + sub) * 3 + c];
                }
            }
        }
    }

    private static void BoxBlurVertical(float[] input, float[] output, int w, int h, int r)
    {
        var norm = 1f / (2 * r + 1);
        for (var x = 0; x < w; x++)
        {
            for (var c = 0; c < 3; c++)
            {
                float sum = 0;
                for (var i = -r; i <= r; i++)
                    sum += input[(Math.Clamp(i, 0, h - 1) * w + x) * 3 + c];
                for (var y = 0; y < h; y++)
                {
                    output[(y * w + x) * 3 + c] = sum * norm;
                    var add = Math.Min(h - 1, y + r + 1);
                    var sub = Math.Max(0, y - r);
                    sum += input[(add * w + x) * 3 + c] - input[(sub * w + x) * 3 + c];
                }
            }
        }
    }




    private void LogFailureOnce(string reason)
    {
        if (_loggedFailure) return;
        _loggedFailure = true;
        AppLog.Write($"BackgroundEffectProcessor: passing webcam through — {reason}");
    }
}