import requests
import base64
from PIL import Image
from io import BytesIO
import argparse
import os

def base64_save(base64_image_data, save_path):
    # 解码Base64数据并保存为PNG文件
    img_data = base64.b64decode(base64_image_data)
    img = Image.open(BytesIO(img_data))
    # 保存为本地PNG文件
    img.save(save_path, "PNG")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="HivisionIDPhotos证件照制作推理程序。")

    parser.add_argument("-u", "--url", help="API服务的URL", default="http://127.0.0.1:8080")
    parser.add_argument("-t", "--type", help="请求API的种类，有idphoto、add_background和generate_layout_photos可选",
                        default="idphoto")
    parser.add_argument("-i", "--input_image_dir", help="输入图像路径", required=True)
    parser.add_argument("-o", "--output_image_dir", help="保存图像路径", required=True)
    parser.add_argument("-s", "--size", help="证件照尺寸", default="(413,295)")
    parser.add_argument("-c", "--color", help="证件照背景色", default="(255,255,255)")

    args = parser.parse_args()

    url = f"{args.url}/{args.type}"  # 替换为实际的接口URL
    files = {'input_image': (open(args.input_image_dir, 'rb'))}  # 替换为实际的文件路径和文件名
    data = {"size": args.size, "color": args.color}

    response = requests.post(url, data=data, files=files)

    if response.status_code == 200:
        # 获取Base64编码的图像数据
        if args.type == "idphoto":
            response_json = response.json()

            status = response_json["status"]
            if status:

                base64_image_data_standard = response_json["img_output_standard"]
                base64_image_data_standard_hd = response_json["img_output_standard_hd"]

                # 解码Base64数据并保存为PNG文件
                base64_save(base64_image_data_standard, args.output_image_dir)

                file_name, file_extension = os.path.splitext(args.output_image_dir)
                # 定义新的文件路径（在原有的文件名后添加"_hd"）
                new_file_name = file_name + "_hd" + file_extension

                base64_save(base64_image_data_standard_hd, new_file_name)


                print(f"标准照保存至'{args.output_image_dir}'，高清照保存至'{new_file_name}'")

            else:
                print('人脸数量不等于1，请上传单张人脸的图像。')

        elif args.type == "add_background":
            response_json = response.json()

            status = response_json["status"]

            if status:
                base64_image_data = response_json["image"]
                base64_save(base64_image_data, args.output_image_dir)
                print(f"增加背景后的照片保存至'{args.output_image_dir}'。")
            else:
                print(f'遇到了一些问题，报错为{response_json["error"]}')

        elif args.type == "generate_layout_photos":
            response_json = response.json()

            status = response_json["status"]

            if status:
                base64_image_data = response_json["image"]
                base64_save(base64_image_data, args.output_image_dir)
                print(f"排版照保存至'{args.output_image_dir}'。")
            else:
                print(f'遇到了一些问题，报错为{response_json["error"]}')

    else:
        print("请求失败")