COM3D2.MeidoPhotoStudio/COM3D2.MeidoPhotoStudio.Plugin/MeidoPhotoStudio/MeidoPhotoStudio.cs

using System;
using System.IO;
using System.Collections;
using System.Collections.Generic;
using UnityEngine;
using UnityEngine.SceneManagement;
using Ionic.Zlib;
using BepInEx;

namespace COM3D2.MeidoPhotoStudio.Plugin
{
    using Input = InputManager;
    [BepInPlugin(pluginGuid, pluginName, pluginVersion)]
    [BepInDependency("org.bepinex.plugins.unityinjectorloader", BepInDependency.DependencyFlags.SoftDependency)]
    public class MeidoPhotoStudio : BaseUnityPlugin
    {
        private static readonly CameraMain mainCamera = GameMain.Instance.MainCamera;
        private static event EventHandler<ScreenshotEventArgs> ScreenshotEvent;
        private const string pluginGuid = "com.habeebweeb.com3d2.meidophotostudio";
        public const string pluginName = "MeidoPhotoStudio";
        public const string pluginVersion = "1.0.0";
        public const int sceneVersion = 1000;
        public const int kankyoMagic = -765;
        public static string pluginString = $"{pluginName} {pluginVersion}";
        public static bool EditMode => currentScene == Constants.Scene.Edit;
        public static event EventHandler<ScreenshotEventArgs> NotifyRawScreenshot;
        private WindowManager windowManager;
        private SceneManager sceneManager;
        private MeidoManager meidoManager;
        private EnvironmentManager environmentManager;
        private MessageWindowManager messageWindowManager;
        private LightManager lightManager;
        private PropManager propManager;
        private EffectManager effectManager;
        private CameraManager cameraManager;
        private static Constants.Scene currentScene;
        private bool initialized;
        private bool active;
        private bool uiActive;

        static MeidoPhotoStudio()
        {
            Input.Register(MpsKey.Screenshot, KeyCode.S, "Take screenshot");
            Input.Register(MpsKey.Activate, KeyCode.F6, "Activate/deactivate MeidoPhotoStudio");
        }

        private void Awake()
        {
            var harmony = HarmonyLib.Harmony.CreateAndPatchAll(typeof(AllProcPropSeqStartPatcher));
            harmony.PatchAll(typeof(BgMgrPatcher));
            ScreenshotEvent += OnScreenshotEvent;
            DontDestroyOnLoad(this);
            UnityEngine.SceneManagement.SceneManager.sceneLoaded += OnSceneLoaded;
            UnityEngine.SceneManagement.SceneManager.activeSceneChanged += OnSceneChanged;
        }

        private void Start()
        {
            Constants.Initialize();
            Translation.Initialize(Translation.CurrentLanguage);
        }

        private void OnSceneLoaded(Scene scene, LoadSceneMode sceneMode)
        {
            currentScene = (Constants.Scene)scene.buildIndex;
        }

        private void OnSceneChanged(Scene current, Scene next)
        {
            if (active) Deactivate(true);
            ResetCalcNearClip();
        }

        public byte[] SerializeScene(bool kankyo = false)
        {
            if (meidoManager.Busy) return null;

            byte[] compressedData;

            using (MemoryStream memoryStream = new MemoryStream())
            using (DeflateStream deflateStream = new DeflateStream(memoryStream, CompressionMode.Compress))
            using (BinaryWriter binaryWriter = new BinaryWriter(deflateStream, System.Text.Encoding.UTF8))
            {
                binaryWriter.Write("MPS_SCENE");
                binaryWriter.Write(sceneVersion);

                binaryWriter.Write(kankyo ? kankyoMagic : meidoManager.ActiveMeidoList.Count);

                effectManager.Serialize(binaryWriter);
                environmentManager.Serialize(binaryWriter);
                lightManager.Serialize(binaryWriter);

                if (!kankyo)
                {
                    messageWindowManager.Serialize(binaryWriter);
                    // meidomanager has to be serialized before propmanager because reattached props will be in the 
                    // wrong place after a maid's pose is deserialized.
                    meidoManager.Serialize(binaryWriter);
                }

                propManager.Serialize(binaryWriter);

                binaryWriter.Write("END");

                deflateStream.Close();

                compressedData = memoryStream.ToArray();
            }

            return compressedData;
        }

