using UnityEngine;

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

    private void SetupCpu(Color32[][] canonical)
    {
        _cpu = new CpuMap[MAP_COUNT];
        for (int mi = 0; mi < MAP_COUNT; mi++)
        {
            var cm = new CpuMap
            {
                tex = new Texture2D(textureSize, textureSize, TextureFormat.RGBA32, false, MapIsLinear(mi))
                {
                    name = "MeshPainter_Canvas_" + mi,
                    wrapMode = TextureWrapMode.Clamp,
                    filterMode = FilterMode.Bilinear
                },
                pixels = SeedCpuPixels(mi, canonical),
                painted = new bool[textureSize * textureSize],
                fillStamp = new int[textureSize * textureSize],
                dilateStamp = 0
            };
            cm.tex.SetPixels32(cm.pixels);
            cm.tex.Apply(false);
            _cpu[mi] = cm;
            AssignMapTexture(mi, cm.tex);
        }
    }

    /// <summary>CPU pixels are stored directly in material layout; just guarantee a sane alpha for the color maps.</summary>
    private Color32[] SeedCpuPixels(int mi, Color32[][] canonical)
    {
        var src = (canonical != null && canonical[mi] != null && canonical[mi].Length == textureSize * textureSize)
            ? canonical[mi] : _seedCanonical[mi];
        var px = (Color32[])src.Clone();
        if (mi != MAP_MSO) // albedo / emission: opaque. The mask map keeps alpha = smoothness.
            for (int i = 0; i < px.Length; i++) px[i].a = 255;
        return px;
    }

    private void TeardownCpu()
    {
        if (_cpu == null) return;
        foreach (var cm in _cpu) if (cm != null && cm.tex) Destroy(cm.tex);
        _cpu = null;
    }

    // --- CPU painting (rasterization + dilation) ---

    private void ResetCpuDirty()
    {
        for (int mi = 0; mi < MAP_COUNT; mi++)
        {
            if (!IsChannelActive(mi)) continue;
            var cm = _cpu[mi];
            cm.dx0 = textureSize; cm.dy0 = textureSize; cm.dx1 = -1; cm.dy1 = -1;
        }
    }

    private void EndCpuStroke()
    {
        for (int mi = 0; mi < MAP_COUNT; mi++)
        {
            if (!IsChannelActive(mi)) continue;
            var cm = _cpu[mi];
            if (paddingTexels > 0) DilatePadding(cm, paddingTexels);
            FlushDirty(cm);
        }
    }

    // sRGB<->linear LUT so the CPU blends the sRGB color maps (albedo/emission) in LINEAR
    // space, matching the GPU's sRGB-framebuffer blend. Without it, partial-opacity texels
    // (soft mask edges, brushColor.a<1) would differ between the GPU and CPU backends.
    private static readonly float[] _srgbToLinear = BuildSrgbToLinearLut();

    private static float[] BuildSrgbToLinearLut()
    {
        var lut = new float[256];
        for (int i = 0; i < 256; i++) lut[i] = Mathf.GammaToLinearSpace(i / 255f);
        return lut;
    }

    /// <summary>Blends a linear source over an sRGB-encoded destination byte in linear space, returns the sRGB byte.</summary>
    private static byte BlendSrgb(byte dstByte, float srcLinear, float invOp, float op)
    {
        float outLinear = _srgbToLinear[dstByte] * invOp + srcLinear * op;
        return (byte)(Mathf.Clamp01(Mathf.LinearToGammaSpace(outLinear)) * 255f + 0.5f);
    }

    /// <summary>
    /// Paints a square stamp in the surface tangent plane: rasterizes triangles in range
    /// and blends every texel whose 3D position lies inside the box (|du|,|dv|,|dn| &lt;=
    /// radius) over each active map, weighted by the mask opacity. Mirrors the GPU brush
    /// path (mask value = per-texel opacity, gray = 50%).
    /// </summary>
    private void PaintStamp3D(Vector3 pLocal, float radiusLocal, Vector3 u, Vector3 v, Vector3 n)
    {
        bool aAlb = paintAlbedo, aMS = paintMetallicSmoothness, aEM = paintEmission, aAO = paintOcclusion;
        if (!aAlb && !aMS && !aEM && !aAO) return;

        // Albedo target (brush color in linear space; blended linear to match the GPU).
        float brLin = Mathf.GammaToLinearSpace(brushColor.r);
        float bgLin = Mathf.GammaToLinearSpace(brushColor.g);
        float bbLin = Mathf.GammaToLinearSpace(brushColor.b);
        float albOpMul = Mathf.Clamp01(brushColor.a);
        CpuMap alb = _cpu[MAP_ALBEDO];

        // Mask map target: R = metallic, G = occlusion, B = 0, A = smoothness (one texture
        // bound to both _MetallicGlossMap and _OcclusionMap).
        float msoR = Mathf.Clamp01(metallic) * 255f;    // metallic
        float msoG = Mathf.Clamp01(occlusion) * 255f;   // occlusion
        float msoA = Mathf.Clamp01(smoothness) * 255f;  // smoothness
        float msOpMul = Mathf.Clamp01(metallicSmoothnessOpacity); // metallic+smoothness paint opacity
        float aoOpMul = Mathf.Clamp01(occlusionOpacity);          // occlusion paint opacity
        CpuMap mso = _cpu[MAP_MSO];
        bool aMSO = aMS || aAO;

        // Emission target (linear, like albedo).
        float erLin = Mathf.GammaToLinearSpace(emissionColor.r);
        float egLin = Mathf.GammaToLinearSpace(emissionColor.g);
        float ebLin = Mathf.GammaToLinearSpace(emissionColor.b);
        float emOpMul = Mathf.Clamp01(emissionColor.a);
        CpuMap em = _cpu[MAP_EM];

        float invR = 1f / radiusLocal;
        int triCount = _tris.Length / 3;
        float ts = textureSize;

        for (int t = 0; t < triCount; t++)
        {
            // Fast culling: stamp box (corner reach = sqrt(3)*radius) vs. the triangle's circumscribed circle.
            float reach = radiusLocal * SquareReach + _triRadius[t];
            if ((_triCenter[t] - pLocal).sqrMagnitude > reach * reach) continue;

            int i0 = _tris[t * 3], i1 = _tris[t * 3 + 1], i2 = _tris[t * 3 + 2];
            Vector3 v0 = _verts[i0], v1 = _verts[i1], v2 = _verts[i2];
            Vector2 uv0 = _uvs[i0], uv1 = _uvs[i1], uv2 = _uvs[i2];

            // UV -> pixel.
            float p0x = uv0.x * ts, p0y = uv0.y * ts;
            float p1x = uv1.x * ts, p1y = uv1.y * ts;
            float p2x = uv2.x * ts, p2y = uv2.y * ts;

            int minX = Mathf.Clamp(Mathf.FloorToInt(Mathf.Min(p0x, Mathf.Min(p1x, p2x))), 0, textureSize - 1);
            int maxX = Mathf.Clamp(Mathf.CeilToInt(Mathf.Max(p0x, Mathf.Max(p1x, p2x))), 0, textureSize - 1);
            int minY = Mathf.Clamp(Mathf.FloorToInt(Mathf.Min(p0y, Mathf.Min(p1y, p2y))), 0, textureSize - 1);
            int maxY = Mathf.Clamp(Mathf.CeilToInt(Mathf.Max(p0y, Mathf.Max(p1y, p2y))), 0, textureSize - 1);
            if (maxX < minX || maxY < minY) continue;

            // Barycentric denominator.
            float d = (p1y - p2y) * (p0x - p2x) + (p2x - p1x) * (p0y - p2y);
            if (Mathf.Abs(d) < 1e-7f) continue;
            float invD = 1f / d;

            const float eps = -1e-4f; // tolerance against 1px cracks between triangles

            for (int y = minY; y <= maxY; y++)
            {
                float py = y + 0.5f;
                int rowStart = y * textureSize;
                for (int x = minX; x <= maxX; x++)
                {
                    float px = x + 0.5f;

                    float a = ((p1y - p2y) * (px - p2x) + (p2x - p1x) * (py - p2y)) * invD;
                    if (a < eps) continue;
                    float b = ((p2y - p0y) * (px - p2x) + (p0x - p2x) * (py - p2y)) * invD;
                    if (b < eps) continue;
                    float c = 1f - a - b;
                    if (c < eps) continue;

                    // 3D position of the texel via barycentric interpolation.
                    float wx = a * v0.x + b * v1.x + c * v2.x;
                    float wy = a * v0.y + b * v1.y + c * v2.y;
                    float wz = a * v0.z + b * v1.z + c * v2.z;

                    // Project the offset onto the stamp frame (right/up/normal).
                    float ddx = wx - pLocal.x, ddy = wy - pLocal.y, ddz = wz - pLocal.z;
                    float du = ddx * u.x + ddy * u.y + ddz * u.z;
                    float dv = ddx * v.x + ddy * v.y + ddz * v.z;
                    float dn = ddx * n.x + ddy * n.y + ddz * n.z;
                    if (du < -radiusLocal || du > radiusLocal) continue;
                    if (dv < -radiusLocal || dv > radiusLocal) continue;
                    if (dn < -radiusLocal || dn > radiusLocal) continue;

                    // Mask -> opacity.
                    float mu = du * invR * 0.5f + 0.5f;
                    float mv = dv * invR * 0.5f + 0.5f;
                    float maskOp = SampleMaskOpacity(mu, mv);
                    if (maskOp <= 0f) continue;

                    int idx = rowStart + x;

                    if (aAlb)
                    {
                        float op = maskOp * albOpMul;
                        if (op > 0f)
                        {
                            Color32 dst = alb.pixels[idx];
                            float it = 1f - op;
                            alb.pixels[idx] = new Color32(
                                BlendSrgb(dst.r, brLin, it, op),
                                BlendSrgb(dst.g, bgLin, it, op),
                                BlendSrgb(dst.b, bbLin, it, op),
                                255);
                            alb.painted[idx] = true;
                            ExpandDirty(alb, x, y);
                        }
                    }

                    if (aMSO)
                    {
                        // Mask map blend. Metallic(R) + smoothness(A) blend with the MS opacity,
                        // occlusion(G) with its own opacity - each = mask * the channel's opacity
                        // multiplier (the albedo brushColor.a equivalent). Untouched channels keep
                        // their previous value; B stays 0.
                        Color32 dst = mso.pixels[idx];
                        float msOp = maskOp * msOpMul;
                        float aoOp = maskOp * aoOpMul;
                        byte mr = aMS ? (byte)(dst.r * (1f - msOp) + msoR * msOp) : dst.r;
                        byte mg = aAO ? (byte)(dst.g * (1f - aoOp) + msoG * aoOp) : dst.g;
                        byte ma = aMS ? (byte)(dst.a * (1f - msOp) + msoA * msOp) : dst.a;
                        mso.pixels[idx] = new Color32(mr, mg, 0, ma);
                        mso.painted[idx] = true;
                        ExpandDirty(mso, x, y);
                    }

                    if (aEM)
                    {
                        float op = maskOp * emOpMul;
                        if (op > 0f)
                        {
                            Color32 dst = em.pixels[idx];
                            float it = 1f - op;
                            em.pixels[idx] = new Color32(
                                BlendSrgb(dst.r, erLin, it, op),
                                BlendSrgb(dst.g, egLin, it, op),
                                BlendSrgb(dst.b, ebLin, it, op),
                                255);
                            em.painted[idx] = true;
                            ExpandDirty(em, x, y);
                        }
                    }

                }
            }
        }
    }

    private static void ExpandDirty(CpuMap cm, int x, int y)
    {
        if (x < cm.dx0) cm.dx0 = x;
        if (x > cm.dx1) cm.dx1 = x;
        if (y < cm.dy0) cm.dy0 = y;
        if (y > cm.dy1) cm.dy1 = y;
    }

    /// <summary>
    /// CPU edge padding against UV seams: bleeds the freshly painted colors outward by
    /// 'pad' texels. The dilation source is only brush-painted texels (cm.painted).
    /// </summary>
    private void DilatePadding(CpuMap cm, int pad)
    {
        if (cm.dx1 < cm.dx0 || cm.dy1 < cm.dy0) return; // nothing painted

        int x0 = Mathf.Max(0, cm.dx0 - pad);
        int y0 = Mathf.Max(0, cm.dy0 - pad);
        int x1 = Mathf.Min(textureSize - 1, cm.dx1 + pad);
        int y1 = Mathf.Min(textureSize - 1, cm.dy1 + pad);

        cm.dilateStamp++;

        for (int iter = 0; iter < pad; iter++)
        {
            _dilateScratch.Clear();

            for (int y = y0; y <= y1; y++)
            {
                int row = y * textureSize;
                for (int x = x0; x <= x1; x++)
                {
                    int idx = row + x;
                    if (cm.painted[idx] || cm.fillStamp[idx] == cm.dilateStamp) continue;

                    if (TryFindPaddingSource(cm, x, y, out Color32 c))
                    {
                        cm.pixels[idx] = c;
                        _dilateScratch.Add(idx);
                        ExpandDirty(cm, x, y);
                    }
                }
            }

            for (int k = 0; k < _dilateScratch.Count; k++)
                cm.fillStamp[_dilateScratch[k]] = cm.dilateStamp;

            if (_dilateScratch.Count == 0) break;
        }
    }

    /// <summary>Searches the 8-neighborhood for a source texel (painted or already padded) and returns its color.</summary>
    private bool TryFindPaddingSource(CpuMap cm, int x, int y, out Color32 col)
    {
        for (int oy = -1; oy <= 1; oy++)
        {
            int ny = y + oy;
            if (ny < 0 || ny >= textureSize) continue;
            int nrow = ny * textureSize;
            for (int ox = -1; ox <= 1; ox++)
            {
                if (ox == 0 && oy == 0) continue;
                int nx = x + ox;
                if (nx < 0 || nx >= textureSize) continue;
                int nidx = nrow + nx;
                if (cm.painted[nidx] || cm.fillStamp[nidx] == cm.dilateStamp)
                {
                    col = cm.pixels[nidx];
                    return true;
                }
            }
        }
        col = default;
        return false;
    }

    /// <summary>Uploads only the region changed this frame to the GPU once (CPU backend).</summary>
    private void FlushDirty(CpuMap cm)
    {
        if (cm.dx1 < cm.dx0 || cm.dy1 < cm.dy0) return; // nothing painted

        int bw = cm.dx1 - cm.dx0 + 1;
        int bh = cm.dy1 - cm.dy0 + 1;
        var block = new Color32[bw * bh];
        for (int y = 0; y < bh; y++)
            System.Array.Copy(cm.pixels, (cm.dy0 + y) * textureSize + cm.dx0, block, y * bw, bw);

        cm.tex.SetPixels32(cm.dx0, cm.dy0, bw, bh, block);
        cm.tex.Apply(false);
    }
}