# #!/usr/bin/env python
# # -*- coding: UTF-8 -*-
# """
# @Project ：trinity_client
# @File    ：service_att_recognition.py
# @Author  ：周成融
# @Date    ：2023/7/26 12:01:05
# @detail  ：
# """
# import json
# import logging
# import time
#
# import cv2
# import numpy as np
# import redis
# import tritonclient.grpc as grpcclient
# from PIL import Image
# from tritonclient.utils import np_to_triton_dtype
#
# from app.core.config import *
# from app.schemas.generate_image import GenerateProductImageModel
# from app.service.generate_image.utils.upload_sd_image import upload_SDXL_image
# from app.service.utils.oss_client import oss_get_image
#
# logger = logging.getLogger()
#
#
# class GenerateProductImage:
#     def __init__(self, request_data):
#         if DEBUG is False:
#             self.connection = pika.BlockingConnection(pika.ConnectionParameters(**RABBITMQ_PARAMS))
#             self.channel = self.connection.channel()
#         # self.connection = pika.BlockingConnection(pika.ConnectionParameters(**RABBITMQ_PARAMS))
#         # self.channel = self.connection.channel()
#         # self.minio_client = Minio(MINIO_URL, access_key=MINIO_ACCESS, secret_key=MINIO_SECRET, secure=MINIO_SECURE)
#         self.grpc_client = grpcclient.InferenceServerClient(url=GPI_MODEL_URL)
#         self.redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True)
#         self.category = "product_image"
#         self.image_strength = request_data.image_strength
#         self.batch_size = 1
#         self.product_type = request_data.product_type
#         self.prompt = request_data.prompt
#         self.image, self.image_size, self.left, self.top = pre_processing_image(request_data.image_url)
#         self.tasks_id = request_data.tasks_id
#         self.user_id = self.tasks_id[self.tasks_id.rfind('-') + 1:]
#         self.gen_product_data = {'tasks_id': self.tasks_id, 'status': 'PENDING', 'message': "pending", 'image_url': ''}
#         self.redis_client.set(self.tasks_id, json.dumps(self.gen_product_data))
#         self.redis_client.expire(self.tasks_id, 600)
#
#     def callback(self, result, error):
#         if error:
#             self.gen_product_data['status'] = "FAILURE"
#             self.gen_product_data['message'] = str(error)
#             self.redis_client.set(self.tasks_id, json.dumps(self.gen_product_data))
#         else:
#             # pil图像转成numpy数组
#             image = result.as_numpy("generated_inpaint_image")
#             image_result = Image.fromarray(np.squeeze(image.astype(np.uint8))).resize(self.image_size)
#             cropped_image = post_processing_image(image_result, self.left, self.top)
#             image_url = upload_SDXL_image(cropped_image, user_id=self.user_id, category=f"{self.category}", file_name=f"{self.tasks_id}.png")
#             self.gen_product_data['status'] = "SUCCESS"
#             self.gen_product_data['message'] = "success"
#             self.gen_product_data['image_url'] = str(image_url)
#             self.redis_client.set(self.tasks_id, json.dumps(self.gen_product_data))
#
#     def read_tasks_status(self):
#         status_data = self.redis_client.get(self.tasks_id)
#         return json.loads(status_data), status_data
#
#     def get_result(self):
#         try:
#             prompts = [self.prompt] * self.batch_size
#             self.image = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB)
#             self.image = cv2.resize(self.image, (1024, 1024))
#             images = [self.image.astype(np.uint8)] * self.batch_size
#
#             if self.product_type == "single":
#                 text_obj = np.array(prompts, dtype="object").reshape(-1, 1)
#                 image_obj = np.array(images, dtype=np.uint8).reshape((-1, 1024, 1024, 3))
#                 image_strength_obj = np.array(self.image_strength, dtype=np.float32).reshape(-1, 1)
#             else:
#                 text_obj = np.array(prompts, dtype="object").reshape((1))
#                 image_obj = np.array(images, dtype=np.uint8).reshape((1024, 1024, 3))
#                 image_strength_obj = np.array(self.image_strength, dtype=np.float32).reshape((1))
#
#             # 假设 prompts、images 和 self.image_strength 已经定义
#
#             input_text = grpcclient.InferInput("prompt", text_obj.shape, np_to_triton_dtype(text_obj.dtype))
#             input_image = grpcclient.InferInput("input_image", image_obj.shape, "UINT8")
#             input_image_strength = grpcclient.InferInput("image_strength", image_strength_obj.shape, np_to_triton_dtype(image_strength_obj.dtype))
#
#             input_text.set_data_from_numpy(text_obj)
#             input_image.set_data_from_numpy(image_obj)
#             input_image_strength.set_data_from_numpy(image_strength_obj)
#
#             inputs = [input_text, input_image, input_image_strength]
#
#             if self.product_type == "single":
#                 ctx = self.grpc_client.async_infer(model_name="stable_diffusion_xl_cnet_inpaint", inputs=inputs, callback=self.callback)
#             else:
#                 ctx = self.grpc_client.async_infer(model_name=GPI_MODEL_NAME_OVERALL, inputs=inputs, callback=self.callback)
#
#             time_out = 600
#             while time_out > 0:
#                 gen_product_data, _ = self.read_tasks_status()
#                 if gen_product_data['status'] in ["REVOKED", "FAILURE"]:
#                     ctx.cancel()
#                     break
#                 elif gen_product_data['status'] == "SUCCESS":
#                     break
#                 time_out -= 1
#                 time.sleep(0.1)
#             gen_product_data, _ = self.read_tasks_status()
#             return gen_product_data
#         except Exception as e:
#             self.gen_product_data['status'] = "FAILURE"
#             self.gen_product_data['message'] = str(e)
#             self.redis_client.set(self.tasks_id, json.dumps(self.gen_product_data))
#             raise Exception(str(e))
#         finally:
#             dict_gen_product_data, str_gen_product_data = self.read_tasks_status()
#             if DEBUG is False:
#                 self.channel.basic_publish(exchange='', routing_key=GPI_RABBITMQ_QUEUES, body=str_gen_product_data)
#             logger.info(f" [x] Sent to： {GPI_RABBITMQ_QUEUES} data：@@@@  {json.dumps(dict_gen_product_data, indent=4)}")
#
#
# def infer_cancel(tasks_id):
#     redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True)
#     data = {'tasks_id': tasks_id, 'status': 'REVOKED', 'message': "revoked", 'data': 'revoked'}
#     gen_product_data = json.dumps(data)
#     redis_client.set(tasks_id, gen_product_data)
#     return data
#
#
# def pre_processing_image(image_url):
#     image = oss_get_image(bucket=image_url.split('/')[0], object_name=image_url[image_url.find('/') + 1:], data_type="PIL")
#     # resize 原图至1024*1024
#     image = image.resize((int(1024 / image.height * image.width), 1024))
#
#     # 原始图片的尺寸
#     width, height = image.size
#
#     new_height, new_width = 1024, 1024
#     # 创建一个新的画布，大小为添加 padding 后的尺寸，并设置为白色背景
#     pad_image = Image.new('RGBA', (new_width, new_height), (0, 0, 0, 0))
#
#     # 将原始图片粘贴到新的画布中心
#     left = (new_width - width) // 2
#     top = (new_height - height) // 2
#     pad_image.paste(image, (left, top))
#
#     # 将画布 resize 成宽度 1024，长度 1024
#     resized_image = pad_image.resize((1024, 1024))
#     image_size = (1024, 1024)
#
#     if resized_image.mode in ('RGBA', 'LA') or (resized_image.mode == 'P' and 'transparency' in resized_image.info):
#         # 创建白色背景
#         background = Image.new("RGB", image_size, (255, 255, 255))
#         # 将图片粘贴到白色背景上
#         background.paste(resized_image, mask=resized_image.split()[3])
#         image = np.array(background)
#         image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
#     return image, image_size, left, top
#
#
# def post_processing_image(image, left, top):
#     resized_image = image.resize((int(image.width * (768 / image.height)), 768))
#     # 计算裁剪的坐标
#     left = (resized_image.width - 512) // 2
#     upper = 0
#     right = left + 512
#     lower = 768
#
#     # 进行裁剪
#     cropped_image = resized_image.crop((left, upper, right, lower))
#     return cropped_image
#
#
# if __name__ == '__main__':
#     rd = GenerateProductImageModel(
#         tasks_id="123-89",
#         # prompt="",
#         image_strength=0.7,
#         prompt="The best quality, masterpiece,outwear, 8K realistic, HUD",
#         image_url="aida-results/result_53381ada-ac64-11ef-ae9d-0242ac150002.png",
#         product_type="overall"
#     )
#     server = GenerateProductImage(rd)
#     print(server.get_result())

