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

from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
import requests
from io import BytesIO
from PIL import Image, ImageFilter
import rembg
import onnxruntime as ort
from concurrent.futures import ThreadPoolExecutor
import asyncio
import gc
from asyncio import Semaphore

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

# Cria uma sessão global do rembg para reutilização
session = rembg.new_session(model_name="u2net", session_options=options)

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

# Semáforo para limitar requisições simultâneas
semaphore = Semaphore(4)  # Ajuste conforme a capacidade da sua CPU/RAM

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_url):
    try:
        # Baixa a imagem da URL fornecida
        response = requests.get(image_url)
        response.raise_for_status()
        
        # Abre a imagem usando Pillow
        image = Image.open(BytesIO(response.content))
        
        # Pré-processamento: apenas redimensiona para 512px
        image = resize_image(image, max_size=512)
        
        # Remove o fundo da imagem usando rembg (reutiliza a sessão global)
        output = rembg.remove(image, session=session)
        
        # 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)
        
        # Libera memória explicitamente
        del image, output
        gc.collect()  # Força a liberação de memória
        
        return img_byte_arr
    except Exception as e:
        raise e

@app.get("/remove-background")
async def remove_background(image_url: str):
    async with semaphore:  # Limita o número de requisições simultâneas
        try:
            # Executa o processamento da imagem em um thread separado
            loop = asyncio.get_event_loop()
            img_byte_arr = await loop.run_in_executor(executor, process_image, image_url)
            
            # Retorna a imagem processada diretamente no navegador
            return StreamingResponse(img_byte_arr, media_type="image/png")
        
        except Exception as e:
            raise HTTPException(status_code=400, detail=str(e))