using UnityEngine;

public partial class MeshPainter
{
    // --- GPU backend ---

    /// <summary>GPU backend per map: a display RenderTexture + scratch for the dilation ping-pong.</summary>
    private class GpuMap
    {
        public RenderTexture display;   // painted into (work RT for the mask map)
        public RenderTexture temp;      // dilation ping-pong scratch
        public RenderTexture final;     // MSO only: packed mask map (R=metallic, G=occlusion, A=smoothness); null otherwise
        public bool drew;               // anything drawn this stroke?
        public bool sRGB;               // sRGB color map (albedo/emission) vs linear data map (MSO)
    }

    public Shader brushShader;               // optional; else Shader.Find("Hidden/MeshPaintBrush")
    public Shader dilateShader;              // optional; else Shader.Find("Hidden/MeshPaintDilate")
    public Shader packMsShader;              // optional; else Shader.Find("Hidden/MeshPaintPackMS")
    private GpuMap[] _gpu;
    private Material _brushMat;
    private Material _dilateMat;
    private Material _packMat;
    private bool _prevSRGB;

    /// <summary>
    /// Sets up the GPU backend (RenderTexture set + materials per map). Returns true on
    /// success; on missing shaders it falls back to the CPU backend and returns false so
    /// the caller leaves _gpuActive=false (and useGpu in sync) instead of crashing later.
    /// </summary>
    private bool SetupGpu(Color32[][] canonical)
    {
        if (!brushShader) brushShader = Shader.Find("Hidden/MeshPaintBrush");
        if (!dilateShader) dilateShader = Shader.Find("Hidden/MeshPaintDilate");
        if (!packMsShader) packMsShader = Shader.Find("Hidden/MeshPaintPackMS");

        if (!brushShader || !dilateShader || !packMsShader)
        {
            Debug.LogError("[MeshPainter] GPU shaders not found (Hidden/MeshPaintBrush, " +
                           "Hidden/MeshPaintDilate or Hidden/MeshPaintPackMS). Falling back to CPU painting.", this);
            useGpu = false;
            SetupCpu(canonical);
            return false;
        }

        _brushMat = new Material(brushShader) { hideFlags = HideFlags.HideAndDontSave };
        _dilateMat = new Material(dilateShader) { hideFlags = HideFlags.HideAndDontSave };
        _packMat = new Material(packMsShader) { hideFlags = HideFlags.HideAndDontSave };

        _gpu = new GpuMap[MAP_COUNT];
        for (int mi = 0; mi < MAP_COUNT; mi++)
        {
            bool linear = MapIsLinear(mi);
            var gm = new GpuMap
            {
                sRGB = !linear,
                display = CreateRT("MeshPainter_DisplayRT_" + mi, linear),
                temp = CreateRT("MeshPainter_TempRT_" + mi, linear),
                final = (mi == MAP_MSO) ? CreateRT("MeshPainter_FinalRT_" + mi, linear) : null
            };
            _gpu[mi] = gm;

            SeedGpuMap(mi, canonical);

            // Color maps render to the display RT; the MS map exposes its packed final RT.
            AssignMapTexture(mi, gm.final != null ? gm.final : gm.display);
        }
        return true;
    }

    /// <summary>Seeds a GPU map's display RT from canonical data (alpha = 0 = unpainted), then packs MS to final.</summary>
    private void SeedGpuMap(int mi, Color32[][] canonical)
    {
        var src = (canonical != null && canonical[mi] != null && canonical[mi].Length == textureSize * textureSize)
            ? canonical[mi] : _seedCanonical[mi];
        var gm = _gpu[mi];

        if (mi == MAP_MSO)
        {
            // Canonical (R=metallic, G=occlusion, A=smoothness) -> work layout
            // (R=metallic, G=smoothness, B=occlusion); alpha = coverage stays 0 (unpainted).
            var work = new Color32[src.Length];
            for (int i = 0; i < src.Length; i++) work[i] = new Color32(src[i].r, src[i].a, src[i].g, 0);
            UploadColorsToRT(work, gm.display, paintedAlpha: false, linear: true);
            PackMs(gm); // build the final mask RT once so the material shows the seeded values
        }
        else
        {
            UploadColorsToRT((Color32[])src.Clone(), gm.display, paintedAlpha: false, linear: gm.sRGB == false);
        }
    }

