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
    
def img_2_furry(prompt, 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))

    input_image = base64.b64encode(image_data).decode('utf-8')
    payload = {
        "init_images": [
            input_image
        ],
        "denoising_strength": 0.76,
        "resize_mode": 1,
        "width": 512,
        "height": 512,
        "prompt": prompt,
        "negative_prompt": "(worst quality, low quality:2),NSFW,monochrome,zombie,overexposure,watermark,text,bad anatomy,bad hand,extra hands,extra fingers,too many fingers,fused fingers,bad arm,distorted arm,extra arms,fused arms,extra legs,missing leg,disembodied leg,extra nipples,detached arm,liquid hand,inverted hand,disembodied limb,small breasts,loli,oversized head,extra body,completely nude,extra navel,EasyNegative,(hair between eyes),sketch,duplicate,ugly,huge eyes,text,logo,worst face,(bad and mutated hands:1.3),(blurry:2),horror,geometry,bad_prompt,(bad hands),(missing fingers),multiple limbs,bad anatomy,(interlocked fingers:1.2),Ugly Fingers,(extra digit and hands and fingers and legs and arms:1.4),((2girl)),(deformed fingers:1.2),(long fingers:1.2),(bad-artist-anime),bad-artist,bad hand,extra legs",
        "batch_size": 1,
        "steps": 20,
        "cfg_scale": 7,
        "sampler_index": "Euler a",
        "alwayson_scripts": {
            "controlnet": {
                "args": [
                    {
                        "input_image": input_image,
                        "module": "openpose_full",
                        "model": "control_v11p_sd15_openpose [cab727d4]",
                        "weight": 0.25,
                        "starting_control_step": 0,
                        "ending_control_step": 0.36,
                    },
                    {
                        "input_image": input_image,
                        "module": "canny",
                        "model": "control_v11p_sd15_canny [d14c016b]",
                        "weight": 0.25,
                        "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