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;

// ReSharper disable once InconsistentNaming
internal class VideoBGPlayer : MonoBehaviour
{
    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");

    [SuppressMessage("ReSharper", "InconsistentNaming")]
    private enum GPUMode
    {
        None,
        YUV420P,
        NV12
    }

    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;

    private double _currentFrameTimestamp;
    private double _nextFrameTimestamp;
    private GPUMode _gpuMode = GPUMode.None;
    private RenderTexture? _targetTexture;
    private bool _isLooping;
    private bool _isInitialized;

    private GameObject? _quad;
    private Material? _material;

    // 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 / 2, h / 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 / 2, h / 2, TextureFormat.R8, false);
                            _textureV = new Texture2D(w / 2, h / 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();
        }
    }

    private static void ConfigureTexture(Texture2D tex)
    {
        tex.wrapMode = TextureWrapMode.Clamp;
        tex.filterMode = FilterMode.Bilinear;
    }

    private void LateUpdate()
    {
        if (!_isInitialized || _decoder == null) return;

        var cond = scrConductor.instance;
        if (!cond || !cond.hasSongStarted) return;
            
        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;

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


        if (ADOBase.customLevel)
        {
            ADOBase.customLevel.ShowTutorialBackground(false);
        }

        try
        {
            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;
        }
    }

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

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

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

    private void OnDestroy()
    {
        CleanUp();
    }

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

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