ADOFAIRenderer/Util/VideoBGPlayer.cs

using System;
using System.Diagnostics.CodeAnalysis;
using System.IO;
using ADOFAIRenderer.UI;
using ADOFAIRenderer.Util.FFmpeg;
using MelonLoader;
using UnityEngine;
using FFmpeg.AutoGen;

namespace ADOFAIRenderer.Util;

/// <summary>
/// Plays a custom level's background video with FFmpeg (FFmpegNativeDecoder) in place of Unity's VideoPlayer while a
/// render is active, so the video frame shown always matches the render clock. Attached to the scnGame.videoBG object
/// by scnGamePatch.ReplaceVideo, which disables the native VideoPlayer; removed again by that patch when no render is
/// active.
/// <para>
/// Every LateUpdate it computes the video time from the conductor (the same formula scrVfxPlus uses to start and
/// position the native video), decodes frames up to that time and draws the newest one either into the VideoPlayer's
/// target RenderTexture or onto a quad at the far plane of the background camera.
/// Decoded YUV frames are converted to RGBA on the GPU (ColorConversion.compute, kernels CSMain_NV12ToRGBA /
/// CSMain_YUV420PToRGBA); other pixel formats, or a missing shader, use the decoder's RGBA output (CPU conversion).
/// Main thread only.
/// </para>
/// </summary>
// ReSharper disable once InconsistentNaming
internal class VideoBGPlayer : MonoBehaviour
{
    // Property IDs of the compute shader parameters: input planes, output texture, frame size in pixels,
    // ColorRange (1 = full range, 0 = limited range) and ColorMatrix (FFmpegNativeDecoder.ColorMatrix).
    private static readonly int YTex = Shader.PropertyToID("_YTex");
    private static readonly int UVTex = Shader.PropertyToID("_UVTex");
    private static readonly int Result = Shader.PropertyToID("_Result");
    private static readonly int Width = Shader.PropertyToID("Width");
    private static readonly int Height = Shader.PropertyToID("Height");
    private static readonly int UTex = Shader.PropertyToID("_UTex");
    private static readonly int VTex = Shader.PropertyToID("_VTex");
    private static readonly int ColorRange = Shader.PropertyToID("ColorRange");
    private static readonly int ColorMatrix = Shader.PropertyToID("ColorMatrix");

    /// <summary>
    /// Conversion path of the current video: <c>None</c> = CPU RGBA texture, <c>YUV420P</c> = separate Y, U, V planes,
    /// <c>NV12</c> = Y plane and interleaved UV plane.
    /// </summary>
    [SuppressMessage("ReSharper", "InconsistentNaming")]
    private enum GPUMode
    {
        None,
        YUV420P,
        NV12
    }

    /// <summary>Decoder of the current video; null when not initialized or after CleanUp.</summary>
    private FFmpegNativeDecoder? _decoder;
        
    // CPU fallback rendering
    // ReSharper disable once InconsistentNaming
    private Texture2D? _textureRGBA;

    // GPU conversion textures
    private Texture2D? _textureY;
    private Texture2D? _textureU;
    private Texture2D? _textureV;
    private Texture2D? _textureUV;

    // Compute Shader variables
    private ComputeShader? _computeShader;
    private RenderTexture? _renderTexture;
    private int _kernelIndex = -1;

    /// <summary>Presentation time in seconds of the frame currently shown (video time, 0 = start of the file).</summary>
    private double _currentFrameTimestamp;
    /// <summary>Video time in seconds from which the next frame is decoded (current timestamp + 1 / FPS).</summary>
    private double _nextFrameTimestamp;
    /// <summary>Conversion path selected by Init.</summary>
    private GPUMode _gpuMode = GPUMode.None;
    /// <summary>Target RenderTexture of the native VideoPlayer; null when the video is shown on the background quad.</summary>
    private RenderTexture? _targetTexture;
    /// <summary>Whether the level loops the video (miscSettings "loopVideo").</summary>
    private bool _isLooping;
    /// <summary>True after a successful Init; set to false by CleanUp or a playback error (playback then stops).</summary>
    private bool _isInitialized;

    /// <summary>Background quad created by SetupBackgroundQuad when there is no target texture; hidden until the first frame.</summary>
    private GameObject? _quad;
    /// <summary>Material of the background quad; its main texture is the converted frame.</summary>
    private Material? _material;