    private void TeardownGpu()
    {
        if (_gpu != null)
        {
            foreach (var gm in _gpu)
            {
                if (gm == null) continue;
                if (gm.display) { gm.display.Release(); Destroy(gm.display); }
                if (gm.temp) { gm.temp.Release(); Destroy(gm.temp); }
                if (gm.final) { gm.final.Release(); Destroy(gm.final); }
            }
            _gpu = null;
        }
        if (_brushMat) { Destroy(_brushMat); _brushMat = null; }
        if (_dilateMat) { Destroy(_dilateMat); _dilateMat = null; }
        if (_packMat) { Destroy(_packMat); _packMat = null; }
    }

    private RenderTexture CreateRT(string rtName, bool linear)
    {
        var rt = new RenderTexture(textureSize, textureSize, 0, RenderTextureFormat.ARGB32,
            linear ? RenderTextureReadWrite.Linear : RenderTextureReadWrite.sRGB)
        {
            name = rtName,
            wrapMode = TextureWrapMode.Clamp,
            filterMode = FilterMode.Bilinear,
            antiAliasing = 1,
            useMipMap = false,
            autoGenerateMips = false
        };
        rt.Create();
        return rt;
    }

    /// <summary>Writes Color32 data into a RenderTexture via a temporary texture.</summary>
    private void UploadColorsToRT(Color32[] colors, RenderTexture rt, bool paintedAlpha, bool linear)
    {
        var tmp = new Texture2D(textureSize, textureSize, TextureFormat.RGBA32, false, linear);
        // Alpha marks painted (255) vs unpainted (0) for the dilation pass.
        byte a = paintedAlpha ? (byte)255 : (byte)0;
        for (int i = 0; i < colors.Length; i++) colors[i].a = a;
        tmp.SetPixels32(colors);
        tmp.Apply(false);

        bool prev = GL.sRGBWrite;
        GL.sRGBWrite = !linear;
        Graphics.Blit(tmp, rt);
        GL.sRGBWrite = prev;
        Destroy(tmp);
    }

    /// <summary>Reads a RenderTexture back as Color32 (for the backend switch).</summary>
    private Color32[] ReadRtPixels(RenderTexture rt, bool linear)
    {
        var prev = RenderTexture.active;
        RenderTexture.active = rt;
        var tex = new Texture2D(textureSize, textureSize, TextureFormat.RGBA32, false, linear);
        tex.ReadPixels(new Rect(0, 0, textureSize, textureSize), 0, 0, false);
        tex.Apply(false);
        RenderTexture.active = prev;
        var px = tex.GetPixels32();
        Destroy(tex);
        return px;
    }

    private void PackMs(GpuMap gm)
    {
        bool prev = GL.sRGBWrite;
        GL.sRGBWrite = false; // linear data, no conversion
        Graphics.Blit(gm.display, gm.final, _packMat, 0);
        GL.sRGBWrite = prev;
    }

    // --- GPU painting ---

    private void BeginGpuStroke()
    {
        _prevSRGB = GL.sRGBWrite;
        for (int mi = 0; mi < MAP_COUNT; mi++)
            if (IsChannelActive(mi)) _gpu[mi].drew = false;
    }

