using System;
using System.Windows;
using System.Windows.Media;
using System.Windows.Media.Imaging;

namespace Character_Builder;

/// <summary>
/// Many part PNGs are exported as opaque RGB with a "fake transparency" checkerboard
/// (two light grays ~#FEFEFE / ~#F1F1F1) baked into the pixels. That opaque background
/// hides layers below it and breaks alpha-based neck detection.
///
/// KeyOut() removes the background by flood-filling from the image borders, so only the
/// background that is *connected to the edge* is cleared — interior white clothing (crop
/// top, hoodie) is kept. If the source already has real transparency, it is returned as-is.
/// </summary>
public static class BgKey
{
    // A pixel counts as background when it is light and near-neutral (low saturation).
    private const int BgMinBrightness = 208;   // min channel value
    private const int BgMaxSpread = 26;         // max(channel) - min(channel)

    // Boundary feather: soften light fringe pixels that touch cleared background.
    private const int FeatherLo = 200;          // fully opaque below this brightness
    private const int FeatherHi = 250;          // fully transparent at/above this brightness

    public static BitmapSource KeyOut(BitmapSource src, bool isHead = false)
    {
        BitmapSource conv = src.Format == PixelFormats.Bgra32
            ? src
            : new FormatConvertedBitmap(src, PixelFormats.Bgra32, null, 0);

        int w = conv.PixelWidth, h = conv.PixelHeight;
        if (w <= 0 || h <= 0) return src;

        int stride = w * 4;
        var px = new byte[h * stride];
        conv.CopyPixels(px, stride, 0);

        // Already has real transparency? Leave it alone.
        long transparent = 0;
        for (int i = 3; i < px.Length; i += 4)
            if (px[i] < 16) { if (++transparent > (long)w * h / 50) break; }
        if (transparent > (long)w * h / 50) return src;

        var seen = new bool[w * h];
        var stack = new int[w * h];
        int sp = 0;

        // Seed the flood fill from every border pixel that looks like background.
        for (int x = 0; x < w; x++)
        {
            TryPush(px, stride, seen, stack, ref sp, x, 0, w);
            TryPush(px, stride, seen, stack, ref sp, x, h - 1, w);
        }
        for (int y = 0; y < h; y++)
        {
            TryPush(px, stride, seen, stack, ref sp, 0, y, w);
            TryPush(px, stride, seen, stack, ref sp, w - 1, y, w);
        }

        // Flood fill (4-connected) over border-connected background; clear its alpha.
        while (sp > 0)
        {
            int idx = stack[--sp];
            int x = idx % w, y = idx / w;
            px[idx * 4 + 3] = 0; // transparent

            TryPush(px, stride, seen, stack, ref sp, x - 1, y, w);
            TryPush(px, stride, seen, stack, ref sp, x + 1, y, w);
            TryPush(px, stride, seen, stack, ref sp, x, y - 1, w);
            TryPush(px, stride, seen, stack, ref sp, x, y + 1, w);
        }

        // The baked "transparency" is a two-tone light-gray checkerboard. Loops/rings (necklace,
        // bracelet, headphone band) enclose background pockets the border fill can't reach; remove
        // them by detecting the checker pattern (both tones present, no dark object pixels nearby).
        RemoveCheckerPockets(px, w, h, stride);

        // Head portraits include a white shirt/shoulders below the neck; that interior white
        // is not border-connected, so it survives the flood fill and looks like an un-keyed
        // white patch once the head is layered on a body. Drop light interior pixels in the
        // lower part of the head silhouette (below the face, so eyes/teeth are untouched).
        if (isHead) DropHeadShirt(px, w, h, stride);

        Feather(px, w, h, stride);

        var wb = new WriteableBitmap(w, h, 96, 96, PixelFormats.Bgra32, null);
        wb.WritePixels(new Int32Rect(0, 0, w, h), px, stride, 0);
        wb.Freeze();
        return wb;
    }

    private static void TryPush(byte[] px, int stride, bool[] seen, int[] stack, ref int sp, int x, int y, int w)
    {
        int h = px.Length / stride;
        if (x < 0 || y < 0 || x >= w || y >= h) return;
        int idx = y * w + x;
        if (seen[idx]) return;
        int p = y * stride + x * 4;
        if (!IsBg(px[p], px[p + 1], px[p + 2])) return;
        seen[idx] = true;
        stack[sp++] = idx;
    }

    private static bool IsBg(byte b, byte g, byte r)
    {
        int mn = Math.Min(r, Math.Min(g, b));
        int mx = Math.Max(r, Math.Max(g, b));
        return mn >= BgMinBrightness && (mx - mn) <= BgMaxSpread;
    }

