LlamaCasty/Services/Compositor/SceneCompositor.cs

using ytLive.Models;

namespace ytLive.Services.Compositor;

/// <summary>
/// The output compositor: renders a scene into the encoder's master frame (tightly-packed
/// BGRA8 <see cref="VideoFrame"/>) exactly as the XAML preview renders it, minus the
/// editing chrome (SelectionOverlay, DimRects, badge, placeholder). The preview stays the
/// editing view; this is the output view — see TASKS.md "Ship step 1 — Scene compositor".
///
/// Frame sources are supplied by a <see cref="Func{SceneElement, VideoFrame}"/> resolver
/// (the caller maps webcam → DeviceId, images → AssetId, background → CaptureKey), keeping
/// the compositor pure and free of WPF and of the capture managers.
/// </summary>
public sealed class SceneCompositor
{
    /// <summary>
    /// Composite <paramref name="scene"/> into the tier's output frame. Layer order (back →
    /// front): live background (the scene's <c>IsBackground</c> source) → background image →
    /// visible elements (z-order = <c>Elements</c> order, mirroring the XAML DataTemplate) →
    /// branding flash → social bar (a global overlay; spans the master width at
    /// <paramref name="socialBarTop"/>). Transparent regions read opaque black.
    /// </summary>
    public VideoFrame Render(
        Scene scene,
        Func<SceneElement, VideoFrame?> frameFor,
        VideoFrame? flashFrame,
        CompositorOptions options,
        VideoFrame? socialBarFrame = null,
        int socialBarTop = 0,
        VideoFrame? staticBase = null,
        int split = -1)
    {
        if (scene == null) throw new ArgumentNullException(nameof(scene));
        if (frameFor == null) throw new ArgumentNullException(nameof(frameFor));
        if (options == null) throw new ArgumentNullException(nameof(options));
        if (options.SourceRectWidth <= 0 || options.SourceRectHeight <= 0)
            throw new ArgumentException("The source rect must be positive.", nameof(options));
        if (options.OutputWidth <= 0 || options.OutputHeight <= 0)
            throw new ArgumentException("The output size must be positive.", nameof(options));

        var cropW = options.SourceRectWidth;
        var cropH = options.SourceRectHeight;
        var buffer = new byte[cropW * cropH * 4];

        // staticBase != null: the caller pre-baked all static layers below `split`.
        // We copy it (never mutate the cached base) and composite elements[split..] on top.
        if (staticBase != null)
        {
            if (staticBase.Width != cropW || staticBase.Height != cropH)
                throw new ArgumentException("The static base must match the source rect.", nameof(staticBase));
            Buffer.BlockCopy(staticBase.BgraPixels, 0, buffer, 0, buffer.Length);
        }
        else
        {
            if (split < 0) split = scene.Elements.Count;

            var elements = scene.Elements;
            var liveCapture = elements.OfType<Source>().FirstOrDefault(s => s.IsBackground);
            var liveCaptureFrame = liveCapture != null && liveCapture.IsVisible ? frameFor(liveCapture) : null;

            // A master-sized live backdrop paints every pixel of the base — the
            // opaque-black pre-fill (2M byte writes) would be dead work. Capture
            // frames are opaque BGRA by contract (the WinRT path ignores alpha),
            // so the 1:1 row walk covers the full rect.
            var covered = liveCaptureFrame != null
                && liveCaptureFrame.Width == cropW && liveCaptureFrame.Height == cropH;
            if (!covered)
                for (var i = 3; i < buffer.Length; i += 4)
                    buffer[i] = 255; // opaque black base

            if (liveCaptureFrame != null)
                BlitContent(buffer, cropW, cropH, 0, 0, cropW, cropH, liveCaptureFrame, 1f, false, false);

            var background = elements.OfType<Source>().FirstOrDefault(s => s.Type == SourceType.Background);
            var backgroundFrame = background != null ? frameFor(background) : null;
            if (backgroundFrame != null)
                BlitContent(buffer, cropW, cropH, 0, 0, cropW, cropH, backgroundFrame, 1f, false, false);

            foreach (var element in elements)
            {
                if (!element.IsVisible) continue;
                switch (element)
                {
                    case Source { IsBackground: true }:
                    case Source { Type: SourceType.Background }:
                    case Source { Type: SourceType.TextOverlay }:
                        continue;
                }
                var fullFrame = frameFor(element);
                if (fullFrame == null) continue;
                BlitBakedElement(buffer, cropW, cropH, options, element, fullFrame);
            }
        }

        // Composite the above-split layers (dynamic + any static above them) per frame.
        for (var i = Math.Max(0, split); i < scene.Elements.Count; i++)
        {
            var element = scene.Elements[i];
            if (!element.IsVisible) continue;
            switch (element)
            {
                case Source { IsBackground: true }:
                case Source { Type: SourceType.Background }:
                case Source { Type: SourceType.TextOverlay }:
                    continue;
            }
            var frame = frameFor(element);
            if (frame == null) continue;
            BlitBakedElement(buffer, cropW, cropH, options, element, frame);
        }

        if (flashFrame != null)
            BlitOverlay(buffer, cropW, cropH, options, flashFrame, 0, 0);

        if (socialBarFrame != null && socialBarFrame.Width > 0 && socialBarFrame.Height > 0)
            BlitOverlay(buffer, cropW, cropH, options, socialBarFrame, 0, socialBarTop);

        return StretchMath.BilinearScale(
            new VideoFrame(cropW, cropH, buffer), options.OutputWidth, options.OutputHeight);
    }

