EditorWindowScreenshotUtility.cs

using System;
using System.Collections.Generic;
using System.IO;
using System.Reflection;
using System.Runtime.ExceptionServices;
using System.Threading;
using MCPForUnity.Runtime.Helpers;
using UnityEditor;
using UnityEngine;

namespace MCPForUnity.Editor.Helpers
{
    /// <summary>
    /// Captures the pixels currently displayed in an editor window viewport.
    /// Uses the editor view's own pixel grab path instead of re-rendering through a Camera.
    /// </summary>
    internal static class EditorWindowScreenshotUtility
    {
        // Keep capture synchronous so callers can immediately return the screenshot payload.
        // The short sleep gives Unity a chance to flush repaint work before GrabPixels reads the viewport.
        private const int RepaintSettlingDelayMs = 75;
        private static readonly HashSet<string> WindowsReservedNames = new HashSet<string>(StringComparer.OrdinalIgnoreCase)
        {
            "CON", "PRN", "AUX", "NUL",
            "COM1", "COM2", "COM3", "COM4", "COM5", "COM6", "COM7", "COM8", "COM9",
            "LPT1", "LPT2", "LPT3", "LPT4", "LPT5", "LPT6", "LPT7", "LPT8", "LPT9",
        };

        /// <summary>
        /// Captures the active Scene View viewport to a PNG asset.
        /// </summary>
        /// <param name="sceneView">Scene View window to capture.</param>
        /// <param name="fileName">Optional file name, defaulting to a timestamped PNG.</param>
        /// <param name="superSize">
        /// Preserved in the result for API consistency, but Scene View capture always uses the current viewport resolution.
        /// </param>
        /// <param name="ensureUniqueFileName">If true, appends a suffix instead of overwriting an existing file.</param>
        /// <param name="includeImage">If true, includes a base64 PNG in the returned result.</param>
        /// <param name="maxResolution">Maximum edge length for the inline image payload.</param>
        /// <param name="viewportWidth">Captured viewport width in pixels.</param>
        /// <param name="viewportHeight">Captured viewport height in pixels.</param>
        public static ScreenshotCaptureResult CaptureSceneViewViewportToProject(
            SceneView sceneView,
            string fileName,
            int superSize,
            bool ensureUniqueFileName,
            bool includeImage,
            int maxResolution,
            out int viewportWidth,
            out int viewportHeight,
            string folderOverride = null)
        {
            if (sceneView == null)
                throw new ArgumentNullException(nameof(sceneView));

            int effectiveSuperSize = NormalizeSceneViewSuperSize(superSize);

            FocusAndRepaint(sceneView);

            Rect viewportRectPixels = GetSceneViewViewportPixelRect(sceneView);
            viewportWidth = Mathf.RoundToInt(viewportRectPixels.width);
            viewportHeight = Mathf.RoundToInt(viewportRectPixels.height);

            if (viewportWidth <= 0 || viewportHeight <= 0)
                throw new InvalidOperationException("Captured Scene view viewport is empty.");

            Texture2D captured = null;
            Texture2D downscaled = null;
            try
            {
                captured = CaptureViewRect(sceneView, viewportRectPixels);

                var result = PrepareCaptureResult(fileName, effectiveSuperSize, ensureUniqueFileName, folderOverride);
                byte[] png = captured.EncodeToPNG();
                File.WriteAllBytes(result.FullPath, png);

                if (includeImage)
                {
                    int targetMax = maxResolution > 0 ? maxResolution : 640;
                    string imageBase64;
                    int imageWidth;
                    int imageHeight;

                    if (captured.width > targetMax || captured.height > targetMax)
                    {
                        downscaled = ScreenshotUtility.DownscaleTexture(captured, targetMax);
                        imageBase64 = Convert.ToBase64String(downscaled.EncodeToPNG());
                        imageWidth = downscaled.width;
                        imageHeight = downscaled.height;
                    }
                    else
                    {
                        imageBase64 = Convert.ToBase64String(png);
                        imageWidth = captured.width;
                        imageHeight = captured.height;
                    }

                    return new ScreenshotCaptureResult(
                        result.FullPath,
                        result.ProjectRelativePath,
                        result.SuperSize,
                        false,
                        imageBase64,
                        imageWidth,
                        imageHeight);
                }

                return result;
            }
            finally
            {
                DestroyTexture(captured);
                DestroyTexture(downscaled);
            }
        }

