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;
            for (var i = 3; i < buffer.Length; i += 4)
                buffer[i] = 255; // opaque black base

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

        var (scale, ox, oy) = StretchMath.UniformToFill(ew, eh, src.Width, src.Height);
        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;

        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>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).</summary>
    private static void BlitOverlay(
        byte[] dst, int dstW, int dstH, CompositorOptions options,
        VideoFrame overlay, int sx0, int sy0)
    {
        for (var y = 0; y < dstH; y++)
        {
            for (var x = 0; x < dstW; x++)
            {
                var sx = x + options.SourceRectX - sx0;
                var sy = y + options.SourceRectY - sy0;
                if (sx < 0 || sy < 0 || sx >= overlay.Width || sy >= overlay.Height) continue;
                var sample = StretchMath.SampleBgra(overlay.BgraPixels, overlay.Width, overlay.Height, sx, sy);
                BlendPixel(dst, (y * dstW + x) * 4, sample, 1f);
            }
        }
    }

    /// <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);
    }
}