using System;
using UnityEditor;
using UnityEngine;
using Unity.InferenceEngine;

namespace VibeCopilot.Editor
{
    // Scoaterea fundalului (matting) LOCAL, gratis, cu Unity AI Inference Engine (fost Sentis).
    // Model: u2netp (U^2-Net small, licenta Apache-2.0 = ok comercial) — o retea care intelege
    // "obiect vs fundal", deci merge pe ORICE fundal (spre deosebire de chroma key, care cerea
    // un fundal uniform pe care AI-ul refuza sa-l faca). Ruleaza o data pe poza de concept in editor.
    public static class VibeMatting
    {
        private const string ModelPath = "Assets/VibeCopilot/Models/u2netp.onnx";
        private const int S = 320;   // u2netp cere intrare 320x320
        // Normalizarea ImageNet cu care a fost antrenat modelul (obligatorie pt masca buna).
        private static readonly float[] Mean = { 0.485f, 0.456f, 0.406f };
        private static readonly float[] Std = { 0.229f, 0.224f, 0.225f };

        private static Model _model;

        private static Model LoadModel()
        {
            if (_model != null) return _model;
            var asset = AssetDatabase.LoadAssetAtPath<ModelAsset>(ModelPath);
            if (asset == null)
            {
                Debug.LogWarning("[VibeCopilot] Matting model not found at " + ModelPath + " (skipping cutout).");
                return null;
            }
            _model = ModelLoader.Load(asset);
            return _model;
        }

