#!/usr/bin/env python3
"""Batch cutout for all 9 products: flat items / model shots -> PNG with smooth edges."""
import numpy as np
from PIL import Image, ImageFilter
from rembg import remove, new_session
import gc, os

session = new_session('isnet-general-use')
OUT = '/root/.hermes/image_cache/zrd_1x1/cutouts/'
os.makedirs(OUT, exist_ok=True)

# slug -> [(output_name, source_path, is_model)]
jobs = {
 'bomber-hood-weather-black': [
   ('bomber_front.png', 'products/bomber-hood-weather-black/zl0a9742-1.jpg', False),
 ],
 'windbreaker-hood-weather': [
   ('windbreaker_front.png', 'products/windbreaker-hood-weather/zl0a9568-1.jpg', False),
 ],
 'pants-reflector-black': [
   ('reflector_front.png', 'products/pants-reflector-black/20260204-dsc04789.jpg', True),
   ('reflector_back.png', 'products/pants-reflector-black/20260204-dsc04801.jpg', True),
 ],
 't-shirt-workflow-black': [
   ('workflow_front.png', 'products/t-shirt-workflow-black/t-shirt-workflow-polochka-.jpg', False),
   ('workflow_back.png', 'products/t-shirt-workflow-black/t-shirt-workflow-spinka-.jpg', False),
 ],
 't-shirt-belyj-russkij-white': [
   ('belyj_front.png', 'products/t-shirt-belyj-russkij-white/futbolka-belaya-razlozhka-1.jpg', False),
   ('belyj_back.png', 'products/t-shirt-belyj-russkij-white/futbolka-belaya-razlozhka-2.jpg', False),
 ],
 'futbolka-zrd-cnt-3': [
   ('cnt_front.png', 'products/futbolka-zrd-cnt-3/futbolka-zrd-cnt-black-.jpg', True),
   ('cnt_back.png', 'products/futbolka-zrd-cnt-3/futbolka-zrd-cnt-black-4.jpg', True),
 ],
 'pants-no-color-black': [
   ('nocolor_front.png', 'products/pants-no-color-black/pants-no-color-black.jpg', False),
 ],
 'jacket-full-grey': [
   ('fullgrey_front.png', 'products/jacket-full-grey/zl0a9995.jpg', False),
 ],
 'pants-nylon-basic-grey': [
   ('nylon_front.png', 'products/pants-nylon-basic-grey/zl0a9976.jpg', False),
 ],
}

def cut(src, out, is_model):
    cell = Image.open(src).convert('RGB')
    # model shots are big; resize to max 2200px
    if max(cell.size) > 2200:
        r = 2200 / max(cell.size)
        cell = cell.resize((int(cell.width*r), int(cell.height*r)), Image.LANCZOS)
    scale = 2.0
    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)
    if is_model:
        # model shots: less erosion, more feather
        filt = ImageFilter.MinFilter(5)
        blur = 2.2
    else:
        filt = ImageFilter.MinFilter(9)
        blur = 3.0
    a_im = Image.fromarray((a0*255).astype(np.uint8), 'L').filter(filt)
    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)

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