    /// <summary>
    /// Releases any previous video and opens <paramref name="videoPath"/>. Allocates the plane textures for the GPU path
    /// of the decoder's pixel format (chroma planes are (width + 1) / 2 x (height + 1) / 2), or an RGBA texture for the
    /// CPU path. Errors are logged; the player then stays uninitialized and shows nothing.
    /// </summary>
    /// <param name="videoPath">Full path of the video file.</param>
    /// <param name="targetTexture">RenderTexture the native VideoPlayer renders into, or null to draw on a background quad.</param>
    /// <param name="loop">Whether the video loops.</param>
    /// <param name="targetCamera">Camera the background quad is attached to when <paramref name="targetTexture"/> is null.</param>
    /// <param name="videoBGGo">VideoPlayer object; its layer is given to the background quad.</param>
    // ReSharper disable once InconsistentNaming
    public void Init(string videoPath, RenderTexture? targetTexture, bool loop, Camera? targetCamera, GameObject videoBGGo)
    {
        CleanUp();

        _targetTexture = targetTexture;
        _isLooping = loop;

        if (_targetTexture == null && targetCamera != null)
        {
            SetupBackgroundQuad(targetCamera, videoBGGo);
        }

        if (!File.Exists(videoPath))
        {
            Melon<RenderMod>.Logger.Warning($"[FFmpegVideoBGPlayer] Video file does not exist: {videoPath}");
            return;
        }

        try
        {
            _decoder = new FFmpegNativeDecoder(videoPath);
            _computeShader = ModAssets.ColorConversionShader;

            var w = _decoder.Width;
            var h = _decoder.Height;
            var pixFmt = _decoder.PixelFormat;

            var gpuInitSuccess = false;

            if (_computeShader != null)
            {
                switch (pixFmt)
                {
                    case AVPixelFormat.AV_PIX_FMT_NV12:
                        try
                        {
                            _kernelIndex = _computeShader.FindKernel("CSMain_NV12ToRGBA");
                            
                            _textureY = new Texture2D(w, h, TextureFormat.R8, false);
                            _textureUV = new Texture2D((w + 1) / 2, (h + 1) / 2, TextureFormat.RG16, false);
                            ConfigureTexture(_textureY);
                            ConfigureTexture(_textureUV);

                            _gpuMode = GPUMode.NV12;
                            gpuInitSuccess = true;
                            Melon<RenderMod>.Logger.Msg("[FFmpegVideoBGPlayer] Initialized Compute Shader NV12 mode successfully");
                        }
                        catch
                        {
                            // ignored
                        }

                        break;
                    case AVPixelFormat.AV_PIX_FMT_YUV420P:
                    case AVPixelFormat.AV_PIX_FMT_YUVJ420P:
                        try
                        {
                            _kernelIndex = _computeShader.FindKernel("CSMain_YUV420PToRGBA");
                            
                            _textureY = new Texture2D(w, h, TextureFormat.R8, false);
                            _textureU = new Texture2D((w + 1) / 2, (h + 1) / 2, TextureFormat.R8, false);
                            _textureV = new Texture2D((w + 1) / 2, (h + 1) / 2, TextureFormat.R8, false);
                            ConfigureTexture(_textureY);
                            ConfigureTexture(_textureU);
                            ConfigureTexture(_textureV);

                            _gpuMode = GPUMode.YUV420P;
                            gpuInitSuccess = true;
                            Melon<RenderMod>.Logger.Msg("[FFmpegVideoBGPlayer] Initialized Compute Shader YUV420P mode successfully");
                        }
                        catch
                        {
                            // ignored
                        }

                        break;
                    default:
                        Melon<RenderMod>.Logger.Msg("[FFmpegVideoBGPlayer] Cannot use Compute Shader for colorspace conversion.");
                        break;
                }
            }

            if (gpuInitSuccess && _computeShader != null)
            {
                _renderTexture = new RenderTexture(w, h, 0, RenderTextureFormat.ARGB32)
                {
                    enableRandomWrite = true
                };
                _renderTexture.Create();
            }
            else
            {
                _textureRGBA = new Texture2D(w, h, TextureFormat.RGBA32, false);
                ConfigureTexture(_textureRGBA);
                _gpuMode = GPUMode.None;
                Melon<RenderMod>.Logger.Warning("[FFmpegVideoBGPlayer] GPU acceleration unsupported or failed to load, falling back to CPU (RGBA).");
            }

            _currentFrameTimestamp = 0.0;
            _nextFrameTimestamp = 0.0;
            _isInitialized = true;

            Melon<RenderMod>.Logger.Msg($"[FFmpegVideoBGPlayer] Successfully initialized - Resolution: {w}x{h}, FPS: {_decoder.FPS:F2}");
        }
        catch (Exception ex)
        {
            Melon<RenderMod>.Logger.Error($"[FFmpegVideoBGPlayer] Error during initialization: {ex}");
            CleanUp();
        }
    }

