LlamaCasty/Services/WebView2Manager.cs

using System;
using System.Collections.Generic;
using System.Drawing;
using System.Runtime.InteropServices;
using System.Threading.Tasks;
using System.Windows;
using System.Windows.Media;
using System.Windows.Media.Imaging;
using System.Windows.Threading;
using Microsoft.Web.WebView2.Core;
using Windows.Graphics.Capture;
using Windows.System;
using Windows.UI.Composition;
using ytLive.Helpers;
using ytLive.Models;

namespace ytLive.Services;

/// <summary>
/// Owns the composition-based web sources (one per widget). A session renders
/// its widget through a <see cref="CoreWebView2CompositionController"/> whose
/// visual tree lives off-screen on the WPF UI thread (OBS web-source model:
/// the page renders at the MASTER CANVAS size 1920×1080, stable, never tracked);
/// the controller's root visual is captured by a <see cref="WebCaptureFrameSource"/>
/// at Chromium's own presentation cadence. The preview is crop-copied to a shared
/// WriteableBitmap on the UI thread (dispatcher-coalesced latest-wins, mirror of
/// ScreenCaptureManager); the FramePump reads <see cref="GetLatestFrame"/>.
///
/// The composition stack needs a DispatcherQueue on this thread (CsWinRT does not
/// project <c>CreateOnCurrentThread</c> on the 19041 SDK — only the dedicated
/// variant — so it is installed via a CoreMessaging P/Invoke, the documented
/// fallback the WebView2CompositionControl / webview_windows embedders use).
/// Everything composition-side therefore runs on the WPF UI thread; only the
/// frame conversion and alpha-bbox scan run on the capture worker.
///
/// Seam: an internal <c>Func&lt;string, IScreenCaptureSource&gt;</c> source
/// factory substitutes the whole WebView2+item creation in tests — the manager
/// then owns the session/coalesce/crop/preview plumbing exactly as in production.
/// </summary>
public sealed class WebView2Manager : IDisposable
{
    private readonly Dispatcher _dispatcher;
    private readonly Func<string, IScreenCaptureSource>? _sessionSourceFactory;
    private readonly Dictionary<string, WebSourceSession> _sessions = new();
    private Task<CoreWebView2Environment>? _environmentTask;
    private DispatcherQueueController? _dispatcherQueueController;
    private Windows.UI.Composition.Compositor? _compositor;

    public event Action<string, WriteableBitmap>? PreviewBitmapChanged;

    private sealed class WebSourceSession
    {
        public required string Key { get; init; }
        public CoreWebView2? Web;
        public CoreWebView2CompositionController? Controller;
        public IScreenCaptureSource? Source;
        public Action<VideoFrame>? FrameHandler;
        public VideoFrame? LatestFrame;
        public WriteableBitmap? PreviewBitmap;

        // Buffer reuse (2026-09-04 lesson): the preview crop copy reuses one
        // scratch array while the widget's extent holds — no per-frame LOH mint.
        public byte[]? CropScratch;
        public bool Disposed;
        public bool Initialized;
        public bool FramePending;
        public bool FirstCaptureLogged;
        public int WidgetLogsRemaining;
    }

    public WebView2Manager(Dispatcher dispatcher)
        : this(dispatcher, null)
    {
    }

    /// <summary>Test seam: the source factory replaces WebView2+composition
    /// creation, so the session/coalesce/crop/preview plumbing is testable
    /// without the runtime.</summary>
    internal WebView2Manager(Dispatcher dispatcher, Func<string, IScreenCaptureSource>? sessionSourceFactory)
    {
        _dispatcher = dispatcher;
        _sessionSourceFactory = sessionSourceFactory;
    }

    public void Register(Source source)
    {
        if (_sessions.ContainsKey(source.Id)) return;
        _sessions[source.Id] = new WebSourceSession { Key = source.Id };
        _ = InitializeAsync(source, _sessions[source.Id]);
    }

    public void Unregister(string sourceId) => RemoveSession(sourceId);

