from PIL import Image
from rembg import remove
from io import BytesIO
import base64
import uuid
import os

class RemoveBG:

    def removebg(path):
        original_image = Image.open(path)
        no_bg_img = remove(original_image)
        return no_bg_img
        

    def init(self, data):
        if isinstance(data['data'], list):
            result = []
            
            files = data['data']
            for file in files:
                image_bytes = file.split(',')[1].encode('utf-8')
                image = Image.open(BytesIO(base64.b64decode(image_bytes)))
                random_filename = str(uuid.uuid4()) + '.png'
                image.save('./tmp/' + random_filename)
                image_path = './tmp/' + random_filename
                
                no_bg_image = RemoveBG.removebg(image_path)
                
                name_nobg = str(uuid.uuid4())
                
                img_nobg_path = f"./tmp_nobg/{name_nobg}.png"
                
                no_bg_image.save(img_nobg_path, format='PNG')
    
                os.remove(image_path)
                
                url = f"/tmp_nobg/{name_nobg}.png"
                result.append(url)
            return result
        else:
            image_bytes = data['data'].split(',')[1].encode('utf-8')
            image = Image.open(BytesIO(base64.b64decode(image_bytes)))
            random_filename = str(uuid.uuid4()) + '.png'
            image.save('./tmp/' + random_filename)
    
            image_path = './tmp/' + random_filename
    
           
    
            # Remova o fundo e obtenha a imagem sem fundo
            no_bg_image = RemoveBG.removebg(image_path)
    
            name_nobg = str(uuid.uuid4())
    
            img_nobg_path = f"./tmp_nobg/{name_nobg}.png"
            # Converta a imagem sem fundo em bytes
            no_bg_image.save(img_nobg_path, format='PNG')
    
            os.remove(image_path)
    
            # Agora 'no_bg_image_bytes' contém os dados da imagem sem fundo em formato PNG
    
            # Se desejar, você pode retornar 'no_bg_image_bytes' ou usá-lo conforme necessário.
            return f"/tmp_nobg/{name_nobg}.png"