    /// <summary>
    /// Render the static base of a scene: the opaque black layer, the live-capture
    /// background, the background image, and every <paramref name="split"/> z-layers
    /// below the split point (all of which are static by construction — see
    /// <see cref="SceneElement.Kind"/>). The result is a source-rect sized frame that
    /// <see cref="Render"/> can start from and composite the dynamic/above-split layers
    /// onto each tick. Returns null when there is no static layer at all.
    /// </summary>
    public static VideoFrame? BakeStaticBase(
        Scene scene,
        int split,
        Func<SceneElement, VideoFrame?> frameFor,
        Func<CompositorOptions> options,
        VideoFrame? socialBarFrame = null,
        int socialBarTop = 0,
        VideoFrame? flashFrame = null)
    {
        if (split <= 0) return null;
        var opts = options();
        if (opts.SourceRectWidth <= 0 || opts.SourceRectHeight <= 0) return null;

        var cropW = opts.SourceRectWidth;
        var cropH = opts.SourceRectHeight;
        var buffer = new byte[cropW * cropH * 4];
        for (var i = 3; i < buffer.Length; i += 4)
            buffer[i] = 255;

        var elements = scene.Elements;

        var liveCapture = elements.OfType<Source>().FirstOrDefault(s => s.IsBackground);
        if (liveCapture != null)
        {
            var liveCaptureFrame = liveCapture.IsVisible ? frameFor(liveCapture) : null;
            if (liveCaptureFrame != null && split > elements.IndexOf(liveCapture))
                BlitContent(buffer, cropW, cropH, 0, 0, cropW, cropH, liveCaptureFrame, 1f, false, false);
        }

        var background = elements.OfType<Source>().FirstOrDefault(s => s.Type == SourceType.Background);
        if (background != null)
        {
            var backgroundFrame = frameFor(background);
            if (backgroundFrame != null && split > elements.IndexOf(background))
                BlitContent(buffer, cropW, cropH, 0, 0, cropW, cropH, backgroundFrame, 1f, false, false);
        }

        for (var i = 0; i < split; i++)
        {
            var element = scene.Elements[i];
            if (!element.IsVisible) continue;
            switch (element)
            {
                case Source { IsBackground: true }:
                case Source { Type: SourceType.Background }:
                case Source { Type: SourceType.TextOverlay }:
                    continue;
            }
            var frame = frameFor(element);
            if (frame == null) continue;
            BlitBakedElement(buffer, cropW, cropH, opts, element, frame);
        }

        var baked = new VideoFrame(cropW, cropH, buffer);
        if (flashFrame != null)
            baked = Overlay(baked, opts, flashFrame, 0, 0);
        if (socialBarFrame != null && socialBarFrame.Width > 0 && socialBarFrame.Height > 0)
            baked = Overlay(baked, opts, socialBarFrame, 0, socialBarTop);
        return baked;
    }

