import io
import cv2
import base64
import requests
from PIL import Image
import numpy as np

"""
    To use this example make sure you've done the following steps before executing:
    1. Ensure automatic1111 is running in api mode with the controlnet extension. 
       Use the following command in your terminal to activate:
            ./webui.sh --no-half --api
    2. Validate python environment meet package dependencies.
       If running in a local repo you'll likely need to pip install cv2, requests and PIL 
"""


class ControlnetRequestImg2Img:
    def __init__(self, prompt, net_prompt):
        self.url = "http://127.0.0.1:7860/sdapi/v1/img2img"
        self.prompt = prompt
        self.neg_prompt = net_prompt
        self.body = None

    def build_body(self, dst_width, dst_height, cfg_scale, base_img):

        self.body = {
            "prompt": self.prompt,
            "negative_prompt": self.neg_prompt,
            "sampler_name": "Restart",
            "batch_size": 1,
            "steps": 30,
            "width": dst_width,
            "height": dst_height,
            "cfg_scale": cfg_scale,
            "seed": -1,
            "init_images": [
                self.encode_image_to_base64(base_img)
            ],
            "denoising_strength": 0.4,
            "alwayson_scripts": {
                "controlnet": {
                    "args": [
                        {
                            "enabled": True,
                            "module": "openpose_full",
                            "model": "openpose",
                            "weight": 1.0,
                            # "image": self.read_image(),
                            "resize_mode": "Crop and Resize",
                            "low_vram": False,
                            "processor_res": 512,
                            "guidance_start": 0.0,
                            "guidance_end": 1.0,
                            "control_mode": "Balanced",
                            "pixel_perfect": True
                        }
                    ]
                }
            }
        }

    def send_request(self):
        response = requests.post(url=self.url, json=self.body)
        return response.json()

    def encode_image_to_base64(self, img):
        retval, bytes = cv2.imencode('.png', img)
        encoded_image = base64.b64encode(bytes).decode('utf-8')
        return encoded_image

    def read_image(self):
        img = cv2.imread(self.img_path)
        retval, bytes = cv2.imencode('.png', img)
        encoded_image = base64.b64encode(bytes).decode('utf-8')
        return encoded_image
    
    def read_mask(self):
        img = cv2.imread(self.mask)
        retval, bytes = cv2.imencode('.png', img)
        encoded_image = base64.b64encode(bytes).decode('utf-8')
        return encoded_image

def encode_image_to_base64(img):
    retval, bytes = cv2.imencode('.jpg', img)
    encoded_image = base64.b64encode(bytes).decode('utf-8')
    return encoded_image



if __name__ == '__main__':
    path = '/home/chinatszrn/Downloads/photo_service/service_data/template_data/template01.png'
    img = cv2.imread(path)
    prompt = '<lora:5b05d5eeee0188f436d7131c4f0ff52b:0.8>,easyphoto_face, easyphoto, 1person,face,suit'
    neg_prompt = '(worst quality:2),(low quality:2),(normal quality:2),lowres,watermark'


    control_net = ControlnetRequestImg2Img(prompt, neg_prompt)
    control_net.build_body(dst_width=img.shape[1], dst_height=img.shape[0], cfg_scale=3.5, base_img=img)
    output = control_net.send_request()
    result = output['images'][0]

    image_array = np.frombuffer(base64.b64decode(result.split(",", 1)[0]), np.uint8)
    image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
    cv2.imshow('image', image)
    cv2.waitKey()