using System;
using System.IO;
using UnityEngine;

[RequireComponent(typeof(Camera))]
public class EditorScreenshot : MonoBehaviour
{
    [SerializeField] string folder = "EditorScreenshot/results";
    [SerializeField] string prefix = "screenshot";
    [SerializeField] bool ensureTransparentBackground = false;

    Camera Camera
    {
        get
        {
            if (camera == null) {
                camera = GetComponent<Camera>();
            }
            return camera;
        }
    }
    new Camera camera;
    int Width
    {
        get => Camera.main.pixelWidth;
    }
    int Height
    {
        get => Camera.main.pixelHeight;
    }

    [ContextMenu("Take Screenshot")]
    public void TakeScreenshot()
    {
        PrepareCamera();
        WriteScreenshot();
    }

    void WriteScreenshot()
    {
        RenderTexture renderTextureBefore = RenderTexture.active;
        RenderTexture.active = camera.targetTexture;

        Debug.Log($"{Width}x{Height}");

        Texture2D screenshot = new Texture2D(Width, Height, TextureFormat.ARGB32, false);
        screenshot.ReadPixels(new Rect(0, 0, Width, Height), 0, 0, false);

        if (QualitySettings.activeColorSpace == ColorSpace.Linear)
        {
            Color[] pixels = screenshot.GetPixels();
            for (int p = 0; p < pixels.Length; p++)
            {
                pixels[p] = pixels[p].gamma;
            }
            screenshot.SetPixels(pixels);
        }
        screenshot.Apply(false);

        folder = GetSafePath(folder.Trim('/'));
        prefix = GetSafeFilename(prefix);

        string dir = $"{Application.dataPath}/{folder}/";
        string filename = $"{prefix}_{DateTime.Now.ToString("yyMMdd_HHmmss")}.png";
        string path = $"{dir}{filename}";

        Directory.CreateDirectory(dir);
        File.WriteAllBytes(path, screenshot.EncodeToPNG());

        camera.targetTexture = null;
        RenderTexture.active = renderTextureBefore;

        Debug.Log("Screenshot saved to:\n" + path);
    }

    void PrepareCamera()
    {
        Camera camera = Camera;
        camera.targetTexture = new RenderTexture(Width, Height, 0, RenderTextureFormat.ARGB32);

        CameraClearFlags clearFlagsBefore = camera.clearFlags;
        Color backgroundColorBefore = camera.backgroundColor;

        if (ensureTransparentBackground)
        {
            camera.clearFlags = CameraClearFlags.SolidColor;
            camera.backgroundColor = new Color();
        }

        camera.Render();

        if (ensureTransparentBackground)
        {
            camera.clearFlags = clearFlagsBefore;
            camera.backgroundColor = backgroundColorBefore;
        }
    }

    public string GetSafePath(string path)
    {
        return string.Join("_", path.Split(Path.GetInvalidPathChars()));
    }

    public string GetSafeFilename(string filename)
    {
        return string.Join("_", filename.Split(Path.GetInvalidFileNameChars()));
    }
}