    /// <summary>The compositor's dynamic/above-split path: start from an existing frame
    /// buffer and composite the z-layers from <paramref name="split"/> onward on top.
    /// When <paramref name="baseFrame"/> is null the caller must fall back to full render.</summary>
    public static VideoFrame CompositeLayers(
        VideoFrame baseFrame,
        Scene scene,
        int split,
        Func<SceneElement, VideoFrame?> frameFor,
        CompositorOptions options,
        VideoFrame? socialBarFrame = null,
        int socialBarTop = 0,
        VideoFrame? flashFrame = null)
    {
        var cropW = options.SourceRectWidth;
        var cropH = options.SourceRectHeight;
        // We must not mutate the cached base: copy it into a fresh buffer.
        var buffer = new byte[baseFrame.BgraPixels.Length];
        Buffer.BlockCopy(baseFrame.BgraPixels, 0, buffer, 0, buffer.Length);

        for (var i = split; i < scene.Elements.Count; i++)
        {
            var element = scene.Elements[i];
            if (!element.IsVisible) continue;
            switch (element)
            {
                case Source { IsBackground: true }:
                case Source { Type: SourceType.Background }:
                case Source { Type: SourceType.TextOverlay }:
                    continue;
            }
            var frame = frameFor(element);
            if (frame == null) continue;
            BlitBakedElement(buffer, cropW, cropH, options, element, frame);
        }
        return StretchMath.BilinearScale(new VideoFrame(cropW, cropH, buffer), options.OutputWidth, options.OutputHeight);
    }

    private static void BlitBakedElement(
        byte[] buffer, int cropW, int cropH, CompositorOptions options, SceneElement element, VideoFrame frame)
    {
        var ex = (float)(element.X - options.SourceRectX);
        var ey = (float)(element.Y - options.SourceRectY);
        var ew = (float)element.Width;
        var eh = (float)element.Height;
        var isRound = element.ClipShape == ClipShape.Round;
        BlitContent(buffer, cropW, cropH, ex, ey, ew, eh, frame,
            (float)element.Opacity, isRound, element.IsMirrored);
        if (element.HasBorder)
            DrawBorder(buffer, cropW, cropH, ex, ey, ew, eh, element, isRound);
    }

    private static VideoFrame Overlay(VideoFrame frame, CompositorOptions options, VideoFrame overlay, int sx0, int sy0)
    {
        var cropW = options.SourceRectWidth;
        var cropH = options.SourceRectHeight;
        var buffer = new byte[frame.BgraPixels.Length];
        Buffer.BlockCopy(frame.BgraPixels, 0, buffer, 0, buffer.Length);
        BlitOverlay(buffer, cropW, cropH, options, overlay, sx0, sy0);
        return new VideoFrame(cropW, cropH, buffer);
    }

    /// <summary>
    /// Render a cross-fade blend between two scenes. Used by <see cref="TransitionService"/>
    /// during fade/move transitions: the caller renders each scene individually, then blends
    /// the two results at the given alpha (0 = fully from, 1 = fully to).
    /// </summary>
    public static VideoFrame BlendFrames(VideoFrame from, VideoFrame to, float alpha)
    {
        if (from.Width != to.Width || from.Height != to.Height)
        {
            if (alpha <= 0.5f) return from;
            return to;
        }

        var count = from.BgraPixels.Length;
        var result = new byte[count];
        var inv = 1f - alpha;

        for (var i = 0; i < count; i += 4)
        {
            result[i]     = (byte)(from.BgraPixels[i]     * inv + to.BgraPixels[i]     * alpha);
            result[i + 1] = (byte)(from.BgraPixels[i + 1] * inv + to.BgraPixels[i + 1] * alpha);
            result[i + 2] = (byte)(from.BgraPixels[i + 2] * inv + to.BgraPixels[i + 2] * alpha);
            result[i + 3] = 255;
        }

        return new VideoFrame(from.Width, from.Height, result);
    }