    /// <summary>Sets clamp wrapping and bilinear filtering.</summary>
    private static void ConfigureTexture(Texture2D tex)
    {
        tex.wrapMode = TextureWrapMode.Clamp;
        tex.filterMode = FilterMode.Bilinear;
    }

    /// <summary>
    /// Main thread, every frame after the song has started. Computes the video time, seeks when it moved backwards or more
    /// than 1 s ahead of the shown frame, decodes frames until the next frame lies after that time, and draws the last
    /// decoded frame. Before the video starts (negative time) the target texture is cleared to black and the quad hidden.
    /// When decoding fails at the end of the file, a looping video seeks to 0; otherwise the last frame stays.
    /// </summary>
    private void LateUpdate()
    {
        if (!_isInitialized || _decoder == null) return;

        var cond = scrConductor.instance;
        if (!cond || !cond.hasSongStarted) return;
            
        // Song position in seconds of song time, as scrConductor computes songposition_minusi from dspTime with the input
        // offset calibration_i = 0 (the render sets the input offset to 0), minus the separate countdown length.
        // scrVfxPlus starts the native video at songposition_minusi - countdown + vidOffset (scrVfxPlus.cs).
        var pitch = cond.song.pitch;
        var separateCountdownOffset = cond.separateCountdownTime ? 
            cond.crotchetAtStart * cond.adjustedCountdownTicks : 0.0;
        var songTime = (cond.dspTime - cond.dspTimeSong) * pitch - cond.addoffset - separateCountdownOffset;

        var vidOffset = scrVfxPlus.instance ? scrVfxPlus.instance.vidOffset : 0.0;
        var targetVideoTime = songTime + vidOffset;

        // When looping is enabled, the playback position is the elapsed time wrapped into [0, Duration).
        // When it is disabled, the position keeps increasing and the last frame stays on screen after the end.
        var videoDuration = _decoder.Duration;
        if (_isLooping && targetVideoTime >= 0 && videoDuration > 0)
        {
            targetVideoTime %= videoDuration;
        }

        if (targetVideoTime < 0)
        {
            if (_targetTexture)
            {
                var active = RenderTexture.active;
                RenderTexture.active = _targetTexture;
                GL.Clear(true, true, Color.black);
                RenderTexture.active = active;
            }
            if (_quad)
            {
                _quad.SetActive(false);
            }
            return;
        }


        // scrVfxPlus hides the tutorial (editor) background when the native video starts (scrVfxPlus.cs).
        if (ADOBase.customLevel)
        {
            ADOBase.customLevel.ShowTutorialBackground(false);
        }

        try
        {
            // A seek is cheaper than decoding many frames when the time jumped forward by more than 1 s.
            if (targetVideoTime < _currentFrameTimestamp || targetVideoTime - _currentFrameTimestamp > 1.0)
            {
                _decoder.SeekToTime(targetVideoTime);
                _currentFrameTimestamp = targetVideoTime;
                _nextFrameTimestamp = targetVideoTime;
            }

            bool frameUpdated = false;
            while (targetVideoTime >= _nextFrameTimestamp)
            {
                double timestamp = 0.0;
                bool success = false;

                if (_gpuMode == GPUMode.NV12 && _textureY && _textureUV)
                {
                    success = _decoder.TryReadNV12Frame(
                        _textureY.GetRawTextureData<byte>(),
                        _textureUV.GetRawTextureData<byte>(),
                        out timestamp
                    );
                }
                else if (_gpuMode == GPUMode.YUV420P && _textureY && _textureU && _textureV)
                {
                    success = _decoder.TryReadYUV420PFrame(
                        _textureY.GetRawTextureData<byte>(),
                        _textureU.GetRawTextureData<byte>(),
                        _textureV.GetRawTextureData<byte>(),
                        out timestamp
                    );
                }
                else
                {
                    if (_textureRGBA != null)
                    {
                        success = _decoder.TryReadFrame(
                            _textureRGBA.GetRawTextureData<byte>(),
                            vflip: true,
                            out timestamp
                        );
                    }
                }

                if (success)
                {
                    _currentFrameTimestamp = timestamp;
                    _nextFrameTimestamp = timestamp + (1.0 / _decoder.FPS);
                    frameUpdated = true;
                }
                else
                {
                    if (_isLooping)
                    {
                        _decoder.SeekToTime(0);
                        _currentFrameTimestamp = 0;
                        _nextFrameTimestamp = 1.0 / _decoder.FPS;
                        break;
                    }
                    else
                    {
                        break;
                    }
                }
            }

            if (frameUpdated)
            {
                RenderToTarget();
            }
        }
        catch (Exception ex)
        {
            Melon<RenderMod>.Logger.Error($"[FFmpegVideoBGPlayer] Error during playback update: {ex}");
            _isInitialized = false;
        }
    }