# 旧版product
# !/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：trinity_client
@File    ：service_att_recognition.py
@Author  ：周成融
@Date    ：2023/7/26 12:01:05
@detail  ：
"""
import json
import logging
import time

import cv2
import numpy as np
import redis
import tritonclient.grpc as grpcclient
from PIL import Image
from tritonclient.utils import np_to_triton_dtype

from app.core.config import *
from app.schemas.generate_image import GenerateProductImageModel
from app.service.generate_image.utils.mq import publish_status
from app.service.generate_image.utils.upload_sd_image import upload_SDXL_image
from app.service.utils.oss_client import oss_get_image

logger = logging.getLogger()


class GenerateProductImage:
    def __init__(self, request_data):
        self.grpc_client = grpcclient.InferenceServerClient(url=GPI_MODEL_URL)
        self.redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True)
        self.category = "product_image"
        self.image_strength = request_data.image_strength
        self.batch_size = 1
        self.product_type = request_data.product_type
        self.prompt = request_data.prompt
        self.image = pre_processing_image(request_data.image_url)
        self.tasks_id = request_data.tasks_id
        self.user_id = self.tasks_id[self.tasks_id.rfind('-') + 1:]
        self.gen_product_data = {'tasks_id': self.tasks_id, 'status': 'PENDING', 'message': "pending", 'image_url': ''}
        self.redis_client.set(self.tasks_id, json.dumps(self.gen_product_data))
        self.redis_client.expire(self.tasks_id, 600)

    def callback(self, result, error):
        if error:
            self.gen_product_data['status'] = "FAILURE"
            if 'mask_list' in str(error):
                self.gen_product_data['status'] = "NO_FACE"
            self.gen_product_data['message'] = str(error)
            self.redis_client.set(self.tasks_id, json.dumps(self.gen_product_data))
        else:
            # pil图像转成numpy数组
            if self.product_type == "single":
                image = result.as_numpy("generated_cnet_image")
            else:
                image = result.as_numpy("generated_inpaint_image")
            image_result = Image.fromarray(np.squeeze(image.astype(np.uint8)))
            # cropped_image = post_processing_image(image_result, self.left, self.top)
            image_url = upload_SDXL_image(image_result, user_id=self.user_id, category=f"{self.category}", file_name=f"{self.tasks_id}.png")
            self.gen_product_data['status'] = "SUCCESS"
            self.gen_product_data['message'] = "success"
            self.gen_product_data['image_url'] = str(image_url)
            self.redis_client.set(self.tasks_id, json.dumps(self.gen_product_data))

    def read_tasks_status(self):
        status_data = self.redis_client.get(self.tasks_id)
        return json.loads(status_data), status_data

    def get_result(self):
        try:
            prompts = [self.prompt] * self.batch_size

            self.image = cv2.cvtColor(self.image, cv2.COLOR_RGBA2RGB)
            # self.image = cv2.resize(self.image, (1024, 1024))
            images = [self.image.astype(np.uint8)] * self.batch_size

            if self.product_type == "single":
                text_obj = np.array(prompts, dtype="object").reshape((-1, 1))
                image_obj = np.array(images, dtype=np.uint8).reshape((-1, 768, 512, 3))
                image_strength_obj = np.array(self.image_strength, dtype=np.float32).reshape((-1, 1))
            else:
                text_obj = np.array(prompts, dtype="object").reshape((1))
                image_obj = np.array(images, dtype=np.uint8).reshape((768, 512, 3))
                image_strength_obj = np.array(self.image_strength, dtype=np.float32).reshape((1))

            # 假设 prompts、images 和 self.image_strength 已经定义

            input_text = grpcclient.InferInput("prompt", text_obj.shape, np_to_triton_dtype(text_obj.dtype))
            input_image = grpcclient.InferInput("input_image", image_obj.shape, "UINT8")
            input_image_strength = grpcclient.InferInput("image_strength", image_strength_obj.shape, np_to_triton_dtype(image_strength_obj.dtype))

            input_text.set_data_from_numpy(text_obj)
            input_image.set_data_from_numpy(image_obj)
            input_image_strength.set_data_from_numpy(image_strength_obj)

            inputs = [input_text, input_image, input_image_strength]

            if self.product_type == "single":
                ctx = self.grpc_client.async_infer(model_name=GPI_MODEL_NAME_SINGLE, inputs=inputs, callback=self.callback, priority=1)
            else:
                ctx = self.grpc_client.async_infer(model_name=GPI_MODEL_NAME_OVERALL, inputs=inputs, callback=self.callback, priority=1)

            time_out = 600
            while time_out > 0:
                gen_product_data, _ = self.read_tasks_status()
                if gen_product_data['status'] in ["REVOKED", "FAILURE", "NO_FACE"]:
                    ctx.cancel()
                    break
                elif gen_product_data['status'] == "SUCCESS":
                    break
                time_out -= 1
                time.sleep(0.1)
            gen_product_data, _ = self.read_tasks_status()
            return gen_product_data
        except Exception as e:
            self.gen_product_data['status'] = "FAILURE"
            self.gen_product_data['message'] = str(e)
            self.redis_client.set(self.tasks_id, json.dumps(self.gen_product_data))
            raise Exception(str(e))
        finally:
            dict_gen_product_data, str_gen_product_data = self.read_tasks_status()
            if not DEBUG:
                publish_status(str_gen_product_data, GPI_RABBITMQ_QUEUES)


def infer_cancel(tasks_id):
    redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB, decode_responses=True)
    data = {'tasks_id': tasks_id, 'status': 'REVOKED', 'message': "revoked", 'data': 'revoked'}
    gen_product_data = json.dumps(data)
    redis_client.set(tasks_id, gen_product_data)
    return data


def pre_processing_image(image_url):
    image = oss_get_image(bucket=image_url.split('/')[0], object_name=image_url[image_url.find('/') + 1:], data_type="PIL")
    # 目标图片的尺寸
    target_width = 512
    target_height = 768

    # 原始图片的尺寸
    original_width, original_height = image.size

    # 计算宽度和高度的缩放比例
    width_ratio = target_width / original_width
    height_ratio = target_height / original_height

    # 选择较小的缩放比例，确保图片能完整放入目标图片中
    scale_ratio = min(width_ratio, height_ratio)

    # 计算调整后的尺寸
    new_width = int(original_width * scale_ratio)
    new_height = int(original_height * scale_ratio)

    # 调整图片大小
    resized_image = image.resize((new_width, new_height))

    # 创建一个 512x768 的透明图片
    result_image = Image.new("RGBA", (target_width, target_height), (255, 255, 255, 0))

    # 计算需要粘贴的位置，使图片居中
    x_offset = (target_width - new_width) // 2
    y_offset = (target_height - new_height) // 2

    # 将调整大小后的图片粘贴到透明图片上
    if resized_image.mode == "RGBA":
        result_image.paste(resized_image, (x_offset, y_offset), mask=resized_image.split()[3])
    else:
        result_image.paste(resized_image, (x_offset, y_offset))

    image = np.array(result_image)

    # image = cv2.cvtColor(image, cv2.COLOR_BGRA2RGBA)
    return image


def post_processing_image(image, left, top):
    resized_image = image.resize((int(image.width * (768 / image.height)), 768))
    # 计算裁剪的坐标
    left = (resized_image.width - 512) // 2
    upper = 0
    right = left + 512
    lower = 768

    # 进行裁剪
    cropped_image = resized_image.crop((left, upper, right, lower))
    return cropped_image


if __name__ == '__main__':
    rd = GenerateProductImageModel(
        tasks_id="123-89",
        # prompt="",
        image_strength=0.7,
        prompt=" The best quality, masterpiece, real image.Outwear,high quality clothing details,8K realistic,HDR",
        image_url="aida-results/result_40b1a2fe-e220-11ef-9bfa-0242ac150003.png",
        product_type="single"
    )
    server = GenerateProductImage(rd)
    print(server.get_result())