    /// <summary>
    /// UniformToFill blit of a source frame into element rect (ex, ey, ew, eh), with
    /// straight-alpha source-over, optional round clip (true circle, hard edge) and
    /// horizontal mirror around the element center.
    /// </summary>
    private static void BlitContent(
        byte[] dst, int dstW, int dstH,
        float ex, float ey, float ew, float eh,
        VideoFrame src, float opacity, bool isRound, bool isMirror)
    {
        if (ew <= 0 || eh <= 0) return;

        var (scale, ox, oy) = StretchMath.UniformToFill(ew, eh, src.Width, src.Height);

        // 1:1 aligned fast path (the full-frame backdrop, master-sized overlays):
        // bilinear sampling at scale 1 with integer offsets is the identity, so the
        // blit is a row walk, not 2M float samples + rounds. Research-first (2026-09-03,
        // take-3 fix; libyuv's row-based BLEND_NONE pattern —
        // https://chromium.googlesource.com/libyuv/libyuv/): the 258ms/frame render
        // that time-lapsed take 3 lived in the general loop below.
        if (!isRound && !isMirror && scale == 1f && ox == 0f && oy == 0f
            && ex == MathF.Floor(ex) && ey == MathF.Floor(ey))
        {
            Blit1To1Rows(dst, dstW, dstH, (int)ex, (int)ey, src, opacity);
            return;
        }

        var drawnW = src.Width * scale;
        var drawnH = src.Height * scale;
        var radius = Math.Min(ew, eh) / 2f;
        var cx = ew / 2f;
        var cy = eh / 2f;

        var x0 = Math.Max(0, (int)Math.Floor(ex));
        var y0 = Math.Max(0, (int)Math.Floor(ey));
        var x1 = Math.Min(dstW - 1, (int)Math.Ceiling(ex + ew));
        var y1 = Math.Min(dstH - 1, (int)Math.Ceiling(ey + eh));
        if (x0 > x1 || y0 > y1) return;

        for (var y = y0; y <= y1; y++)
        {
            for (var x = x0; x <= x1; x++)
            {
                var px = x - ex; // element space
                var py = y - ey;
                if (px < ox || px > ox + drawnW || py < oy || py > oy + drawnH) continue;
                if (isRound)
                {
                    var dx = px - cx;
                    var dy = py - cy;
                    if (dx * dx + dy * dy > radius * radius) continue;
                }

                var sx = (px - ox) / scale;
                var sy = (py - oy) / scale;
                if (isMirror) sx = src.Width - 1 - sx;
                var sample = StretchMath.SampleBgra(src.BgraPixels, src.Width, src.Height, sx, sy);
                BlendPixel(dst, (y * dstW + x) * 4, sample, opacity);
            }
        }
    }

    /// <summary>Row walk for a 1:1 aligned blit at integer offset (ex, ey): per-pixel
    /// branch on source alpha — opaque rows copy straight, transparent pixels are
    /// skipped, partial alpha blends in integer fixed-point (libyuv's pattern). The
    /// master buffer's destination alpha is always opaque (black base paints 255 and
    /// every blend writes 255), so dst alpha is never read — matching BlendPixel's
    /// source-over result for opaque dst within ±1.</summary>
    private static void Blit1To1Rows(byte[] dst, int dstW, int dstH, int ex, int ey, VideoFrame src, float opacity)
    {
        var op = (int)(Math.Clamp(opacity, 0f, 1f) * 256);
        if (op <= 0) return;

        var x0 = Math.Max(0, ex);
        var y0 = Math.Max(0, ey);
        var x1 = Math.Min(dstW, ex + src.Width);
        var y1 = Math.Min(dstH, ey + src.Height);
        if (x0 >= x1 || y0 >= y1) return;

        var rowLen = x1 - x0;
        for (var y = y0; y < y1; y++)
        {
            var si = ((y - ey) * src.Width + (x0 - ex)) * 4;
            var di = (y * dstW + x0) * 4;
            if (op == 256) BlendRowOpaque(dst, di, src.BgraPixels, si, rowLen);
            else BlendRowWeighted(dst, di, src.BgraPixels, si, rowLen, op);
        }
    }

    private static void BlendRowOpaque(byte[] dst, int di, byte[] src, int si, int count)
    {
        for (var n = 0; n < count; n++)
        {
            var a = src[si + 3];
            if (a == 255)
            {
                dst[di] = src[si];
                dst[di + 1] = src[si + 1];
                dst[di + 2] = src[si + 2];
                dst[di + 3] = 255;
            }
            else if (a != 0)
            {
                var inv = 255 - a;
                dst[di] = (byte)((src[si] * a + dst[di] * inv + 127) / 255);
                dst[di + 1] = (byte)((src[si + 1] * a + dst[di + 1] * inv + 127) / 255);
                dst[di + 2] = (byte)((src[si + 2] * a + dst[di + 2] * inv + 127) / 255);
                dst[di + 3] = 255;
            }
            si += 4;
            di += 4;
        }
    }