        public static byte[] DecompressScene(string filePath)
        {
            if (!File.Exists(filePath))
            {
                Utility.LogWarning($"Scene file '{filePath}' does not exist.");
                return null;
            }

            byte[] compressedData;

            using (FileStream fileStream = File.OpenRead(filePath))
            {
                if (Utility.IsPngFile(fileStream))
                {
                    if (!Utility.SeekPngEnd(fileStream))
                    {
                        Utility.LogWarning($"'{filePath}' is not a PNG file");
                        return null;
                    }

                    if (fileStream.Position == fileStream.Length)
                    {
                        Utility.LogWarning($"'{filePath}' contains no scene data");
                        return null;
                    }

                    int dataLength = (int)(fileStream.Length - fileStream.Position);

                    compressedData = new byte[dataLength];

                    fileStream.Read(compressedData, 0, dataLength);
                }
                else
                {
                    compressedData = new byte[fileStream.Length];
                    fileStream.Read(compressedData, 0, compressedData.Length);
                }
            }

            return DeflateStream.UncompressBuffer(compressedData);
        }

        public void ApplyScene(string filePath)
        {
            if (meidoManager.Busy) return;

            byte[] sceneBinary = DecompressScene(filePath);

            if (sceneBinary == null) return;

            string header = string.Empty;
            string previousHeader = string.Empty;

            using (MemoryStream memoryStream = new MemoryStream(sceneBinary))
            using (BinaryReader binaryReader = new BinaryReader(memoryStream, System.Text.Encoding.UTF8))
            {
                try
                {
                    if (binaryReader.ReadString() != "MPS_SCENE")
                    {
                        Utility.LogWarning($"'{filePath}' is not a {pluginName} scene");
                        return;
                    }

                    var version = binaryReader.ReadInt32();

                    if (version > sceneVersion)
                    {
                        Utility.LogWarning($"'{filePath}' is newer than the current version"
                        + $"Scene's version: {version}, current version: {sceneVersion}");
                        return;
                    }

                    binaryReader.ReadInt32(); // Number of Maids

                    while ((header = binaryReader.ReadString()) != "END")
                    {
                        switch (header)
                        {
                            case MessageWindowManager.header:
                                messageWindowManager.Deserialize(binaryReader);
                                break;
                            case EnvironmentManager.header:
                                environmentManager.Deserialize(binaryReader);
                                break;
                            case CameraManager.header:
                                environmentManager.Deserialize(binaryReader);
                                break;
                            case MeidoManager.header:
                                meidoManager.Deserialize(binaryReader);
                                break;
                            case PropManager.header:
                                propManager.Deserialize(binaryReader);
                                break;
                            case LightManager.header:
                                lightManager.Deserialize(binaryReader);
                                break;
                            case EffectManager.header:
                                effectManager.Deserialize(binaryReader);
                                break;
                            default: throw new Exception($"Unknown header '{header}'");
                        }
                        previousHeader = header;
                    }
                }
                catch (Exception e)
                {
                    Utility.LogError(
                        $"Failed to deserialize scene '{filePath}' because {e.Message}"
                        + $"\nCurrent header: '{header}'. Last header: '{previousHeader}'"
                    );
                    Utility.LogError(e.StackTrace);
                    return;
                }
            }
        }

        public static void TakeScreenshot(ScreenshotEventArgs args) => ScreenshotEvent?.Invoke(null, args);

        public static void TakeScreenshot(string path = "", int superSize = -1, bool hideMaids = false)
        {
            TakeScreenshot(new ScreenshotEventArgs() { Path = path, SuperSize = superSize, HideMaids = hideMaids });
        }

        private void OnScreenshotEvent(object sender, ScreenshotEventArgs args)
        {
            StartCoroutine(Screenshot(args));
        }

        private void Update()
        {
            if (currentScene == Constants.Scene.Daily || currentScene == Constants.Scene.Edit)
            {
                if (Input.GetKeyDown(MpsKey.Activate))
                {
                    if (active) Deactivate();
                    else Activate();
                }

                if (active)
                {
                    if (!Input.Control && !Input.GetKey(MpsKey.CameraLayer) && Input.GetKeyDown(MpsKey.Screenshot))
                    {
                        TakeScreenshot();
                    }

                    meidoManager.Update();
                    cameraManager.Update();
                    windowManager.Update();
                    effectManager.Update();
                    sceneManager.Update();
                }
            }
        }