    public void OnWebUriChanged(Source source)
    {
        if (!_sessions.TryGetValue(source.Id, out var session)) return;
        if (session.Disposed || !session.Initialized) return;

        if (string.IsNullOrEmpty(source.WebUri))
        {
            session.LatestFrame = null;
            session.Web?.NavigateToString(
                "<html><body style='background:transparent;margin:0'></body></html>");
            return;
        }

        Navigate(source.WebUri, session);
    }

    public VideoFrame? GetLatestFrame(string sourceId)
    {
        return _sessions.TryGetValue(sourceId, out var session) ? session.LatestFrame : null;
    }

    // Injected BEFORE the document is parsed (AddScriptToExecuteOnDocumentCreatedAsync) so
    // the widget page's own CSS/scripts cannot repaint html/body opaque — the reason the
    // late NavigationStarting/NavigationCompleted ExecuteScriptAsync injection lost the
    // fight (page styles ran first). A transparent capture REQUIRES the transparency in
    // place before the page paints, on the composition surface too.
    //
    // The OLD inline `element.style.background='transparent'` form lost to any page CSS:
    // a widget that paints a background-COLOR on a container (or html/body with stronger
    // specificity) after connect re-opaques the capture → the "black opaque box" in the
    // recording while the early frames were transparent. This is a solved problem in the
    // overlay ecosystem — OBS browser sources use a Custom CSS override and the decade-
    // validated formula for arbitrary pages is a `!important` background-color wipe
    // (OBS Forums 2016 "body { background-color: rgba(0,0,0,0) !important }",
    // https://obsproject.com/forum/threads/translucent-transparent-browser-source.59549/,
    // and the div-level variant for stubborn widgets, woahtech.com OBS custom-CSS guide).
    // A `<style>` node injected before parse, with `!important`, outranks every page rule
    // (only an author !important beats a later document-order !important; ours runs first).
    // Only background-COLOR was wiped at first — the 15:34 take showed the widget's own
    // full-canvas backdrop surviving as the box's black void: that paint is a CSS
    // background-IMAGE (gradient/backdrop), and the OBS fix for the same symptom is
    // `background: none !important` / background-image none, not color alone
    // (obsproject/obs-studio#6659 — "set the CSS for html and body to background: none
    // !important"). Wipe background-image too; HTML overlay art (the widget's content) is
    // <img>/DOM/CSS shapes and survives — only container decorations are removed, which is
    // OBS-browser-source semantics for a transparent overlay.
    internal const string TransparentBackgroundScript =
        "(function(){" +
        "if(document.getElementById('ytl-transparent-bg'))return;" +
        "var s=document.createElement('style');" +
        "s.id='ytl-transparent-bg';" +
        "s.appendChild(document.createTextNode(" +
        "'html,body{background-color:transparent!important;background-image:none!important;margin:0!important;overflow:hidden!important;}'" +
        "'html,body,html *{background-color:transparent!important;background-image:none!important;}'" +
        "));" +
        "(document.head||document.documentElement).appendChild(s);" +
        "})();";