        // Intoarce o COPIE a texturii cu fundalul scos (alpha din masca modelului).
        // La orice eroare / model lipsa -> intoarce textura originala (nu stricam fluxul).
        public static Texture2D Cutout(Texture2D src)
        {
            if (src == null) return src;
            Model model = LoadModel();
            if (model == null) return src;

            Color32[] px;
            try { px = src.GetPixels32(); }
            catch { return src; }   // textura ne-citibila
            int w = src.width, h = src.height;
            if (w <= 0 || h <= 0) return src;

            Worker worker = null;
            Tensor<float> input = null;
            try
            {
                // ---- preprocesare: resize la 320x320, RGB, normalizare ImageNet, layout NCHW ----
                float[] data = new float[3 * S * S];
                int plane = S * S;
                for (int ty = 0; ty < S; ty++)
                {
                    int sy = (int)((ty + 0.5f) * h / S); if (sy >= h) sy = h - 1;
                    for (int tx = 0; tx < S; tx++)
                    {
                        int sx = (int)((tx + 0.5f) * w / S); if (sx >= w) sx = w - 1;
                        Color32 c = px[sy * w + sx];
                        int o = ty * S + tx;
                        data[o] = (c.r / 255f - Mean[0]) / Std[0];               // canal R
                        data[plane + o] = (c.g / 255f - Mean[1]) / Std[1];        // canal G
                        data[2 * plane + o] = (c.b / 255f - Mean[2]) / Std[2];    // canal B
                    }
                }

                input = new Tensor<float>(new TensorShape(1, 3, S, S), data);
                worker = new Worker(model, BackendType.GPUCompute);
                worker.Schedule(input);

                var output = worker.PeekOutput() as Tensor<float>;
                if (output == null) return src;
                float[] mask = output.DownloadToArray();   // (1,1,320,320) -> 320*320 valori
                if (mask == null || mask.Length < plane) return src;

                // Masca u2net nu e in [0,1] garantat -> normalizare min-max in nrm[] (320x320).
                float mn = float.MaxValue, mx = float.MinValue;
                for (int i = 0; i < plane; i++) { float v = mask[i]; if (v < mn) mn = v; if (v > mx) mx = v; }
                float range = Mathf.Max(1e-5f, mx - mn);
                float[] nrm = new float[plane];
                for (int i = 0; i < plane; i++) nrm[i] = (mask[i] - mn) / range;

                // AUTO-INVERT: fundalul e mereu pe MARGINILE imaginii. Vrem fundal = alpha MIC.
                // Daca media pe marginea mastii e MARE, modelul a scos-o invers -> o inversam.
                // (Asa reparam cazul in care obiectul iesea transparent in loc de fundal.)
                double edge = 0; int ec = 0;
                for (int i = 0; i < S; i++) { edge += nrm[i]; edge += nrm[(S - 1) * S + i]; edge += nrm[i * S]; edge += nrm[i * S + (S - 1)]; ec += 4; }
                bool invert = (edge / ec) > 0.5;
                if (invert) for (int i = 0; i < plane; i++) nrm[i] = 1f - nrm[i];

                // DOAR CEA MAI MARE BUCATA: pastram obiectul principal, aruncam "fantomele" (ex. un
                // al doilea avion sters in fundal pe care Nano Banana l-a pus). Binarizam, gasim
                // componenta conexa cea mai mare (4-vecini), o dilatam putin (ca sa nu taiem marginea
                // moale a obiectului), si stergem tot ce e in afara.
                {
                    const float T = 0.45f;
                    int[] label = new int[plane];
                    int[] stack = new int[plane];
                    int best = 0, bestSize = 0, cur = 0;
                    for (int i = 0; i < plane; i++)
                    {
                        if (nrm[i] <= T || label[i] != 0) continue;
                        cur++; int sp = 0; stack[sp++] = i; label[i] = cur; int size = 0;
                        while (sp > 0)
                        {
                            int p = stack[--sp]; size++;
                            int py = p / S, pxi = p - py * S;
                            if (pxi > 0)     { int q = p - 1; if (nrm[q] > T && label[q] == 0) { label[q] = cur; stack[sp++] = q; } }
                            if (pxi < S - 1) { int q = p + 1; if (nrm[q] > T && label[q] == 0) { label[q] = cur; stack[sp++] = q; } }
                            if (py > 0)      { int q = p - S; if (nrm[q] > T && label[q] == 0) { label[q] = cur; stack[sp++] = q; } }
                            if (py < S - 1)  { int q = p + S; if (nrm[q] > T && label[q] == 0) { label[q] = cur; stack[sp++] = q; } }
                        }
                        if (size > bestSize) { bestSize = size; best = cur; }
                    }
                    if (best > 0)
                    {
                        const int R = 3;
                        bool[] keep = new bool[plane];
                        for (int y = 0; y < S; y++)
                            for (int x = 0; x < S; x++)
                            {
                                if (label[y * S + x] != best) continue;
                                int ay0 = Mathf.Max(0, y - R), ay1 = Mathf.Min(S - 1, y + R);
                                int ax0 = Mathf.Max(0, x - R), ax1 = Mathf.Min(S - 1, x + R);
                                for (int yy = ay0; yy <= ay1; yy++)
                                    for (int xx = ax0; xx <= ax1; xx++)
                                        keep[yy * S + xx] = true;
                            }
                        for (int i = 0; i < plane; i++) if (!keep[i]) nrm[i] = 0f;
                    }
                }

                // ---- aplicam masca drept alpha pe textura originala (bilinear = margini netede) ----
                Color32[] outPx = new Color32[px.Length];
                long fg = 0;   // cati pixeli raman OPACI -> ca sa prindem o masca degenerata
                for (int y = 0; y < h; y++)
                {
                    float fy = (y + 0.5f) * S / (float)h - 0.5f;
                    int y0 = Mathf.Clamp((int)Mathf.Floor(fy), 0, S - 1);
                    int y1 = Mathf.Min(y0 + 1, S - 1);
                    float wy = Mathf.Clamp01(fy - y0);
                    for (int x = 0; x < w; x++)
                    {
                        float fx = (x + 0.5f) * S / (float)w - 0.5f;
                        int x0 = Mathf.Clamp((int)Mathf.Floor(fx), 0, S - 1);
                        int x1 = Mathf.Min(x0 + 1, S - 1);
                        float wx = Mathf.Clamp01(fx - x0);

                        float m = Mathf.Lerp(
                            Mathf.Lerp(nrm[y0 * S + x0], nrm[y0 * S + x1], wx),
                            Mathf.Lerp(nrm[y1 * S + x0], nrm[y1 * S + x1], wx), wy);

                        // Prag mai strans: zonele slabe (pata din fundalul dintre aripi) -> transparent,
                        // pastrand obiectul opac. (u2netp nu e perfect pe zone complexe.)
                        m = Mathf.SmoothStep(0f, 1f, Mathf.Clamp01((m - 0.45f) / 0.22f));
                        if (m > 0.5f) fg++;

                        Color32 c = px[y * w + x];
                        c.a = (byte)(m * 255f);
                        outPx[y * w + x] = c;
                    }
                }

                // PLASA DE SIGURANTA: daca masca e degenerata (aproape tot transparent sau aproape
                // tot opac = n-a taiat nimic util) -> aratam poza ORIGINALA (nu riscam invizibil).
                long total = (long)w * h;
                if (fg < total * 0.02 || fg > total * 0.98)
                {
                    Debug.LogWarning("[VibeCopilot] Matting mask degenerate (fg=" + fg + "/" + total + ") — showing image as-is.");
                    return src;
                }

                // mipmaps + filtrare trilineara -> afisare NETEDA cand cardul e mai mic decat poza
                Texture2D outTex = new Texture2D(w, h, TextureFormat.RGBA32, true);
                outTex.SetPixels32(outPx);
                outTex.Apply(true);
                outTex.filterMode = FilterMode.Trilinear;
                outTex.anisoLevel = 4;
                return outTex;
            }
            catch (Exception e)
            {
                Debug.LogWarning("[VibeCopilot] Matting failed (" + e.Message + ") — showing image as-is.");
                return src;
            }
            finally
            {
                if (input != null) input.Dispose();
                if (worker != null) worker.Dispose();
            }
        }