        private static void FocusAndRepaint(SceneView sceneView)
        {
            try
            {
                sceneView.Focus();
            }
            catch (Exception ex)
            {
                McpLog.Debug($"[EditorWindowScreenshotUtility] SceneView focus failed: {ex.Message}");
            }

            try
            {
                sceneView.Repaint();
                InvokeMethodIfExists(sceneView, "RepaintImmediately");
                SceneView.RepaintAll();
                UnityEditorInternal.InternalEditorUtility.RepaintAllViews();
                EditorApplication.QueuePlayerLoopUpdate();
                Thread.Sleep(RepaintSettlingDelayMs);
            }
            catch (Exception ex)
            {
                McpLog.Debug($"[EditorWindowScreenshotUtility] SceneView repaint failed: {ex.Message}");
            }
        }

        private static Rect GetSceneViewViewportPixelRect(SceneView sceneView)
        {
            float pixelsPerPoint = EditorGUIUtility.pixelsPerPoint;
            Rect viewportLocalPoints = GetViewportLocalRectPoints(sceneView, pixelsPerPoint);
            if (viewportLocalPoints.width <= 0f || viewportLocalPoints.height <= 0f)
                throw new InvalidOperationException("Failed to resolve Scene view viewport rect.");

            return new Rect(
                Mathf.Round(viewportLocalPoints.x * pixelsPerPoint),
                Mathf.Round(viewportLocalPoints.y * pixelsPerPoint),
                Mathf.Round(viewportLocalPoints.width * pixelsPerPoint),
                Mathf.Round(viewportLocalPoints.height * pixelsPerPoint));
        }

        private static Rect GetViewportLocalRectPoints(SceneView sceneView, float pixelsPerPoint)
        {
            Rect? cameraViewport = GetRectProperty(sceneView, "cameraViewport");
            if (cameraViewport.HasValue && cameraViewport.Value.width > 0f && cameraViewport.Value.height > 0f)
            {
                return cameraViewport.Value;
            }

            Camera camera = sceneView.camera;
            if (camera == null)
                throw new InvalidOperationException("Active Scene View has no camera to derive viewport size from.");

            float viewportWidth = camera.pixelWidth / Mathf.Max(0.0001f, pixelsPerPoint);
            float viewportHeight = camera.pixelHeight / Mathf.Max(0.0001f, pixelsPerPoint);
            Rect windowRect = sceneView.position;

            return new Rect(
                0f,
                Mathf.Max(0f, windowRect.height - viewportHeight),
                Mathf.Min(windowRect.width, viewportWidth),
                Mathf.Min(windowRect.height, viewportHeight));
        }

