Files
2026-07-13 12:49:17 +08:00

428 lines
17 KiB
C#

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