    private async Task InitializeAsync(Source source, WebSourceSession session)
    {
        try
        {
            IScreenCaptureSource sourceImpl;
            if (_sessionSourceFactory != null)
            {
                // Test path: the factory hands over a fully-configured source; the
                // branch has no awaits, so Register() completes the session inline.
                sourceImpl = _sessionSourceFactory(source.Id);
            }
            else
            {
                var controller = await CreateControllerAsync();
                if (controller == null)
                {
                    RemoveSession(source.Id);
                    return;
                }

                session.Controller = controller;
                session.Web = controller.CoreWebView2;
                if (session.Web == null)
                {
                    RemoveSession(source.Id);
                    return;
                }

                // Pre-page hook: runs after the global object exists but BEFORE the document
                // parses / any page script runs — must complete before navigating away from
                // about:blank for it to apply to the widget document.
                await session.Web.AddScriptToExecuteOnDocumentCreatedAsync(TransparentBackgroundScript);

                session.Web.NavigationStarting += (_, _) =>
                {
                    _ = session.Web?.ExecuteScriptAsync(TransparentBackgroundScript);
                };

                session.Web.NavigationCompleted += (_, e) =>
                {
                    _ = session.Web?.ExecuteScriptAsync(TransparentBackgroundScript);
                    if (e.IsSuccess && session.Web?.Source is { } src && src != "about:blank")
                    {
                        // Log-only diagnostic on the REAL widget document (a few post-paint
                        // frames so styles/backgrounds settle). The first-capture dump only
                        // ever saw about:blank; the old PNG widget dumps are gone with the
                        // PNG path — alpha stats now ride the frame-driven channel.
                        session.WidgetLogsRemaining = 3;
                        var rearm = session;
                        _ = Task.Run(async () =>
                        {
                            await Task.Delay(TimeSpan.FromSeconds(30));
                            if (!rearm.Disposed) rearm.WidgetLogsRemaining = 3;
                        });
                    }
                };

                session.Web.NavigateToString(
                    "<html><body style='background:transparent;margin:0'></body></html>");

                if (!string.IsNullOrEmpty(source.WebUri))
                    Navigate(source.WebUri, session);

                session.Initialized = true;

                sourceImpl = new WebCaptureFrameSource(source.Id, CreateCaptureItem(controller));
            }

            session.Source = sourceImpl;
            session.FrameHandler = frame => OnFrameAvailable(session, frame);
            sourceImpl.FrameAvailable += session.FrameHandler;
            await sourceImpl.StartAsync();
            session.Initialized = true;
        }
        catch (Exception ex)
        {
            AppLog.Write($"WebView2Manager: init failed for '{source.Name}': {ex.Message}");
            RemoveSession(source.Id);
        }
    }

    // The environment MUST exist before the controller is created — it's created once
    // and shared by every session (one browser process). These flags (the
    // Electron-class embedder answer to Chromium throttling hidden pages — the moment
    // the app window is unfocused or covered, rAF parks and timers clamp to ~1s) keep
    // the renderer at full rate: --disable-backgrounding-occluded-windows,
    // --disable-renderer-backgrounding, and disable the native-window occlusion
    // detector that marks the page hidden.
    private Task<CoreWebView2Environment> GetEnvironmentAsync()
    {
        if (_environmentTask != null) return _environmentTask;
        var options = new CoreWebView2EnvironmentOptions
        {
            AdditionalBrowserArguments =
                "--disable-backgrounding-occluded-windows --disable-renderer-backgrounding --disable-features=CalculateNativeWinOcclusion",
        };
        return _environmentTask = CoreWebView2Environment.CreateAsync(null, null, options);
    }

    /// <summary>Creates the composition controller off-screen under the main
    /// window (the parent HWND is required, the visual never mounts a window
    /// target). Set up is copied from the WebView2CompositionControl /
    /// webview_windows embedders: raw-pixel bounds at the master canvas,
    /// fixed rasterization scale, transparent default background.</summary>
    private async Task<CoreWebView2CompositionController?> CreateControllerAsync()
    {
        var mainWindow = Application.Current?.MainWindow;
        var hwnd = mainWindow == null
            ? IntPtr.Zero
            : new System.Windows.Interop.WindowInteropHelper(mainWindow).Handle;
        if (hwnd == IntPtr.Zero)
        {
            AppLog.Write("WebView2Manager: main window handle unavailable — web source not created");
            return null;
        }

        var env = await GetEnvironmentAsync();
        if (_compositor == null) _compositor = CreateCompositorForCurrentThread();

        var controller = await env.CreateCoreWebView2CompositionControllerAsync(hwnd);
        controller.Bounds = new Rectangle(0, 0, 1920, 1080);
        controller.BoundsMode = CoreWebView2BoundsMode.UseRawPixels;
        controller.ShouldDetectMonitorScaleChanges = false;
        controller.RasterizationScale = 1.0;
        controller.IsVisible = true;
        controller.DefaultBackgroundColor = System.Drawing.Color.Transparent;
        return controller;
    }

