import os
os.environ["NUMBA_CACHE_DIR"] = "/tmp/numba_cache"
os.environ["U2NET_HOME"] = "/tmp/u2net"

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from io import BytesIO
from PIL import Image, ImageFilter
import rembg
import onnxruntime as ort
from concurrent.futures import ThreadPoolExecutor
import asyncio
import gc
import base64

app = FastAPI()

# Configurações do onnxruntime para CPU
options = ort.SessionOptions()
options.intra_op_num_threads = 2  # Limita o número de threads para evitar sobrecarga
options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL  # Execução sequencial para melhor desempenho em CPU

# Pool de threads para executar tarefas bloqueantes
executor = ThreadPoolExecutor(max_workers=4)

def resize_image(image, max_size=512):
    """Redimensiona a imagem para uma largura máxima de 512px, mantendo a proporção."""
    width, height = image.size
    if width > max_size or height > max_size:
        ratio = min(max_size / width, max_size / height)
        new_size = (int(width * ratio), int(height * ratio))
        image = image.resize(new_size, Image.Resampling.LANCZOS)
    return image

def process_image(image_data):
    # Decodifica a imagem base64
    image = Image.open(BytesIO(base64.b64decode(image_data)))
    
    # Pré-processamento: apenas redimensiona para 512px
    image = resize_image(image, max_size=512)
    
    # Remove o fundo da imagem usando rembg
    output = rembg.remove(image, session_options=options)
    
    # Pós-processamento: suaviza as bordas
    output = output.filter(ImageFilter.SMOOTH_MORE)
    
    # Converte a imagem de volta para bytes
    img_byte_arr = BytesIO()
    output.save(img_byte_arr, format='PNG')
    img_byte_arr.seek(0)
    
    # Codifica a imagem processada em base64
    return base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')

class ImageBase64(BaseModel):
    image: str  # A imagem em base64

@app.post("/remove-background")
async def remove_background(image_data: ImageBase64):
    try:
        # Executa o processamento da imagem em um thread separado
        loop = asyncio.get_event_loop()
        processed_image_base64 = await loop.run_in_executor(executor, process_image, image_data.image)
        
        # Retorna a imagem processada em base64
        return {"processed_image": processed_image_base64}
    
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))
    
    finally:
        # Força a coleta de lixo após cada requisição
        gc.collect()