"""BG Remover Web - image engine.

The image functions below (remove_solid, remove_ink, analyze_image, refine_edges ...) are copied unchanged from the
desktop app, so the results are the same. Only the AI model runner and the per-image pipeline (run) are new.
"""
import os
import threading

import numpy as np
from PIL import Image, ImageDraw, ImageFilter

import config

os.environ.setdefault("U2NET_HOME", config.MODELS_DIR)

MODEL_NAME = "u2net"
MODE_AUTO, MODE_AI, MODE_SOLID, MODE_INK = "auto", "ai", "solid", "ink"


def available_models():
    """models folder-e jototuku .onnx file ache sheigulo."""
    try:
        names = sorted(f[:-5] for f in os.listdir(config.MODELS_DIR) if f.lower().endswith(".onnx"))
    except Exception:
        names = []
    return names or [MODEL_NAME]


# ======== copied unchanged from the desktop app ========
def alpha_cov(img):
    """Result-er koto ongsho dekha jachhe (0..1). 0 mane shob faka."""
    a = np.asarray(img.getchannel("A"))
    return float((a > 127).mean())


def dominant_bg(img):
    """Chhobir shobcheye beshi thaka rong (kinar-e onno jinis thakleo background thik dhore)."""
    im = img.convert("RGB")
    im.thumbnail((256, 256))
    arr = np.asarray(im, dtype=np.int16)
    q = arr >> 4
    key = q[..., 0] * 256 + q[..., 1] * 16 + q[..., 2]
    top = int(np.bincount(key.ravel(), minlength=4096).argmax())
    return np.median(arr[key == top].reshape(-1, 3), axis=0).astype(np.float32)


def remove_solid(img, tol, bg=None):
    """Plain background (shada ba onno solid color) puro image theke remove kore,
    stamp-er bhitorer faka jayga shoho. tol: 1..100"""
    rgb = np.asarray(img.convert("RGB"), dtype=np.float32)
    if bg is None:
        border = np.concatenate([rgb[0, :, :], rgb[-1, :, :], rgb[:, 0, :], rgb[:, -1, :]], axis=0)
        bg = np.median(border, axis=0)                  # background color (kinar theke)
    d = np.max(np.abs(rgb - bg), axis=2)                # 0..255
    hi = max(float(tol) * 2.55, 8.0)
    lo = hi * 0.35
    alpha = np.clip((d - lo) / (hi - lo), 0.0, 1.0)
    a3 = np.maximum(alpha, 1e-3)[..., None]
    color = np.where(alpha[..., None] < 1.0, bg + (rgb - bg) / a3, rgb)  # edge-er color thik kora
    color = np.clip(color, 0, 255)
    out = np.dstack([color, alpha * 255.0]).astype(np.uint8)
    return Image.fromarray(out, "RGBA")


def _lum(rgb):
    return rgb[..., 0] * 0.299 + rgb[..., 1] * 0.587 + rgb[..., 2] * 0.114


def _ink_stats(rgb):
    """Kagoj-er rong, kagoj-er ujjolota ar kalir (lekha) gaarota ber kore."""
    L = _lum(rgb)
    m = L >= np.percentile(L, 60)
    paper = np.median(rgb[m].reshape(-1, 3), axis=0)
    pl = float(_lum(paper[None, None, :])[0, 0])
    ink = max(float(np.percentile(L, 1)), 0.0)
    contrast = max(pl - ink, 30.0)
    return L, paper, pl, ink, contrast


def ink_kept(img, raw):
    """AI result lekha-ke (tarikh/signature/text) koto ta rekheche: 0..1. None = eta kagoj-er lekha-r chhobi noy."""
    try:
        w = min(img.width, 512)
        h = max(1, int(round(img.height * w / img.width)))
        rgb = np.asarray(img.convert("RGB").resize((w, h)), dtype=np.float32)
        L, paper, pl, ink, contrast = _ink_stats(rgb)
        if pl < 140:                                   # kagoj ujjol noy
            return None
        mask = (pl - L) > contrast * 0.5               # gaaro (lekha) pixel
        frac = float(mask.mean())
        if frac < 0.002 or frac > 0.35:                # lekha nei, ba chhobir onek ongsho gaaro (photo)
            return None
        a = np.asarray(raw.convert("RGBA").getchannel("A").resize((w, h)), dtype=np.float32) / 255.0
        return float(a[mask].mean())                   # ghost/halka hole 1-er cheye onek kom
    except Exception:
        return None