    private void Navigate(string url, WebSourceSession session)
    {
        try { session.Web?.Navigate(url); }
        catch (Exception ex) { AppLog.Write($"WebView2Manager: navigate failed: {ex.Message}"); }
    }

    // The visual tree webview_windows wires (offscreen mode): a root container
    // visual sized to the master canvas (the captured subject — the Flutter
    // plugin calls GraphicsCaptureItem.CreateFromVisual on this root), with the
    // controller's root visual target set to a child that fills the root via
    // RelativeSizeAdjustment. BoundsMode raw pixels + RasterizationScale 1.0 keep
    // the capture at exactly 1920×1080 regardless of monitor scale.
    private GraphicsCaptureItem CreateCaptureItem(CoreWebView2CompositionController controller)
    {
        var compositor = _compositor!;
        var root = compositor.CreateContainerVisual();
        root.Size = new System.Numerics.Vector2(1920f, 1080f);
        root.IsVisible = true;

        var webViewVisual = compositor.CreateContainerVisual();
        webViewVisual.RelativeSizeAdjustment = new System.Numerics.Vector2(1f, 1f);
        root.Children.InsertAtTop(webViewVisual);

        controller.RootVisualTarget = webViewVisual;
        return GraphicsCaptureItem.CreateFromVisual(root);
    }

    // CsWinRT does not project DispatcherQueueController.CreateOnCurrentThread
    // on the 19041 SDK (only CreateOnDedicatedThread [CS0117, verified by probe]);
    // the documented fallback — "how to create a DispatcherQueue on the current
    // thread" (learn.microsoft.com CoreMessaging docs) — is CreateDispatcherQueueController
    // from CoreMessaging.dll, wrapped via FromAbi (the CaptureInterop pattern). One
    // per UI thread, then new Compositor() (which requires the DQ) works here.
    private Windows.UI.Composition.Compositor CreateCompositorForCurrentThread()
    {
        var options = new DispatcherQueueOptions
        {
            DwSize = Marshal.SizeOf<DispatcherQueueOptions>(),
            ThreadType = 2,    // DQTYPE_THREAD_CURRENT
            ApartmentType = 2, // DQTAT_COM_STA — the WPF UI thread is STA
        };
        var hr = CreateDispatcherQueueController(options, out var controllerPtr);
        if (hr != 0)
            throw Marshal.GetExceptionForHR(hr)!;
        _dispatcherQueueController = DispatcherQueueController.FromAbi(controllerPtr);
        return new Windows.UI.Composition.Compositor();
    }

    [StructLayout(LayoutKind.Sequential)]
    private struct DispatcherQueueOptions
    {
        public int DwSize;
        public int ThreadType;
        public int ApartmentType;
    }

    [DllImport("CoreMessaging.dll")]
    private static extern int CreateDispatcherQueueController(
        DispatcherQueueOptions options, out IntPtr dispatcherQueueController);

    private void OnFrameAvailable(WebSourceSession session, VideoFrame frame)
    {
        if (session.Disposed) return;
        session.LatestFrame = frame;

        if (_dispatcher.CheckAccess())
        {
            UpdatePreview(session);
            return;
        }
        if (session.FramePending) return;
        session.FramePending = true;
        _dispatcher.BeginInvoke(() =>
        {
            session.FramePending = false;
            if (!session.Disposed)
                UpdatePreview(session);
        }, DispatcherPriority.Render);
    }