    /// <summary>
    /// Uploads the decoded planes, converts them to RGBA with the compute shader (or uploads the RGBA texture), and shows
    /// the result: blitted into the target texture, or set as the quad's main texture (the quad is activated).
    /// </summary>
    private void RenderToTarget()
    {
        if (_decoder == null) return;

        int w = _decoder.Width;
        int h = _decoder.Height;

        if (_gpuMode == GPUMode.NV12 && _textureY && _textureUV && _computeShader && _renderTexture)
        {
            _textureY.Apply();
            _textureUV.Apply();

            _computeShader.SetTexture(_kernelIndex, YTex, _textureY);
            _computeShader.SetTexture(_kernelIndex, UVTex, _textureUV);
            _computeShader.SetTexture(_kernelIndex, Result, _renderTexture);
            _computeShader.SetInt(Width, w);
            _computeShader.SetInt(Height, h);
            _computeShader.SetInt(ColorRange, _decoder.ColorRange == AVColorRange.AVCOL_RANGE_JPEG ? 1 : 0);
            _computeShader.SetInt(ColorMatrix, _decoder.ColorMatrix);

            // One thread group per 8 x 8 output pixels, rounded up (group size expected by the ColorConversion.compute kernels).
            int threadGroupsX = (w + 7) / 8;
            int threadGroupsY = (h + 7) / 8;
            _computeShader.Dispatch(_kernelIndex, threadGroupsX, threadGroupsY, 1);

            if (_targetTexture)
            {
                Graphics.Blit(_renderTexture, _targetTexture);
            }
            else if (_material)
            {
                if (_material.mainTexture != _renderTexture)
                {
                    _material.mainTexture = _renderTexture;
                }
                if (_quad != null && !_quad.activeSelf)
                {
                    _quad.SetActive(true);
                }
            }
        }
        else if (_gpuMode == GPUMode.YUV420P && _textureY && _textureU && _textureV && _computeShader && _renderTexture)
        {
            _textureY.Apply();
            _textureU.Apply();
            _textureV.Apply();

            _computeShader.SetTexture(_kernelIndex, YTex, _textureY);
            _computeShader.SetTexture(_kernelIndex, UTex, _textureU);
            _computeShader.SetTexture(_kernelIndex, VTex, _textureV);
            _computeShader.SetTexture(_kernelIndex, Result, _renderTexture);
            _computeShader.SetInt(Width, w);
            _computeShader.SetInt(Height, h);
            _computeShader.SetInt(ColorRange, _decoder.ColorRange == AVColorRange.AVCOL_RANGE_JPEG ? 1 : 0);
            _computeShader.SetInt(ColorMatrix, _decoder.ColorMatrix);

            // One thread group per 8 x 8 output pixels, rounded up (group size expected by the ColorConversion.compute kernels).
            int threadGroupsX = (w + 7) / 8;
            int threadGroupsY = (h + 7) / 8;
            _computeShader.Dispatch(_kernelIndex, threadGroupsX, threadGroupsY, 1);

            if (_targetTexture)
            {
                Graphics.Blit(_renderTexture, _targetTexture);
            }
            else if (_material)
            {
                if (_material.mainTexture != _renderTexture)
                {
                    _material.mainTexture = _renderTexture;
                }
                if (_quad != null && !_quad.activeSelf)
                {
                    _quad.SetActive(true);
                }
            }
        }
        else if (_textureRGBA)
        {
            _textureRGBA.Apply();
            if (_targetTexture)
            {
                Graphics.Blit(_textureRGBA, _targetTexture);
            }
            else if (_material)
            {
                if (_material.mainTexture != _textureRGBA)
                {
                    _material.mainTexture = _textureRGBA;
                }
                if (_quad != null && !_quad.activeSelf)
                {
                    _quad.SetActive(true);
                }
            }
        }
    }

    /// <summary>Releases the decoder and GPU resources when the component is destroyed.</summary>
    private void OnDestroy()
    {
        CleanUp();
    }