def ink_dropped(img, raw):
    """True hole AI lekha muchhe/halka kore felechhe -> Ink mode dorkar."""
    k = ink_kept(img, raw)
    return k is not None and k < 0.85


def remove_ink(img, strength=60):
    """Kagoj-er opor kalir/pen-er lekha (tarikh, signature) rakhe, kagoj ar halka rong (golapi/nil) muchhe dey.
    strength: 5..90 (beshi hole aro halka dag muchhe jay)."""
    rgb = np.asarray(img.convert("RGB"), dtype=np.float32)
    L, paper, pl, ink, contrast = _ink_stats(rgb)
    d = pl - L
    s = float(strength) / 100.0
    lo = contrast * (0.12 + 0.45 * s)
    hi = lo + contrast * 0.22
    alpha = np.clip((d - lo) / (hi - lo), 0.0, 1.0) ** 0.7
    dark = rgb * (ink / max(pl, 1.0) + 0.35)
    color = np.clip(np.where(alpha[..., None] < 1.0, rgb * 0.55 + dark * 0.45, rgb), 0, 255)
    out = np.dstack([color, alpha * 255.0]).astype(np.uint8)
    return Image.fromarray(out, "RGBA")


def _border(arr):
    return np.concatenate([arr[0, :, :], arr[-1, :, :], arr[:, 0, :], arr[:, -1, :]], axis=0)


def analyze_image(img):
    """Chhobi dekhe bujhe: seal/logo/signature (solid) na photo (AI), ar manush ache kina."""
    im = img.convert("RGB")
    im.thumbnail((384, 384))
    arr = np.asarray(im, dtype=np.int16)
    border = _border(arr)
    bg = np.median(border, axis=0)
    dev = np.abs(border - bg).max(axis=1)
    plain_bg = float(np.percentile(dev, 60)) <= 40       # kinar-er beshirbhag ek-i rong

    ycc = np.asarray(im.convert("YCbCr"), dtype=np.int16)
    cb, cr = ycc[..., 1], ycc[..., 2]
    skin = (cb >= 77) & (cb <= 127) & (cr >= 133) & (cr <= 173) & (ycc[..., 0] > 40)
    ratio = img.width / max(1, img.height)
    # lomba/chokro strip (jemon haater lekha tarikh, signature) manush-er photo noy, golapi kagoj-o skin-er moto dekhay
    q = (arr >> 4).reshape(-1, 3)
    key = q[:, 0] * 256 + q[:, 1] * 16 + q[:, 2]
    counts = np.sort(np.bincount(key, minlength=4096))[::-1]
    flat8 = float(counts[:8].sum() / key.size) >= 0.80       # khub kom rong = stamp/seal/logo, manush noy
    person = float(skin.mean()) >= 0.04 and 0.4 <= ratio <= 2.5 and not flat8

    # kom rong = seal/logo/signature (scan/JPEG noise thakleo top 24 rong 90%+ hole-o hobe)
    few_colors = flat8 or float(counts[:24].sum() / key.size) >= 0.90

    return {"solid": bool(plain_bg and (not person) and few_colors), "person": bool(person)}


def auto_tolerance(img):
    """Background-er noise (JPEG dag) dekhe Solid mode-er tolerance thik kore."""
    rgb = np.asarray(img.convert("RGB"), dtype=np.float32)
    border = _border(rgb)
    bg = np.median(border, axis=0)
    noise = float(np.percentile(np.max(np.abs(border - bg), axis=1), 90))
    hi = (noise + 3.0) / 0.35
    return float(min(60.0, max(25.0, hi / 2.55)))


def pick_model(person):
    """Je model-gulo models folder-e ache tar moddhe shobcheye upojogi ta bachhe."""
    names = available_models()
    prefs = (["u2net_human_seg", "isnet-general-use", "u2net"] if person
             else ["isnet-general-use", "u2net", "u2net_human_seg"])
    for n in prefs:
        if n in names:
            return n
    return names[0]