        private static Texture2D CaptureViewRect(SceneView sceneView, Rect viewportRectPixels)
        {
            object hostView = GetHostView(sceneView);
            if (hostView == null)
                throw new InvalidOperationException("Failed to resolve Scene view host view.");

            // GrabPixels is an internal extern on GUIView (parent of HostView), present since at least Unity 2021.1.
            // See: UnityCsReference/Editor/Mono/GUIView.bindings.cs — `internal extern void GrabPixels(RenderTexture, Rect)`
            // If Unity removes this, the MissingMethodException below keeps the failure explicit.
            MethodInfo grabPixels = hostView.GetType().GetMethod(
                "GrabPixels",
                BindingFlags.Instance | BindingFlags.Public | BindingFlags.NonPublic,
                null,
                new[] { typeof(RenderTexture), typeof(Rect) },
                null);

            if (grabPixels == null)
                throw new MissingMethodException($"{hostView.GetType().FullName}.GrabPixels(RenderTexture, Rect)");

            int width = Mathf.RoundToInt(viewportRectPixels.width);
            int height = Mathf.RoundToInt(viewportRectPixels.height);

            RenderTexture rt = null;
            RenderTexture previousActive = RenderTexture.active;
            try
            {
                rt = new RenderTexture(width, height, 0, RenderTextureFormat.ARGB32)
                {
                    antiAliasing = 1,
                    filterMode = FilterMode.Bilinear,
                    hideFlags = HideFlags.HideAndDontSave,
                };
                rt.Create();

                grabPixels.Invoke(hostView, new object[] { rt, viewportRectPixels });

                RenderTexture.active = rt;
                var texture = new Texture2D(width, height, TextureFormat.RGBA32, false);
                texture.ReadPixels(new Rect(0, 0, width, height), 0, 0);
                texture.Apply();
                FlipTextureVertically(texture);
                return texture;
            }
            catch (TargetInvocationException ex)
            {
                ExceptionDispatchInfo.Capture(ex.InnerException ?? ex).Throw();
                throw;
            }
            finally
            {
                RenderTexture.active = previousActive;
                if (rt != null)
                {
                    rt.Release();
                    UnityEngine.Object.DestroyImmediate(rt);
                }
            }
        }

        private static object GetHostView(EditorWindow window)
        {
            if (window == null)
                return null;

            Type windowType = typeof(EditorWindow);
            FieldInfo parentField = windowType.GetField("m_Parent", BindingFlags.Instance | BindingFlags.NonPublic);
            if (parentField != null)
            {
                object parent = parentField.GetValue(window);
                if (parent != null)
                    return parent;
            }

            PropertyInfo hostViewProperty = windowType.GetProperty("hostView", BindingFlags.Instance | BindingFlags.Public | BindingFlags.NonPublic);
            return hostViewProperty?.GetValue(window, null);
        }

        private static Rect? GetRectProperty(object instance, string propertyName)
        {
            if (instance == null)
                return null;

            Type type = instance.GetType();
            PropertyInfo property = type.GetProperty(propertyName, BindingFlags.Instance | BindingFlags.Public | BindingFlags.NonPublic);
            if (property == null || property.PropertyType != typeof(Rect))
                return null;

            try
            {
                return (Rect)property.GetValue(instance, null);
            }
            catch (Exception ex)
            {
                McpLog.Debug($"[EditorWindowScreenshotUtility] Failed to read rect property '{propertyName}': {ex.Message}");
                return null;
            }
        }

        private static void InvokeMethodIfExists(object instance, string methodName)
        {
            if (instance == null)
                return;

            MethodInfo method = instance.GetType().GetMethod(methodName, BindingFlags.Instance | BindingFlags.Public | BindingFlags.NonPublic);
            if (method == null || method.GetParameters().Length != 0)
                return;

            try
            {
                method.Invoke(instance, null);
            }
            catch (Exception ex)
            {
                McpLog.Debug($"[EditorWindowScreenshotUtility] Best-effort invoke of '{methodName}' failed: {ex.Message}");
            }
        }

        private static void FlipTextureVertically(Texture2D texture)
        {
            if (texture == null)
                return;

            int width = texture.width;
            int height = texture.height;
            Color32[] pixels = texture.GetPixels32();
            var temp = new Color32[width];

            for (int y = 0; y < height / 2; y++)
            {
                int topRow = y * width;
                int bottomRow = (height - 1 - y) * width;
                Array.Copy(pixels, topRow, temp, 0, width);
                Array.Copy(pixels, bottomRow, pixels, topRow, width);
                Array.Copy(temp, 0, pixels, bottomRow, width);
            }

            texture.SetPixels32(pixels);
            texture.Apply();
        }