    /// <summary>
    /// Disposes the decoder and destroys the textures, the RenderTexture, the background quad and its material.
    /// The player is uninitialized afterwards. Called by Init, on errors and on destroy.
    /// </summary>
    public void CleanUp()
    {
        _isInitialized = false;

        if (_decoder != null)
        {
            _decoder.Dispose();
            _decoder = null;
            Melon<RenderMod>.Logger.Msg("[FFmpegVideoBGPlayer] Decoder successfully disposed.");
        }

        if (_textureRGBA != null)
        {
            Destroy(_textureRGBA);
            _textureRGBA = null;
        }

        if (_textureY != null) Destroy(_textureY);
        if (_textureU != null) Destroy(_textureU);
        if (_textureV != null) Destroy(_textureV);
        if (_textureUV != null) Destroy(_textureUV);

        _textureY = null;
        _textureU = null;
        _textureV = null;
        _textureUV = null;

        if (_renderTexture != null)
        {
            _renderTexture.Release();
            Destroy(_renderTexture);
            _renderTexture = null;
        }

        if (_quad != null)
        {
            Destroy(_quad);
            _quad = null;
        }

        if (_material != null)
        {
            Destroy(_material);
            _material = null;
        }
    }

    /// <summary>
    /// Creates a quad as a child of <paramref name="targetCamera"/>, 0.5 units in front of its far clip plane, scaled to
    /// fill the view (orthographic size or field of view), on the layer of <paramref name="videoBGGo"/>. Its shader is
    /// taken from scrCamera's camera quad material when available, otherwise the first of Sprites/Default, Unlit/Texture
    /// and Standard that exists. The quad starts inactive.
    /// </summary>
    // ReSharper disable once InconsistentNaming
    private void SetupBackgroundQuad(Camera targetCamera, GameObject videoBGGo)
    {
        _quad = GameObject.CreatePrimitive(PrimitiveType.Quad);
        _quad.name = "FFmpegVideoBGQuad";
            
        // ReSharper disable once Unity.UnresolvedComponentOrScriptableObject
        var col = _quad.GetComponent("MeshCollider");
        if (col) Destroy(col);
            
        _quad.transform.parent = targetCamera.transform;
            
        var zPos = targetCamera.farClipPlane - 0.5f;
        _quad.transform.localPosition = new Vector3(0f, 0f, zPos);
        _quad.transform.localRotation = Quaternion.identity;
            
        float height;
        float width;
        if (targetCamera.orthographic)
        {
            height = 2.0f * targetCamera.orthographicSize;
            width = height * targetCamera.aspect;
        }
        else
        {
            height = 2.0f * zPos * Mathf.Tan(targetCamera.fieldOfView * 0.5f * Mathf.Deg2Rad);
            width = height * targetCamera.aspect;
        }
        _quad.transform.localScale = new Vector3(width, height, 1.0f);
        _quad.layer = videoBGGo.layer;

        // Set vertex colors to white so Sprites/Default works properly (otherwise it may draw transparently)
        var meshFilter = _quad.GetComponent<MeshFilter>();
        if (meshFilter != null && meshFilter.mesh != null)
        {
            var mesh = meshFilter.mesh;
            var colors = new Color[mesh.vertexCount];
            for (var i = 0; i < colors.Length; i++)
            {
                colors[i] = Color.white;
            }
            mesh.colors = colors;
        }

        Shader? shader = null;
        if (scrCamera.instance != null)
        {
            try
            {
                var camQuadMeshField = typeof(scrCamera).GetField("camQuadMesh", System.Reflection.BindingFlags.Public | System.Reflection.BindingFlags.NonPublic | System.Reflection.BindingFlags.Instance);
                if (camQuadMeshField != null)
                {
                    var renderer = camQuadMeshField.GetValue(scrCamera.instance) as MeshRenderer;
                    if (renderer != null && renderer.sharedMaterial != null)
                    {
                        shader = renderer.sharedMaterial.shader;
                    }
                }
            }
            catch (Exception)
            {
                // ignored
            }
        }
            
        if (shader == null) shader = Shader.Find("Sprites/Default");
        if (shader == null) shader = Shader.Find("Unlit/Texture");
        if (shader == null) shader = Shader.Find("Standard");
            
        _material = new Material(shader)
        {
            color = Color.white
        };

        _quad.GetComponent<MeshRenderer>().material = _material;
            
        // Set inactive initially to prevent white flash
        _quad.SetActive(false);
    }
}