    private void DrawBrushGpu(Vector3 pLocal, float radiusLocal, Vector3 u, Vector3 v, Vector3 n)
    {
        for (int mi = 0; mi < MAP_COUNT; mi++)
        {
            if (!IsChannelActive(mi)) continue;
            var gm = _gpu[mi];

            Graphics.SetRenderTarget(gm.display);
            GL.sRGBWrite = gm.sRGB; // sRGB color map vs linear data map (MS)

            SetBrushColorForMap(mi);
            // Per-channel write mask: lets the shared mask map (R=metallic, G=smoothness work,
            // B=occlusion) be painted by metallic/smoothness and occlusion INDEPENDENTLY -
            // untouched channels keep their value while coverage (A) still accumulates.
            _brushMat.SetFloat(ColorWriteMaskId, BrushColorMaskForMap(mi));
            _brushMat.SetTexture(MaskTexId, _resolvedMask);
            _brushMat.SetVector(BrushCenterId, pLocal);
            _brushMat.SetFloat(BrushRadiusId, radiusLocal);
            _brushMat.SetVector(BrushUId, u);
            _brushMat.SetVector(BrushVId, v);
            _brushMat.SetVector(BrushNormalId, n);
            if (!_brushMat.SetPass(0)) continue;

            int sub = _bakedMesh.subMeshCount;
            for (int i = 0; i < sub; i++)
                Graphics.DrawMeshNow(_bakedMesh, Matrix4x4.identity, i);

            gm.drew = true;
        }
    }

    private void EndGpuStroke()
    {
        Graphics.SetRenderTarget(null);

        for (int mi = 0; mi < MAP_COUNT; mi++)
        {
            if (!IsChannelActive(mi)) continue;
            var gm = _gpu[mi];

            // Edge padding against UV seams: dilate the display RT into temp, then swap.
            if (gm.drew && paddingTexels > 0)
            {
                GL.sRGBWrite = gm.sRGB;
                _dilateMat.SetInt(PadId, paddingTexels);
                Graphics.Blit(gm.display, gm.temp, _dilateMat, 0);

                var swap = gm.display; gm.display = gm.temp; gm.temp = swap;
                if (gm.final == null) AssignMapTexture(mi, gm.display);
            }

            // MS: pack the work layout (R=metallic, G=smoothness) into the URP final layout.
            if (gm.final != null) PackMs(gm);
        }

        GL.sRGBWrite = _prevSRGB;
    }

    /// <summary>Sets _BrushColor for the given map. Albedo/emission are colors; the mask map packs (metallic, smoothness, occlusion) as raw data.</summary>
    private void SetBrushColorForMap(int mi)
    {
        switch (mi)
        {
            case MAP_ALBEDO:
                _brushMat.SetColor(BrushColorId, brushColor);
                break;
            case MAP_MSO:
                // Raw data via SetVector. Work layout: R=metallic, G=smoothness, B=occlusion
                // (ColorMask picks the channels). _BrushColor is a Color property, so Unity
                // applies sRGB->linear on upload in a linear project - we pre-apply the inverse
                // so the data round-trips EXACTLY and matches the CPU's raw byte writes.
                _brushMat.SetVector(BrushColorId, new Vector4(
                    EncodeData(metallic), EncodeData(smoothness), EncodeData(occlusion), 1f));
                break;
            default: // emission
                _brushMat.SetColor(BrushColorId, emissionColor);
                break;
        }
    }

    /// <summary>
    /// ColorMask bitmask for the brush draw (ColorWriteMask: R=8, G=4, B=2, A=1). Color maps
    /// write everything; the mask map writes only the active sub-channels so metallic/smoothness
    /// and occlusion can be painted independently into one texture. Coverage (A) is always written.
    /// </summary>
    /// <summary>
    /// Pre-compensates a raw data value (0..1) for the brush's Color property, which Unity
    /// converts sRGB->linear on upload in a linear project. Applying the inverse here makes
    /// the value round-trip exactly (so metallic/smoothness/occlusion are stored faithfully
    /// and match the CPU byte writes). In gamma color space no conversion happens, so pass through.
    /// </summary>
    private static float EncodeData(float v)
    {
        v = Mathf.Clamp01(v);
        return QualitySettings.activeColorSpace == ColorSpace.Linear ? Mathf.LinearToGammaSpace(v) : v;
    }

    private float BrushColorMaskForMap(int mi)
    {
        if (mi != MAP_MSO) return 15f; // RGBA
        bool ms = paintMetallicSmoothness, occ = paintOcclusion;
        if (ms && occ) return 15f;     // R|G|B|A
        if (ms) return 13f;            // R|G|A  (skip B = occlusion)
        return 3f;                     // B|A    (skip R,G = metallic/smoothness) -> occlusion only
    }
}