        private static ScreenshotCaptureResult PrepareCaptureResult(string fileName, int superSize, bool ensureUniqueFileName, string folderOverride)
        {
            int size = Mathf.Max(1, superSize);
            string resolvedName = BuildFileName(fileName);
            string folder = ScreenshotUtility.ResolveFolderAbsolute(folderOverride);
            Directory.CreateDirectory(folder);

            string fullPath = Path.Combine(folder, resolvedName);
            if (ensureUniqueFileName)
            {
                fullPath = EnsureUnique(fullPath);
            }

            string normalizedFullPath = fullPath.Replace('\\', '/');
            string projectRoot = Path.GetFullPath(Path.Combine(Application.dataPath, "..")).Replace('\\', '/');
            string normalizedRoot = projectRoot.EndsWith("/") ? projectRoot : projectRoot + "/";
            string projectRelativePath = normalizedFullPath.StartsWith(normalizedRoot, StringComparison.OrdinalIgnoreCase)
                ? normalizedFullPath.Substring(normalizedRoot.Length)
                : normalizedFullPath;
            return new ScreenshotCaptureResult(normalizedFullPath, projectRelativePath, size, false);
        }

        private static string BuildFileName(string fileName)
        {
            string baseName = string.IsNullOrWhiteSpace(fileName)
                ? $"screenshot-{DateTime.Now:yyyyMMdd-HHmmss}.png"
                : SanitizeFileName(fileName);

            if (!baseName.EndsWith(".png", StringComparison.OrdinalIgnoreCase))
                baseName += ".png";

            return baseName;
        }

        private static int NormalizeSceneViewSuperSize(int superSize)
        {
            if (superSize > 1)
            {
                McpLog.Warn("[EditorWindowScreenshotUtility] Scene View capture ignores superSize and uses the displayed viewport resolution.");
                return 1;
            }

            return Mathf.Max(1, superSize);
        }

        private static string SanitizeFileName(string fileName)
        {
            string trimmed = (fileName ?? string.Empty).Trim();
            if (string.IsNullOrEmpty(trimmed))
                return $"screenshot-{DateTime.Now:yyyyMMdd-HHmmss}.png";

            string candidate = trimmed;
            string normalizedSeparators = candidate.Replace('\\', '/');
            if (Path.IsPathRooted(candidate) || normalizedSeparators.Contains("/") || normalizedSeparators.Contains(".."))
            {
                string[] pathParts = normalizedSeparators.Split(new[] { '/' }, StringSplitOptions.RemoveEmptyEntries);
                candidate = pathParts.Length > 0 ? pathParts[pathParts.Length - 1] : string.Empty;
            }

            if (string.IsNullOrWhiteSpace(candidate) || candidate == "." || candidate == "..")
                candidate = $"screenshot-{DateTime.Now:yyyyMMdd-HHmmss}.png";

            char[] invalidChars = Path.GetInvalidFileNameChars();
            foreach (char invalidChar in invalidChars)
            {
                candidate = candidate.Replace(invalidChar, '_');
            }

            string extension = Path.GetExtension(candidate);
            string stem = Path.GetFileNameWithoutExtension(candidate);
            extension = extension.TrimEnd(' ', '.');
            stem = stem.TrimEnd(' ', '.');
            if (WindowsReservedNames.Contains(stem))
            {
                candidate = $"_{stem}{extension}";
            }

            return candidate;
        }

        private static string EnsureUnique(string fullPath)
        {
            if (!File.Exists(fullPath))
                return fullPath;

            string directory = Path.GetDirectoryName(fullPath) ?? string.Empty;
            string fileNameWithoutExtension = Path.GetFileNameWithoutExtension(fullPath);
            string extension = Path.GetExtension(fullPath);

            for (int i = 1; i < 10000; i++)
            {
                string candidate = Path.Combine(directory, $"{fileNameWithoutExtension}-{i}{extension}");
                if (!File.Exists(candidate))
                    return candidate;
            }

            throw new IOException($"Could not generate a unique screenshot filename for '{fullPath}'.");
        }

        private static void DestroyTexture(Texture2D texture)
        {
            if (texture == null)
                return;

            UnityEngine.Object.DestroyImmediate(texture);
        }
    }
}