LlamaCasty/Services/MediaVideoSourceManager.cs

using System;
using System.Collections.Generic;
using System.Threading.Tasks;
using System.Windows;
using System.Windows.Media;
using System.Windows.Media.Imaging;
using System.Windows.Threading;
using ytLive.Helpers;

namespace ytLive.Services;

/// <summary>
/// TASK 21: owns media-file decode sessions app-wide, refcounted by file path.
/// One file = at most one decoder, one shared WriteableBitmap; the last release
/// stops and disposes the source. Frames arrive on a worker thread and are
/// coalesced onto the UI dispatcher (at most one pending copy per session, using
/// the latest frame). Mirrors ScreenCaptureManager/CameraManager.
/// </summary>
public sealed class MediaVideoSourceManager : IDisposable
{
    private sealed class MediaSession
    {
        public string Key { get; }
        public IMediaFrameSource Source { get; }
        public Action<VideoFrame>? FrameHandler;
        public int RefCount;
        public bool Started;
        public WriteableBitmap? PreviewBitmap;
        public VideoFrame? LatestFrame;
        public bool FramePending;

        public MediaSession(string key, IMediaFrameSource source)
        {
            Key = key;
            Source = source;
            RefCount = 1;
        }
    }

    private readonly Func<string, IMediaFrameSource?> _sourceFactory;
    private readonly Dispatcher? _uiDispatcher;
    private readonly Func<string, bool>? _loopProvider;
    private readonly Dictionary<string, MediaSession> _sessions = new();
    private readonly object _gate = new();

    /// <summary>Raised on the UI thread when a media file's shared preview bitmap is first created.</summary>
    public event Action<string, WriteableBitmap>? PreviewBitmapChanged;

    /// <summary>Raised when a media decoder cannot be created (bad path).</summary>
    public event Action<string, string>? MediaFailed;

    public MediaVideoSourceManager(
        Func<string, IMediaFrameSource?> sourceFactory,
        Dispatcher? uiDispatcher = null,
        Func<string, bool>? loopProvider = null)
    {
        _sourceFactory = sourceFactory;
        _uiDispatcher = uiDispatcher;
        _loopProvider = loopProvider;
    }

    /// <summary>Increments the refcount for a file, starting the decoder the first time.</summary>
    public async Task<bool> AcquireAsync(string key)
    {
        if (string.IsNullOrWhiteSpace(key)) return false;

        MediaSession session;
        bool shouldStart;
        lock (_gate)
        {
            if (_sessions.TryGetValue(key, out var existing))
            {
                existing.RefCount++;
                shouldStart = false;
                session = existing;
            }
            else
            {
                IMediaFrameSource? source;
                try
                {
                    source = _sourceFactory(key);
                }
                catch (Exception ex)
                {
                    AppLog.Write($"MediaSource: creating decoder '{key}' failed: {ex.Message}");
                    MediaFailed?.Invoke(key, ex.Message);
                    return false;
                }
                if (source == null)
                {
                    MediaFailed?.Invoke(key, "No decoder for this file");
                    return false;
                }
                session = new MediaSession(key, source);
                session.FrameHandler = frame => OnFrameAvailable(session, frame);
                session.Source.FrameAvailable += session.FrameHandler;
                // Loop intent lives on the Source (per scene), so the caller supplies the
                // per-path answer ("any live reference loops"). Applied before Start so a
                // zero-length or instantly-EOF file cannot slip past a later assignment.
                session.Source.Looping = _loopProvider?.Invoke(key) ?? false;
                _sessions[key] = session;
                shouldStart = true;
            }
        }

        if (!shouldStart) return session.Started;

        try
        {
            await session.Source.StartAsync();
            session.Started = true;
            return true;
        }
        catch (Exception ex)
        {
            lock (_gate)
                _sessions.Remove(key);
            session.Source.FrameAvailable -= session.FrameHandler;
            AppLog.Write($"MediaSource: failed to start decoder '{key}': {ex.Message}");
            MediaFailed?.Invoke(key, ex.Message);
            await SafeStopAsync(session.Source);
            return false;
        }
    }

    public async Task ReleaseAsync(string key)
    {
        MediaSession? toStop = null;
        lock (_gate)
        {
            if (!_sessions.TryGetValue(key, out var session)) return;
            if (--session.RefCount > 0) return;
            _sessions.Remove(key);
            toStop = session;
        }

        if (toStop == null) return;
        toStop.Source.FrameAvailable -= toStop.FrameHandler;
        await SafeStopAsync(toStop.Source);
        toStop.PreviewBitmap = null;
    }