        // DIFFERENCE MATTING: doua randari ale ACELUIASI subiect, una pe ALB una pe NEGRU (fundalul
        // se anuleaza matematic) -> alpha EXACT, margini/monturi perfecte, mai bun ca u2netp.
        // Compositing: negru b = a*F ; alb w = a*F + (1-a)  =>  a = 1 - (w - b),  F = b/a.
        // Plasa de siguranta: la nealiniere / rezultat ciudat -> cade pe Cutout(white) (u2netp).
        public static Texture2D DifferenceMatte(Texture2D white, Texture2D black)
        {
            if (white == null) return white;
            if (black == null) return Cutout(white);
            if (black.width != white.width || black.height != white.height) return Cutout(white);
            Color32[] w, b;
            try { w = white.GetPixels32(); b = black.GetPixels32(); } catch { return Cutout(white); }
            if (w.Length != b.Length || w.Length == 0) return Cutout(white);

            int W = white.width, H = white.height;
            Color32[] outPx = new Color32[w.Length];
            long fgCount = 0;
            for (int i = 0; i < w.Length; i++)
            {
                float dr = (w[i].r - b[i].r) / 255f;
                float dg = (w[i].g - b[i].g) / 255f;
                float db = (w[i].b - b[i].b) / 255f;
                float a = Mathf.Clamp01(1f - (dr + dg + db) / 3f);
                // F = negru / alpha (unpremultiplied) -> culoare curata pe margini, fara halou alb
                float inv = a > 0.004f ? 1f / a : 0f;
                byte fr = (byte)Mathf.Clamp(b[i].r * inv, 0f, 255f);
                byte fgc = (byte)Mathf.Clamp(b[i].g * inv, 0f, 255f);
                byte fb = (byte)Mathf.Clamp(b[i].b * inv, 0f, 255f);
                if (a > 0.5f) fgCount++;
                outPx[i] = new Color32(fr, fgc, fb, (byte)(a * 255f));
            }

            // Colturile = fundal -> trebuie sa fie transparente. Daca nu (subiect mutat la editare) =>
            // nealiniere -> rezultatul ar fi fantomatic -> cadem pe u2netp. Idem daca e degenerat.
            long total = w.Length;
            float corner = (outPx[0].a + outPx[W - 1].a + outPx[(H - 1) * W].a + outPx[w.Length - 1].a) / (4f * 255f);
            if (corner > 0.35f || fgCount < total * 0.01 || fgCount > total * 0.99)
            {
                Debug.LogWarning("[VibeCopilot] DifferenceMatte off (corner=" + corner.ToString("0.00") + ", fg=" + fgCount + "/" + total + ") — using u2netp.");
                return Cutout(white);
            }

            Texture2D outTex = new Texture2D(W, H, TextureFormat.RGBA32, true);
            outTex.SetPixels32(outPx);
            outTex.Apply(true);
            outTex.filterMode = FilterMode.Trilinear;
            outTex.anisoLevel = 4;
            return outTex;
        }
    }
}
