import base64

import requests
import os
from PIL import Image
import io

def inpaint(img_path:str,mask_path:str)->"img content (resp.content)":
    image_bytes = open(img_path, 'rb')
    mask_bytes = open(mask_path, 'rb')
    # 将字节数据转换为Base64编码的字符串

    files = {
        "image": image_bytes,
        "mask":mask_bytes
    }
    payload = {
        "ldmSteps": 25,
        "ldmSampler": "plms",
        "zitsWireframe": True,
        "hdStrategy": "Crop",
        "hdStrategyCropMargin": 196,
        "hdStrategyCropTrigerSize": 800,
        "hdStrategyResizeLimit": 2048,
        "prompt": "",
        "negativePrompt": "",
        "croperX": 307,
        "croperY": 544,
        "croperHeight": 512,
        "croperWidth": 512,
        "useCroper": False,
        "sdMaskBlur": 5,
        "sdStrength": 0.75,
        "sdSteps": 50,
        "sdGuidanceScale": 7.5,
        "sdSampler": "uni_pc",
        "sdSeed": -1,
        "sdMatchHistograms": False,
        "sdScale": 1,
        "cv2Radius": 5,
        "cv2Flag": "INPAINT_NS",
        "paintByExampleSteps": 50,
        "paintByExampleGuidanceScale": 7.5,
        "paintByExampleSeed": -1,
        "paintByExampleMaskBlur": 5,
        "paintByExampleMatchHistograms": False,
        "p2pSteps": 50,
        "p2pImageGuidanceScale": 1.5,
        "p2pGuidanceScale": 7.5,
        "controlnet_conditioning_scale": 0.4,
        "controlnet_method": "control_v11p_sd15_canny"
    }#payload用data

    #不使用header

    # resp = requests.post("https://sanster-lama-cleaner-lama.hf.space/inpaint",data= payload,files=files)##huggingface版本
    resp = requests.post("https://sanster-lama-cleaner-lama.hf.space/inpaint", data=payload, files=files)
    return resp.content

def save_img(img_content:"要处理的图片数据",new_save_path:"新文件的保存路径(包含后缀)",old_img_path:"旧文件路径")->"void生成新的文件保存 ,传入旧文件路径是为了删除有问题的旧文件":
    print(new_save_path)
    try:
        img = Image.open(io.BytesIO(img_content))
        # 如果需要指定图像格式，可以在保存时指定
        img.save(new_save_path, format="JPEG")
    except Exception as e:
        #对于可能异常的图片->比如因为不合规导致resp.content没有正常返回
        print(e,new_save_path,"图片返回有问题,跳过并删除图片.这里的路径是新保存路径")
        os.remove(old_img_path)





if __name__ == '__main__':
    # 获取当前目录的子目录的路径
    img_path = 'manga'
    subdir_path = os.path.join(os.getcwd(), img_path)

    # 图片素材获取（包含子目录下所有图片）
    image_files = []
    for root, dirs, files in os.walk(subdir_path):
        for file in files:
            if file.endswith(".jpg") or file.endswith(".png"):
                image_files.append(os.path.relpath(os.path.join(root, file)))

    # 创建处理后的子目录在与image_files同级目录下
    processed_subdir_path = os.path.join(os.path.dirname(subdir_path), f"{img_path}1")
    os.makedirs(processed_subdir_path, exist_ok=True)

    # 对image_files进行某种处理，生成新图片，并保存在处理后的子目录中
    for img_file in image_files:
        # 处理图片的代码（这里仅作示例）
        # 假设处理后的图片为new_img
        img_dir = os.path.dirname(img_file)
        new_img_dir = os.path.join(processed_subdir_path, img_dir)
        os.makedirs(new_img_dir, exist_ok=True)

        new_img_path = os.path.join(new_img_dir, os.path.basename(img_file))

        if not os.path.exists(new_img_path):
            #如果已经处理过那么跳过
            # 处理图片并保存
            img_inpainted = inpaint(img_path=img_file, mask_path='all_mask.jpg')
            save_img(img_content=img_inpainted, new_save_path=new_img_path,old_img_path=img_file)
        else:
            print(f"Skipping {new_img_path} as it already exists.")