    /// <summary>Releases a file unconditionally (every ref) — used on layout reloads.</summary>
    public async Task ReleaseAllAsync(string key)
    {
        MediaSession? toStop = null;
        lock (_gate)
        {
            if (!_sessions.TryGetValue(key, out var session)) return;
            session.RefCount = 0;
            _sessions.Remove(key);
            toStop = session;
        }

        if (toStop == null) return;
        toStop.Source.FrameAvailable -= toStop.FrameHandler;
        await SafeStopAsync(toStop.Source);
        toStop.PreviewBitmap = null;
    }

    /// <summary>The most recent decoded frame for a file, or null before the first
    /// frame (or if the file has no session). The live compositor reads media from here.</summary>
    public VideoFrame? GetLatestFrame(string key)
    {
        lock (_gate)
            return _sessions.TryGetValue(key, out var session) ? session.LatestFrame : null;
    }

    /// <summary>True while a decode session exists for the file — the caller uses this to
    /// keep one acquire per path instead of stacking refcounts on every layout pass.</summary>
    public bool HasSession(string key)
    {
        lock (_gate)
            return _sessions.ContainsKey(key);
    }

    /// <summary>The files with a live session, for reconciling a layout load.</summary>
    public IReadOnlyList<string> GetActiveKeys()
    {
        lock (_gate)
            return new List<string>(_sessions.Keys);
    }

    /// <summary>Applies a loop change to a running session (the creator flipped the
    /// toggle). No-op when the file has no session.</summary>
    public void SetLooping(string key, bool looping)
    {
        lock (_gate)
        {
            if (_sessions.TryGetValue(key, out var session))
                session.Source.Looping = looping;
        }
    }

    /// <summary>The shared preview bitmap for a file (so a config created after the
    /// first frame can adopt the running bitmap), or null.</summary>
    public WriteableBitmap? GetPreviewBitmap(string key)
    {
        lock (_gate)
            return _sessions.TryGetValue(key, out var session) ? session.PreviewBitmap : null;
    }

    public void Dispose()
    {
        List<MediaSession> sessions;
        lock (_gate)
        {
            sessions = new List<MediaSession>(_sessions.Values);
            _sessions.Clear();
        }

        foreach (var session in sessions)
        {
            session.Source.FrameAvailable -= session.FrameHandler;
            _ = SafeStopAsync(session.Source);
        }
    }

    private static async Task SafeStopAsync(IMediaFrameSource source)
    {
        try
        {
            await source.StopAsync();
        }
        catch (Exception ex)
        {
            AppLog.Write($"MediaSource: stopping decoder failed: {ex.Message}");
        }
    }

    private bool TryGetActiveSession(MediaSession session)
    {
        lock (_gate)
            return _sessions.TryGetValue(session.Key, out var current) && ReferenceEquals(current, session);
    }

    private void OnFrameAvailable(MediaSession session, VideoFrame frame)
    {
        if (!TryGetActiveSession(session)) return;
        session.LatestFrame = frame;

        if (session.PreviewBitmap == null)
        {
            if (_uiDispatcher == null) return;
            if (_uiDispatcher.CheckAccess())
                EnsurePreviewBitmap(session);
            else
                _uiDispatcher.BeginInvoke(() =>
                {
                    if (TryGetActiveSession(session))
                        EnsurePreviewBitmap(session);
                }, DispatcherPriority.Render);
            return;
        }

        if (_uiDispatcher == null || _uiDispatcher.CheckAccess())
        {
            CopyFrame(session);
            return;
        }

        if (session.FramePending) return;
        session.FramePending = true;
        _uiDispatcher.BeginInvoke(() =>
        {
            session.FramePending = false;
            if (TryGetActiveSession(session) && session.PreviewBitmap != null)
                CopyFrame(session);
        }, DispatcherPriority.Render);
    }

    private void EnsurePreviewBitmap(MediaSession session)
    {
        if (session.PreviewBitmap != null || session.LatestFrame == null) return;
        var frame = session.LatestFrame;
        var bitmap = new WriteableBitmap(frame.Width, frame.Height, 96, 96, PixelFormats.Bgra32, null);
        bitmap.WritePixels(new Int32Rect(0, 0, frame.Width, frame.Height), frame.BgraPixels, frame.Stride, 0);
        session.PreviewBitmap = bitmap;
        PreviewBitmapChanged?.Invoke(session.Key, bitmap);
    }

    private void CopyFrame(MediaSession session)
    {
        var bitmap = session.PreviewBitmap;
        var frame = session.LatestFrame;
        if (bitmap == null || frame == null) return;
        if (frame.Width != bitmap.PixelWidth || frame.Height != bitmap.PixelHeight) return;
        bitmap.WritePixels(new Int32Rect(0, 0, frame.Width, frame.Height), frame.BgraPixels, frame.Stride, 0);
    }
}