import os, io, traceback, datetime

# مسیر cache مدل (قابل نوشتن)
os.environ.setdefault("U2NET_HOME", "/home/alireaco/rembg/.u2net")
os.environ.setdefault("XDG_CACHE_HOME", "/home/alireaco/rembg/.cache")

# محدود کردن thread ها برای هاست اشتراکی
os.environ.setdefault("OMP_NUM_THREADS", "1")
os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
os.environ.setdefault("MKL_NUM_THREADS", "1")
os.environ.setdefault("NUMEXPR_NUM_THREADS", "1")
os.environ.setdefault("VECLIB_MAXIMUM_THREADS", "1")
os.environ.setdefault("ONNXRUNTIME_INTRA_OP_NUM_THREADS", "1")
os.environ.setdefault("ONNXRUNTIME_INTER_OP_NUM_THREADS", "1")

from flask import Flask, request, jsonify, Response
from PIL import Image

app = Flask(__name__)

# ثابت‌ها (طبق خواسته شما)
OUT_SIZE = 1200
MARGIN = 0.08          # 8% حاشیه
JPG_QUALITY = 92

_SESSION = None

def get_session():
    global _SESSION
    if _SESSION is None:
        from rembg.session_factory import new_session
        _SESSION = new_session("u2netp")
    return _SESSION

def log_exc():
    with open("/home/alireaco/rembg/app_error.log", "a", encoding="utf-8") as f:
        f.write("\n\n=== %s ===\n" % datetime.datetime.now().isoformat())
        f.write(traceback.format_exc())

def square_white_jpg(img_rgba: Image.Image) -> bytes:
    # بوم سفید
    canvas = Image.new("RGBA", (OUT_SIZE, OUT_SIZE), (255, 255, 255, 255))

    # باندینگ باکس سوژه بر اساس آلفا
    alpha = img_rgba.getchannel("A")
    bbox = alpha.getbbox()
    if not bbox:
        # اگر چیزی نبود، خروجی سفید بده
        out = io.BytesIO()
        canvas.convert("RGB").save(out, format="JPEG", quality=JPG_QUALITY, optimize=True)
        return out.getvalue()

    cut = img_rgba.crop(bbox)
    cw, ch = cut.size

    inner = int(OUT_SIZE * (1.0 - 2.0 * MARGIN))
    inner = max(32, min(OUT_SIZE, inner))

    scale = min(inner / cw, inner / ch)
    nw, nh = max(1, int(cw * scale)), max(1, int(ch * scale))
    cut = cut.resize((nw, nh), Image.LANCZOS)

    x = (OUT_SIZE - nw) // 2
    y = (OUT_SIZE - nh) // 2
    canvas.alpha_composite(cut, (x, y))

    out = io.BytesIO()
    canvas.convert("RGB").save(out, format="JPEG", quality=JPG_QUALITY, optimize=True)
    return out.getvalue()

@app.get("/")
def home():
    return "OK - background removal + 1200x1200 white JPG"

@app.get("/health")
def health():
    return "ok"

@app.post("/api/remove")
def remove_bg_square():
    try:
        if "image" not in request.files:
            return jsonify({"error": "form-data field 'image' is required"}), 400

        f = request.files["image"]
        original = Image.open(f.stream).convert("RGBA")

        # rembg ورودی بایت PNG می‌خواهد
        buf = io.BytesIO()
        original.save(buf, format="PNG")
        inp = buf.getvalue()

        from rembg.bg import remove
        out_bytes = remove(inp, session=get_session())
        cutout = Image.open(io.BytesIO(out_bytes)).convert("RGBA")

        jpg_bytes = square_white_jpg(cutout)
        return Response(jpg_bytes, mimetype="image/jpeg")

    except Exception:
        log_exc()
        return "ERROR. Check /home/alireaco/rembg/app_error.log", 500

application = app