def defringe(img, band=None):
    """Edge-er pixel-e original background-er (shada) rong lege thake -> Photoshop-e nil/onno background-er
    opor shada rekha/dag dekhay. Edge-er rong bhitorer asol (chul/shorir-er) rong diye bodle dey."""
    try:
        if band is None:
            band = int(min(6, max(2, round(max(img.size) / 400))))
        a = img.getchannel("A")
        core = a.filter(ImageFilter.MinFilter(2 * band + 1)).point(lambda v: 255 if v >= 250 else 0)
        m = np.asarray(core, dtype=np.float32) / 255.0
        if float(m.sum()) < 10:
            return img
        rgb = np.asarray(img.convert("RGB"), dtype=np.float32)
        mimg = Image.fromarray((m * 255).astype(np.uint8))
        prem = [Image.fromarray(np.clip(rgb[..., c] * m, 0, 255).astype(np.uint8)) for c in range(3)]
        est = rgb.copy()
        todo = (m < 0.5) & (np.asarray(a) > 0)
        for rad in (band * 1.5, band * 4.0, band * 9.0):   # chhoto theke boro radius-e bhitorer rong chhorie dey
            den = np.asarray(mimg.filter(ImageFilter.GaussianBlur(rad)), dtype=np.float32) / 255.0
            ok = todo & (den > 0.1)
            if ok.any():
                num = np.stack([np.asarray(p.filter(ImageFilter.GaussianBlur(rad)), dtype=np.float32)
                                for p in prem], axis=-1)
                est = np.where(ok[..., None], np.clip(num / np.maximum(den, 1e-3)[..., None], 0, 255), est)
                todo = todo & ~ok
            if not todo.any():
                break
        out = est.astype(np.uint8)
        res = Image.fromarray(out, "RGB").convert("RGBA")
        res.putalpha(a)
        return res
    except Exception:
        return img


def refine_edges(img, shrink, smooth):
    """AI result-er edge porishkar kore: shrink = edge ke bhitore-r dike kombe (kalo outline/halo jay),
    smooth = edge naram/gol kore (khoshkhoshe bhab jay)."""
    shrink = int(shrink)
    smooth = float(smooth)
    if shrink <= 0 and smooth <= 0:
        return defringe(img)
    r, g, b, a = img.split()
    if shrink > 0:
        a = a.filter(ImageFilter.MinFilter(2 * shrink + 1))
    if smooth > 0:
        a = a.filter(ImageFilter.GaussianBlur(smooth))
        lut = [max(0, min(255, int((v - 40) * 255 / 175))) for v in range(256)]
        a = a.point(lut)
    return defringe(Image.merge("RGBA", (r, g, b, a)))


# ======== new for the website: AI model runner ========
try:
    import fcntl
except ImportError:                       # not Linux (only used when testing on another OS)
    fcntl = None

_sessions = {}
_thread_lock = threading.Lock()


class Busy:
    """Only one image is processed at a time (all processes), so 1.5 GB of RAM is never exceeded."""

    def __enter__(self):
        _thread_lock.acquire()
        self.f = None
        if fcntl:
            try:
                os.makedirs(config.DATA_DIR, exist_ok=True)
                self.f = open(os.path.join(config.DATA_DIR, "engine.lock"), "w")
                fcntl.flock(self.f, fcntl.LOCK_EX)
            except Exception:
                self.f = None
        return self

    def __exit__(self, *a):
        try:
            if self.f:
                fcntl.flock(self.f, fcntl.LOCK_UN)
                self.f.close()
        finally:
            _thread_lock.release()


class OnnxSession:
    """Runs a U2-Net / ISNet .onnx file with plain onnxruntime (same steps rembg uses, but a much lighter install)."""

    def __init__(self, model):
        import onnxruntime as ort
        path = os.path.join(config.MODELS_DIR, model + ".onnx")
        if not os.path.isfile(path):
            raise FileNotFoundError("AI model file not found on the server: models/" + model + ".onnx")
        so = ort.SessionOptions()
        so.intra_op_num_threads = int(config.ONNX_THREADS)
        so.enable_cpu_mem_arena = False          # lower RAM use
        so.enable_mem_pattern = False
        self.sess = ort.InferenceSession(path, sess_options=so, providers=["CPUExecutionProvider"])
        self.input = self.sess.get_inputs()[0].name
        if model.startswith("isnet"):
            self.size, self.std = 1024, (1.0, 1.0, 1.0)
        else:
            self.size, self.std = 320, (0.229, 0.224, 0.225)
        self.mean = (0.485, 0.456, 0.406)

    def mask(self, img):
        im = img.convert("RGB").resize((self.size, self.size), Image.LANCZOS)
        a = np.asarray(im, dtype=np.float32)
        a = a / max(float(a.max()), 1e-6)
        for c in range(3):
            a[..., c] = (a[..., c] - self.mean[c]) / self.std[c]
        x = np.expand_dims(a.transpose((2, 0, 1)), 0).astype(np.float32)
        pred = self.sess.run(None, {self.input: x})[0][:, 0, :, :]
        ma, mi = float(pred.max()), float(pred.min())
        pred = (pred - mi) / max(ma - mi, 1e-6)
        pred = np.squeeze(pred)
        m = Image.fromarray((pred * 255).astype("uint8"), mode="L")
        return m.resize(img.size, Image.LANCZOS)