    /// <summary>Remove enclosed checkerboard-background pockets while keeping solid light objects.</summary>
    private static void RemoveCheckerPockets(byte[] px, int w, int h, int stride)
    {
        int W1 = w + 1;
        var li = new int[(h + 1) * W1];   // integral of "light" checker tone (~254)
        var mi = new int[(h + 1) * W1];   // integral of "mid" checker tone (~241)
        var di = new int[(h + 1) * W1];   // integral of dark (object) pixels
        var neutral = new bool[w * h];

        for (int y = 0; y < h; y++)
        {
            int row = y * stride, cur = (y + 1) * W1, prev = y * W1;
            int la = 0, ma = 0, da = 0;
            for (int x = 0; x < w; x++)
            {
                int p = row + x * 4; int b = px[p], g = px[p + 1], r = px[p + 2];
                int mn = Math.Min(r, Math.Min(g, b)), mx = Math.Max(r, Math.Max(g, b));
                double v = (r + g + b) / 3.0;
                bool nt = mn >= 224 && (mx - mn) <= 14;
                neutral[y * w + x] = nt;
                if (nt && v >= 249) la++;
                if (nt && v >= 234 && v <= 246) ma++;
                if (v < 205) da++;
                li[cur + x + 1] = li[prev + x + 1] + la;
                mi[cur + x + 1] = mi[prev + x + 1] + ma;
                di[cur + x + 1] = di[prev + x + 1] + da;
            }
        }

        const int R = 18;
        int area = (2 * R + 1) * (2 * R + 1);
        int need = (int)(0.06 * area);
        for (int y = 0; y < h; y++)
        {
            int row = y * stride;
            for (int x = 0; x < w; x++)
            {
                int p = row + x * 4;
                if (px[p + 3] == 0 || !neutral[y * w + x]) continue;
                int y0 = Math.Max(0, y - R), y1 = Math.Min(h, y + R + 1);
                int x0 = Math.Max(0, x - R), x1 = Math.Min(w, x + R + 1);
                if (Win(di, W1, y0, y1, x0, x1) > 2) continue;      // an object edge is near → keep
                if (Win(li, W1, y0, y1, x0, x1) >= need && Win(mi, W1, y0, y1, x0, x1) >= need)
                    px[p + 3] = 0;                                   // both checker tones present → background
            }
        }
    }

    private static int Win(int[] I, int W1, int y0, int y1, int x0, int x1)
        => I[y1 * W1 + x1] - I[y0 * W1 + x1] - I[y1 * W1 + x0] + I[y0 * W1 + x0];

    /// <summary>Clear light, low-saturation interior pixels (the shirt) in the lower head region.</summary>
    private static void DropHeadShirt(byte[] px, int w, int h, int stride)
    {
        int minY = int.MaxValue, maxY = -1;
        for (int y = 0; y < h; y++)
        {
            int row = y * stride;
            for (int x = 0; x < w; x++)
                if (px[row + x * 4 + 3] >= 24) { if (y < minY) minY = y; if (y > maxY) maxY = y; break; }
        }
        if (maxY < 0) return;
        int yThr = minY + (int)((maxY - minY) * 0.58); // below the face

        for (int y = yThr; y < h; y++)
        {
            int row = y * stride;
            for (int x = 0; x < w; x++)
            {
                int p = row + x * 4;
                if (px[p + 3] == 0) continue;
                int b = px[p], g = px[p + 1], r = px[p + 2];
                int mn = Math.Min(r, Math.Min(g, b));
                int mx = Math.Max(r, Math.Max(g, b));
                if (mn >= 222 && (mx - mn) <= 16) px[p + 3] = 0;
            }
        }
    }

    /// <summary>Soften light, low-saturation pixels that border a cleared area (anti-halo).</summary>
    private static void Feather(byte[] px, int w, int h, int stride)
    {
        for (int y = 0; y < h; y++)
        {
            int row = y * stride;
            for (int x = 0; x < w; x++)
            {
                int p = row + x * 4;
                if (px[p + 3] == 0) continue;

                bool touchesCleared =
                    (x > 0 && px[p - 4 + 3] == 0) ||
                    (x < w - 1 && px[p + 4 + 3] == 0) ||
                    (y > 0 && px[p - stride + 3] == 0) ||
                    (y < h - 1 && px[p + stride + 3] == 0);
                if (!touchesCleared) continue;

                int b = px[p], g = px[p + 1], r = px[p + 2];
                int mn = Math.Min(r, Math.Min(g, b));
                int mx = Math.Max(r, Math.Max(g, b));
                if (mn < FeatherLo || (mx - mn) > BgMaxSpread) continue;

                int t = Math.Min(255, Math.Max(0, (mn - FeatherLo) * 255 / (FeatherHi - FeatherLo)));
                px[p + 3] = (byte)(255 - t);
            }
        }
    }
}