import os
import sys
from PIL import Image
from io import BytesIO

# Desabilita JIT e cache do Numba completamente
os.environ["NUMBA_DISABLE_JIT"] = "1"
os.environ["NUMBA_CACHE_DIR"] = "/tmp/numba_cache"
os.makedirs("/tmp/numba_cache", exist_ok=True)

# Pasta temporária para modelos do rembg
cache_dir = "/tmp/u2net"
os.makedirs(cache_dir, exist_ok=True)
os.environ["POOCH"] = cache_dir
os.environ["XDG_CACHE_HOME"] = "/tmp/cache"

from rembg import remove, new_session

# Cria sessão com modelo leve e rápido
session = new_session("u2netp")

input_path = sys.argv[1]
output_path = sys.argv[2]

# Abre a imagem com PIL
with Image.open(input_path) as img:
    # Calcula altura proporcional para largura de 300px
    w_percent = 300 / float(img.width)
    h_size = int(float(img.height) * w_percent)
    img_resized = img.resize((300, h_size), Image.Resampling.LANCZOS)

    # Converte para bytes para enviar ao rembg
    img_bytes = BytesIO()
    img_resized.save(img_bytes, format="PNG")
    img_bytes.seek(0)

    # Remove o fundo
    result = remove(img_bytes.read(), session=session, post_process_mask=False)

    # Salva resultado
    with open(output_path, "wb") as o:
        o.write(result)