        private IEnumerator Screenshot(ScreenshotEventArgs args)
        {
            // Hide UI and dragpoints
            GameObject gameMain = GameMain.Instance.gameObject;
            GameObject editUI = UTY.GetChildObject(GameObject.Find("UI Root"), "Camera");
            GameObject fpsViewer = UTY.GetChildObject(gameMain, "SystemUI Root/FpsCounter");
            GameObject sysDialog = UTY.GetChildObject(gameMain, "SystemUI Root/SystemDialog");
            GameObject sysShortcut = UTY.GetChildObject(gameMain, "SystemUI Root/SystemShortcut");

            // CameraUtility can hide the edit UI so keep its state for later
            bool editUIWasActive = editUI.activeSelf;

            uiActive = false;
            editUI.SetActive(false);
            fpsViewer.SetActive(false);
            sysDialog.SetActive(false);
            sysShortcut.SetActive(false);

            // Hide maid dragpoints and maids
            List<Meido> activeMeidoList = meidoManager.ActiveMeidoList;
            bool[] isIK = new bool[activeMeidoList.Count];
            bool[] isVisible = new bool[activeMeidoList.Count];
            for (int i = 0; i < activeMeidoList.Count; i++)
            {
                Meido meido = activeMeidoList[i];
                isIK[i] = meido.IK;
                if (meido.IK) meido.IK = false;

                // Hide the maid if needed
                if (args.HideMaids)
                {
                    isVisible[i] = meido.Maid.Visible;
                    meido.Maid.Visible = false;
                }
            }

            // Hide other drag points
            bool[] isCubeActive = {
                MeidoDragPointManager.CubeActive,
                PropManager.CubeActive,
                LightManager.CubeActive,
                EnvironmentManager.CubeActive
            };

            MeidoDragPointManager.CubeActive = false;
            PropManager.CubeActive = false;
            LightManager.CubeActive = false;
            EnvironmentManager.CubeActive = false;

            // hide gizmos
            GizmoRender.UIVisible = false;

            yield return new WaitForEndOfFrame();

            Texture2D rawScreenshot = null;

            if (args.InMemory)
            {
                // Take a screenshot directly to a Texture2D for immediate processing
                RenderTexture renderTexture = new RenderTexture(Screen.width, Screen.height, 24);
                RenderTexture.active = renderTexture;
                mainCamera.camera.targetTexture = renderTexture;
                mainCamera.camera.Render();

                rawScreenshot = new Texture2D(Screen.width, Screen.height, TextureFormat.RGB24, false);
                rawScreenshot.ReadPixels(new Rect(0, 0, renderTexture.width, renderTexture.height), 0, 0, false);
                rawScreenshot.Apply();

                mainCamera.camera.targetTexture = null;
                RenderTexture.active = null;
                DestroyImmediate(renderTexture);
            }
            else
            {
                // Take Screenshot
                int[] defaultSuperSize = new[] { 1, 2, 4 };
                int selectedSuperSize = args.SuperSize < 1
                    ? defaultSuperSize[(int)GameMain.Instance.CMSystem.ScreenShotSuperSize]
                    : args.SuperSize;

                string path = string.IsNullOrEmpty(args.Path)
                    ? Utility.ScreenshotFilename()
                    : args.Path;

                Application.CaptureScreenshot(path, selectedSuperSize);
            }
            GameMain.Instance.SoundMgr.PlaySe("se022.ogg", false);

            yield return new WaitForEndOfFrame();

            // Show UI and dragpoints
            uiActive = true;
            editUI.SetActive(editUIWasActive);
            fpsViewer.SetActive(GameMain.Instance.CMSystem.ViewFps);
            sysDialog.SetActive(true);
            sysShortcut.SetActive(true);

            for (int i = 0; i < activeMeidoList.Count; i++)
            {
                Meido meido = activeMeidoList[i];
                if (isIK[i]) meido.IK = true;
                if (args.HideMaids && isVisible[i]) meido.Maid.Visible = true;
            }

            MeidoDragPointManager.CubeActive = isCubeActive[0];
            PropManager.CubeActive = isCubeActive[1];
            LightManager.CubeActive = isCubeActive[2];
            EnvironmentManager.CubeActive = isCubeActive[3];

            GizmoRender.UIVisible = true;

            if (args.InMemory && rawScreenshot)
                NotifyRawScreenshot?.Invoke(null, new ScreenshotEventArgs() { Screenshot = rawScreenshot });
        }

        private void OnGUI()
        {
            if (uiActive)
            {
                windowManager.DrawWindows();

                if (DropdownHelper.Visible) DropdownHelper.HandleDropdown();
                if (Modal.Visible) Modal.Draw();
            }
        }

