import os
import subprocess
import torch
import requests
from PIL import Image
from io import BytesIO

just_get_sd_mask_function = None

print(os.listdir('/usr/local/'))
print(torch.version.cuda)

class EndpointHandler():
    def __init__(self, path="."):
        global just_get_sd_mask_function
        
        is_production = True
        
        if False:
            return

        os.chdir(path)
        
        os.environ['AM_I_DOCKER'] = 'False'
        os.environ['BUILD_WITH_CUDA'] = 'True'
        os.environ['CUDA_HOME'] = '/usr/local/cuda-11.7/' if is_production else '/usr/local/cuda-12.1/'

        # Install Segment Anything
        subprocess.run(["python", "-m", "pip", "install", "-e", "segment_anything"])

        # Install Grounding DINO
        subprocess.run(["python", "-m", "pip", "install", "-e", "GroundingDINO"])
    
        subprocess.run("wget https://huggingface.co/Uminosachi/sam-hq/resolve/main/sam_hq_vit_h.pth -O ./sam_hq_vit_h.pth", shell=True)

        # Install diffusers
        subprocess.run(["pip", "install", "--upgrade", "diffusers[torch]"])

        # Install osx
        subprocess.run(["git", "submodule", "update", "--init", "--recursive"])
        subprocess.run(["bash", "grounded-sam-osx/install.sh"], cwd="grounded-sam-osx")

        # Install RAM & Tag2Text
        subprocess.run(["git", "clone", "https://github.com/xinyu1205/recognize-anything.git"])
        subprocess.run(["pip", "install", "-r", "./recognize-anything/requirements.txt"])
        subprocess.run(["pip", "install", "-e", "./recognize-anything/"])

        from test import just_get_sd_mask
        just_get_sd_mask_function = just_get_sd_mask

    def __call__(self, data):
        mask_pil = just_get_sd_mask(Image.open("assets/demo1.jpg"), "bear", 10)

        if mask_pil.mode != 'RGB':
            mask_pil = mask_pil.convert('RGB')

        # Convert PIL image to byte array
        img_byte_arr = BytesIO()
        mask_pil.save(img_byte_arr, format='JPEG')
        img_byte_arr = img_byte_arr.getvalue()

        # Upload to file.io
        response = requests.post("https://file.io/", files={"file": img_byte_arr})
        url = response.json().get('link')

        return {"url": url}