#!/usr/bin/env python3
"""Re-cut problem items: pad canvas, stronger erosion, more blur."""
import sys, os
import numpy as np
from PIL import Image, ImageFilter
from rembg import remove, new_session
import gc

session = new_session('isnet-general-use')
OUT = '/root/.hermes/image_cache/zrd_1x1/cutouts/'

def cut(src, out, is_model, erode=11, blur=3.5, pad=40):
    cell = Image.open(src).convert('RGB')
    # pad with white to avoid edge-cut artifacts
    if pad:
        c = Image.new('RGB', (cell.width+2*pad, cell.height+2*pad), (255,255,255))
        c.paste(cell, (pad, pad))
        cell = c
    if max(cell.size) > 2000:
        r = 2000 / max(cell.size)
        cell = cell.resize((int(cell.width*r), int(cell.height*r)), Image.LANCZOS)
    scale = 1.5
    cell_big = cell.resize((int(cell.width*scale), int(cell.height*scale)), Image.LANCZOS)
    rgba = remove(cell_big, session=session, post_process_mask=True)
    gc.collect()
    a0 = np.array(rgba)[:, :, 3].astype(np.float32) / 255.0
    rgb0 = np.array(rgba)[:, :, :3]
    a = np.clip(a0, 0.0, 1.0)[..., None]
    rgb_dc = (rgb0.astype(np.float32) - 255.0 * (1.0 - a)) / np.maximum(a, 1e-6)
    rgb_dc = np.clip(rgb_dc, 0, 255).astype(np.uint8)
    a_im = Image.fromarray((a0*255).astype(np.uint8), 'L').filter(ImageFilter.MinFilter(erode))
    a_im = a_im.filter(ImageFilter.GaussianBlur(blur))
    a2 = np.array(a_im).astype(np.float32) / 255.0
    mix = np.clip(a2 / np.maximum(a0, 1e-6), 0, 1)[..., None]
    rgb_final = (rgb_dc * mix + rgb0 * (1 - mix)).astype(np.uint8)
    out_im = Image.fromarray(np.dstack([rgb_final, (a2*255).astype(np.uint8)]), 'RGBA')
    bbox = out_im.getbbox()
    if bbox:
        xa, ya, xb, yb = bbox
        m = 3
        out_im = out_im.crop((max(0,xa-m), max(0,ya-m), min(out_im.width,xb+m), min(out_im.height,yb+m)))
    out_im.save(OUT + out)
    print(out, out_im.size, flush=True)

# (src, out, is_model, erode, blur)
jobs = [
    ('products/t-shirt-workflow-black/t-shirt-workflow-polochka-.jpg', 'workflow_front.png', 0, 11, 3.5),
    ('products/t-shirt-workflow-black/t-shirt-workflow-spinka-.jpg', 'workflow_back.png', 0, 11, 3.5),
    ('products/pants-no-color-black/pants-no-color-black.jpg', 'nocolor_front.png', 0, 11, 3.5),
    ('products/bomber-hood-weather-black/zl0a9742-1.jpg', 'bomber_front.png', 0, 11, 3.5),
    ('products/windbreaker-hood-weather/zl0a9568-1.jpg', 'windbreaker_front.png', 0, 11, 3.5),
    ('products/pants-reflector-black/20260204-dsc04789.jpg', 'reflector_front.png', 1, 9, 3.0),
    ('products/pants-reflector-black/20260204-dsc04801.jpg', 'reflector_back.png', 1, 9, 3.0),
    ('products/futbolka-zrd-cnt-3/futbolka-zrd-cnt-black-.jpg', 'cnt_front.png', 1, 9, 3.0),
    ('products/futbolka-zrd-cnt-3/futbolka-zrd-cnt-black-4.jpg', 'cnt_back.png', 1, 9, 3.0),
]

for src, out, is_model, erode, blur in jobs:
    p = os.path.join('/root/.hermes/image_cache/zrd_1x1', src)
    try:
        cut(p, out, is_model, erode=erode, blur=blur)
    except Exception as e:
        print(out, 'FAIL', e)
print('done')