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, ImageEnhance
import rembg
import onnxruntime as ort
import cv2
import numpy as np

app = FastAPI()

# Limita o número de threads do onnxruntime
options = ort.SessionOptions()
options.intra_op_num_threads = 2

def resize_image(image, max_size=1024):
    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 adjust_brightness_contrast(image, brightness=1.2, contrast=1.2):
    enhancer = ImageEnhance.Brightness(image)
    image = enhancer.enhance(brightness)
    enhancer = ImageEnhance.Contrast(image)
    image = enhancer.enhance(contrast)
    return image

def remove_shadows(image):
    # Converte a imagem para escala de cinza
    gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGBA2GRAY)
    
    # Aplica limiarização para remover áreas escuras (sombras)
    _, mask = cv2.threshold(gray, 240, 255, cv2.THRESH_BINARY)
    
    # Aplica a máscara na imagem original
    result = cv2.bitwise_and(np.array(image), np.array(image), mask=mask)
    return Image.fromarray(result)

@app.get("/remove-background")
async def remove_background(image_url: str):
    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: redimensiona e ajusta brilho/contraste
        image = resize_image(image, max_size=1024)
        image = adjust_brightness_contrast(image)
        
        # Remove o fundo da imagem usando rembg
        output = rembg.remove(image, model="u2net_human_seg", session_options=options)
        
        # Remove sombras
        output = remove_shadows(output)
        
        # Suaviza a imagem
        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)
        
        # 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))