def _use_rembg():
    if config.ENGINE == "onnx":
        return False
    try:
        import rembg  # noqa: F401
        return True
    except Exception:
        if config.ENGINE == "rembg":
            raise
        return False


def get_session(model):
    if model not in _sessions:
        _sessions.clear()                          # keep only ONE model in memory (RAM)
        import gc
        gc.collect()
        if _use_rembg():
            from rembg import new_session
            _sessions[model] = ("rembg", new_session(model))
        else:
            _sessions[model] = ("onnx", OnnxSession(model))
    return _sessions[model]


def ai_remove(img, model):
    """img RGBA -> RGBA with the background removed (same output form as rembg.remove)."""
    kind, sess = get_session(model)
    if kind == "rembg":
        from rembg import remove
        return remove(img, session=sess)
    m = sess.mask(img)
    empty = Image.new("RGBA", img.size, 0)
    return Image.composite(img.convert("RGBA"), empty, m)


def alt_models(current):
    names = available_models()
    return [n for n in ("isnet-general-use", "u2net") if n in names and n != current]


def run(img, mode, tol, model, shrink, smooth, tol_manual=False, model_manual=False):
    """The per-image pipeline of the desktop app. Returns (result RGBA, raw AI RGBA or None, description)."""
    auto = mode == MODE_AUTO
    if auto:
        info = analyze_image(img)
        solid = info["solid"]
        model_i = model if model_manual else pick_model(info["person"])
        tol_i = tol if tol_manual else auto_tolerance(img)
        picked = "Auto: Solid color (seal/logo)" if solid else "Auto: AI " + model_i
    else:
        solid = mode in (MODE_SOLID, MODE_INK)
        model_i, tol_i = model, tol
        picked = "Ink / handwriting" if mode == MODE_INK else ("Solid color" if solid else "AI " + model_i)
    if mode == MODE_INK:
        raw = None
        out = remove_ink(img, tol if tol_manual else 60)
    elif solid:
        raw = None
        out = remove_solid(img, tol_i)
        if auto:
            cov = alpha_cov(out)
            if cov < 0.0003 or cov > 0.995:       # shob faka ba kichui muchhe nai -> bhul background dhorechhe
                try:
                    out2 = remove_solid(img, tol_i, bg=dominant_bg(img))
                    if 0.0003 <= alpha_cov(out2) <= 0.995:
                        out = out2
                except Exception:
                    pass
    else:
        raw = ai_remove(img, model_i)
        if alpha_cov(raw) < 0.01:                 # AI kichui pay nai -> onno model, tarpor solid
            alt_raw = None
            for alt in alt_models(model_i):
                try:
                    r2 = ai_remove(img, alt)
                except Exception:
                    continue
                if 0.01 <= alpha_cov(r2) <= 0.98:
                    alt_raw, model_i = r2, alt
                    break
            if alt_raw is not None:
                raw = alt_raw
                picked = ("Auto: AI " if auto else "AI ") + model_i
            else:
                raw = None
        kept = ink_kept(img, raw) if (raw is not None and auto) else None
        if kept is not None and kept < 0.85:
            raw = None
            out = remove_ink(img, tol if tol_manual else 60)
            picked = "Auto: Ink / handwriting (AI kept only %d%% of writing)" % int(kept * 100)
        elif raw is None:
            out = remove_solid(img, tol_i, bg=dominant_bg(img))
            picked = "Solid color (AI found no subject)"
        else:
            out = refine_edges(raw, shrink, smooth)
            if kept is not None:
                picked += " (writing kept %d%%)" % int(kept * 100)
    return out, raw, picked


def compose_bg(img, rgb):
    base = Image.new("RGBA", img.size, tuple(rgb) + (255,))
    base.alpha_composite(img.convert("RGBA"))
    return base.convert("RGB")


def dpi_for(width_px, inch):
    """Save-er somoy DPI thik kore, jate Word-e chhobi dorkari size-e boshe."""
    try:
        inch = float(str(inch).strip())
        if inch > 0:
            d = max(1.0, width_px / inch)
            return (d, d)
    except Exception:
        pass
    return (300, 300)