        private void Initialize()
        {
            if (initialized) return;
            initialized = true;

            meidoManager = new MeidoManager();
            environmentManager = new EnvironmentManager();
            messageWindowManager = new MessageWindowManager();
            lightManager = new LightManager();
            propManager = new PropManager(meidoManager);
            sceneManager = new SceneManager(this);
            cameraManager = new CameraManager();

            effectManager = new EffectManager();
            effectManager.AddManager<BloomEffectManager>();
            effectManager.AddManager<DepthOfFieldEffectManager>();
            effectManager.AddManager<FogEffectManager>();
            effectManager.AddManager<VignetteEffectManager>();
            effectManager.AddManager<SepiaToneEffectManger>();
            effectManager.AddManager<BlurEffectManager>();

            meidoManager.BeginCallMeidos += (s, a) => uiActive = false;
            meidoManager.EndCallMeidos += (s, a) => uiActive = true;

            MaidSwitcherPane maidSwitcherPane = new MaidSwitcherPane(meidoManager);

            SceneWindow sceneWindow = new SceneWindow(sceneManager);

            windowManager = new WindowManager()
            {
                [Constants.Window.Main] = new MainWindow(meidoManager, propManager, lightManager)
                {
                    [Constants.Window.Call] = new CallWindowPane(meidoManager),
                    [Constants.Window.Pose] = new PoseWindowPane(meidoManager, maidSwitcherPane),
                    [Constants.Window.Face] = new FaceWindowPane(meidoManager, maidSwitcherPane),
                    [Constants.Window.BG] = new BGWindowPane(
                        environmentManager, lightManager, effectManager, sceneWindow, cameraManager
                    ),
                    [Constants.Window.BG2] = new BG2WindowPane(meidoManager, propManager),
                    [Constants.Window.Settings] = new SettingsWindowPane()
                },
                [Constants.Window.Message] = new MessageWindow(messageWindowManager),
                [Constants.Window.Save] = sceneWindow
            };
        }

        private void Activate()
        {
            if (!GameMain.Instance.SysDlg.IsDecided) return;

            if (!initialized) Initialize();
            else
            {
                meidoManager.Activate();
                environmentManager.Activate();
                cameraManager.Activate();
                propManager.Activate();
                lightManager.Activate();
                effectManager.Activate();
                messageWindowManager.Activate();
                windowManager.Activate();
            }

            SetNearClipPlane();

            uiActive = true;
            active = true;

            if (!EditMode)
            {
                GameObject dailyPanel = GameObject.Find("UI Root")?.transform.Find("DailyPanel")?.gameObject;
                if (dailyPanel) dailyPanel.SetActive(false);
            }
            else meidoManager.CallMeidos();
        }

        private void Deactivate(bool force = false)
        {
            if (meidoManager.Busy || SceneManager.Busy) return;

            SystemDialog sysDialog = GameMain.Instance.SysDlg;

            void exit()
            {
                sysDialog.Close();
                ResetCalcNearClip();

                meidoManager.Deactivate();
                environmentManager.Deactivate();
                cameraManager.Deactivate();
                propManager.Deactivate();
                lightManager.Deactivate();
                effectManager.Deactivate();
                messageWindowManager.Deactivate();
                windowManager.Deactivate();
                Input.Deactivate();

                Modal.Close();

                if (!EditMode)
                {
                    GameObject dailyPanel = GameObject.Find("UI Root")?.transform.Find("DailyPanel")?.gameObject;
                    dailyPanel?.SetActive(true);
                }

                Configuration.Config.Save();
            }

            if (sysDialog.IsDecided || EditMode || force)
            {
                uiActive = false;
                active = false;

                if (EditMode || force) exit();
                else
                {
                    string exitMessage = string.Format(Translation.Get("systemMessage", "exitConfirm"), pluginName);
                    sysDialog.Show(exitMessage, SystemDialog.TYPE.OK_CANCEL,
                        f_dgOk: exit,
                        f_dgCancel: () =>
                        {
                            sysDialog.Close();
                            uiActive = true;
                            active = true;
                        }
                    );
                }
            }
        }

        private void SetNearClipPlane()
        {
            mainCamera.StopAllCoroutines();
            mainCamera.m_bCalcNearClip = false;
            mainCamera.camera.nearClipPlane = 0.01f;
        }

        private void ResetCalcNearClip()
        {
            if (mainCamera.m_bCalcNearClip) return;
            mainCamera.StopAllCoroutines();
            mainCamera.m_bCalcNearClip = true;
            mainCamera.Start();
        }
    }

    public class ScreenshotEventArgs : EventArgs
    {
        public string Path { get; set; } = string.Empty;
        public int SuperSize { get; set; } = -1;
        public bool HideMaids { get; set; }
        public bool InMemory { get; set; } = false;
        public Texture2D Screenshot { get; set; }
    }
}