    private void UpdatePreview(WebSourceSession session)
    {
        var frame = session.LatestFrame;
        if (frame == null) return;

        var (cropX, cropY, cropW, cropH) = frame.CropBounds ?? (0, 0, frame.Width, frame.Height);
        if (cropW < 1 || cropH < 1)
        {
            cropX = 0; cropY = 0;
            cropW = frame.Width; cropH = frame.Height;
        }

        // The preview is sized to the widget's alpha bounding box (the compositor
        // gets the FULL canvas + CropBounds; the crop is only for the preview).
        // A fresh box (widget resizes / first content paints) mints a new bitmap
        // once — WPF re-renders a reused mutable bitmap after WritePixels, so
        // later frames only re-copy into the same one.
        if (session.PreviewBitmap == null
            || session.PreviewBitmap.PixelWidth != cropW
            || session.PreviewBitmap.PixelHeight != cropH)
        {
            session.PreviewBitmap = new WriteableBitmap(cropW, cropH, 96, 96, PixelFormats.Bgra32, null);
            PreviewBitmapChanged?.Invoke(session.Key, session.PreviewBitmap);
        }

        var need = cropW * cropH * 4;
        var scratch = session.CropScratch is { } s && s.Length >= need
            ? s
            : (session.CropScratch = new byte[need]);
        var stride = frame.Stride;
        if (cropW == frame.Width && cropH == frame.Height && cropX == 0 && cropY == 0)
            Buffer.BlockCopy(frame.BgraPixels, 0, scratch, 0, need);
        else
            for (var r = 0; r < cropH; r++)
                Buffer.BlockCopy(frame.BgraPixels, (cropY + r) * stride + cropX * 4, scratch, r * cropW * 4, cropW * 4);

        session.PreviewBitmap.WritePixels(new Int32Rect(0, 0, cropW, cropH), scratch, cropW * 4, 0);

        // alpha[min/max/mean/zero%] — the one number that decides whether a widget's
        // capture really has transparent margins (0% = opaque, ~100% if the page
        // painted transparent-only). Sampled every 16th pixel like the original log.
        if (!session.FirstCaptureLogged)
        {
            session.FirstCaptureLogged = true;
            AppLog.Write(
                $"WebView2Manager: first capture for '{session.Key}' {frame.Width}x{frame.Height} " +
                $"{AlphaStats(frame.BgraPixels, frame.Width, frame.Height, stride)} " +
                $"contentBounds=({cropX},{cropY},{cropW},{cropH})");
        }

        if (session.WidgetLogsRemaining > 0)
        {
            var dumpNo = 3 - session.WidgetLogsRemaining + 1;
            session.WidgetLogsRemaining--;
            AppLog.Write(
                $"WebView2Manager: widget capture [{dumpNo}/3] for '{session.Key}' {frame.Width}x{frame.Height} " +
                $"{AlphaStats(frame.BgraPixels, frame.Width, frame.Height, stride)} " +
                $"contentBounds=({cropX},{cropY},{cropW},{cropH})");
        }
    }

    private static string AlphaStats(byte[] pixels, int pixW, int pixH, int stride)
    {
        int minA = 255, maxA = 0;
        long sumA = 0, nA = 0, zeroA = 0;
        for (var y = 0; y < pixH; y += 16)
        {
            var row = y * stride;
            for (var x = 0; x < pixW; x += 16)
            {
                var a = pixels[row + x * 4 + 3];
                if (a < minA) minA = a;
                if (a > maxA) maxA = a;
                sumA += a;
                if (a == 0) zeroA++;
                nA++;
            }
        }
        return $"alpha[min={minA},max={maxA},mean={sumA / (double)nA:F1},zero={100.0 * zeroA / nA:F1}%]";
    }

    private void RemoveSession(string sourceId)
    {
        if (!_sessions.TryGetValue(sourceId, out var session)) return;
        session.Disposed = true;

        var src = session.Source;
        if (src != null)
        {
            if (session.FrameHandler != null)
                src.FrameAvailable -= session.FrameHandler;
            _ = SafeStopAsync(src);
        }
        try { session.Controller?.Close(); } catch { }
        session.Web = null;
        _sessions.Remove(sourceId);
    }

    private static async Task SafeStopAsync(IScreenCaptureSource source)
    {
        try { await source.StopAsync(); }
        catch (Exception ex) { AppLog.Write($"WebView2Manager: stopping web capture failed: {ex.Message}"); }
    }

    public void Dispose()
    {
        foreach (var id in _sessions.Keys.ToList())
            RemoveSession(id);
    }
}