    private static void BlendRowWeighted(byte[] dst, int di, byte[] src, int si, int count, int op)
    {
        for (var n = 0; n < count; n++)
        {
            var a = src[si + 3] * op >> 8;
            if (a >= 255)
            {
                dst[di] = src[si];
                dst[di + 1] = src[si + 1];
                dst[di + 2] = src[si + 2];
                dst[di + 3] = 255;
            }
            else if (a > 0)
            {
                var inv = 255 - a;
                dst[di] = (byte)((src[si] * a + dst[di] * inv + 127) / 255);
                dst[di + 1] = (byte)((src[si + 1] * a + dst[di + 1] * inv + 127) / 255);
                dst[di + 2] = (byte)((src[si + 2] * a + dst[di + 2] * inv + 127) / 255);
                dst[di + 3] = 255;
            }
            si += 4;
            di += 4;
        }
    }

    /// <summary>Centered OBS-style border stroke: rect ring (Traditional) or circle ring (Round).</summary>
    private static void DrawBorder(
        byte[] dst, int dstW, int dstH,
        float ex, float ey, float ew, float eh,
        SceneElement element, bool isRound)
    {
        if (!element.TryGetBorderColor(out var r, out var g, out var b)) return;
        var half = element.BorderWidth / 2f;
        var alpha = (float)(element.Opacity * element.BorderOpacity);
        if (half <= 0 || alpha <= 0) return;

        var cx = ex + ew / 2f;
        var cy = ey + eh / 2f;
        var radius = Math.Min(ew, eh) / 2f;

        var x0 = Math.Max(0, (int)Math.Floor(ex - half));
        var y0 = Math.Max(0, (int)Math.Floor(ey - half));
        var x1 = Math.Min(dstW - 1, (int)Math.Ceiling(ex + ew + half));
        var y1 = Math.Min(dstH - 1, (int)Math.Ceiling(ey + eh + half));

        for (var y = y0; y <= y1; y++)
        {
            for (var x = x0; x <= x1; x++)
            {
                float d = isRound
                    ? MathF.Sqrt((x - cx) * (x - cx) + (y - cy) * (y - cy)) - radius
                    : MathF.Max(MathF.Max(ex - x, x - (ex + ew)), MathF.Max(ey - y, y - (ey + eh)));
                if (MathF.Abs(d) <= half)
                    BlendPixel(dst, (y * dstW + x) * 4, (b, g, r, (byte)255), alpha);
            }
        }
    }

    /// <summary>1:1 copy of a master-sized overlay (flash, social bar) cropped to the
    /// active source rect. The overlay is positioned in master space by (sx0, sy0).
    /// Iterates only the intersection rect — the old loop walked every destination
    /// pixel to blit a ~64px strip (2M iterations of bounds checks per overlay).</summary>
    private static void BlitOverlay(
        byte[] dst, int dstW, int dstH, CompositorOptions options,
        VideoFrame overlay, int sx0, int sy0)
    {
        var dx0 = Math.Max(0, sx0 - options.SourceRectX);
        var dy0 = Math.Max(0, sy0 - options.SourceRectY);
        var dx1 = Math.Min(dstW, sx0 - options.SourceRectX + overlay.Width);
        var dy1 = Math.Min(dstH, sy0 - options.SourceRectY + overlay.Height);
        for (var y = dy0; y < dy1; y++)
        {
            var si = ((y + options.SourceRectY - sy0) * overlay.Width
                + dx0 + options.SourceRectX - sx0) * 4;
            BlendRowOpaque(dst, (y * dstW + dx0) * 4, overlay.BgraPixels, si, dx1 - dx0);
        }
    }

    /// <summary>Straight-alpha source-over blend; the frame's alpha (and the layer opacity) drives coverage.</summary>
    private static void BlendPixel(byte[] dst, int di, (byte B, byte G, byte R, byte A) src, float opacity)
    {
        var sa = src.A / 255f * opacity;
        if (sa <= 0) return;
        var da = dst[di + 3] / 255f;
        var outA = sa + da * (1 - sa);
        if (outA <= 0) return;
        dst[di] = (byte)Math.Round((src.B * sa + dst[di] * da * (1 - sa)) / outA);
        dst[di + 1] = (byte)Math.Round((src.G * sa + dst[di + 1] * da * (1 - sa)) / outA);
        dst[di + 2] = (byte)Math.Round((src.R * sa + dst[di + 2] * da * (1 - sa)) / outA);
        dst[di + 3] = (byte)Math.Round(outA * 255);
    }
}