import json
import os
import requests
import io
import base64
from PIL import Image, PngImagePlugin
from dotenv import load_dotenv
from utils.logger import logger

load_dotenv() # 加载环境变量
img_url = os.getenv('img_url')

def img_2_img(prompt, mask_img, img_key, img_path, output_directory):
    with open(img_path, 'rb') as f:
        image_data = f.read()
    input_image_pil = Image.open(io.BytesIO(image_data))
    width, height = input_image_pil.size

    input_image = base64.b64encode(image_data).decode('utf-8')
    mask = base64.b64encode(mask_img).decode('utf-8')
    payload = {
        "init_images": [
            input_image
        ],
        "denoising_strength": 0.96,
        "width": width,
        "height": height,
        "mask": mask,
        "mask_blur": 8,
        "prompt": prompt,
        "negative_prompt": "EasyNegativeV2,(badhandv4:1.2)",
        "batch_size": 1,
        "inpainting_mask_invert": 1,
        "steps": 30,
        "cfg_scale": 7,
        "sampler_index": "DPM++ 2M Karras",
        "alwayson_scripts": {
            "ADetailer": {
                "args": [
                    True,
                    {
                        "ad_model": "face_yolov8n.pt",
                        "ad_prompt": prompt,
                        "ad_negative_prompt": "EasyNegativeV2,(badhandv4:1.2)",
                    },
                    {
                        "ad_model": "hand_yolov8n.pt",
                        "ad_prompt": prompt,
                        "ad_negative_prompt": "EasyNegativeV2,(badhandv4:1.2)",
                    }
                ]
            },
            "controlnet": {
                "args": [
                    {
                        "input_image": input_image,
                        "module": "openpose_full",
                        "model": "control_v11p_sd15_openpose [cab727d4]",
                    },
                    {
                        "input_image": input_image,
                        "module": "canny",
                        "model": "control_v11p_sd15_canny [d14c016b]",
                        "weight": 0.36,
                        "starting_control_step": 0,
                        "ending_control_step": 0.36,
                    }
                ]
            }
        }
    }
    response = requests.post(url=f'{img_url}/sdapi/v1/img2img', json=payload)

    r = response.json()
    output_path = output_directory + f"/{img_key}-output.jpg"
    try:
        result = r['images'][0]
        image = Image.open(io.BytesIO(base64.b64decode(result.split(",", 1)[0])))
        image.save(output_path)
        return output_path
    except:
        logger.error('img2img error' + str(r))
        return None