# 旧版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 cv2
import numpy as np
import tritonclient.grpc as grpcclient
from PIL import Image
from celery import Celery
from tritonclient.utils import np_to_triton_dtype

from app.core.config import *
from app.schemas.generate_image import BatchGenerateProductImageModel, ProductItemModel
from app.service.generate_image.utils.upload_sd_image import upload_SDXL_image
from app.service.utils.oss_client import oss_get_image

celery_app = Celery('product_tasks', broker=f'amqp://rabbit:123456@18.167.251.121:5672//', backend='rpc://', BROKER_CONNECTION_RETRY_ON_STARTUP=True)
celery_app.conf.task_default_queue = 'queue_product'
celery_app.conf.worker_log_format = '%(asctime)s %(filename)s [line:%(lineno)d] %(levelname)s %(message)s'
celery_app.conf.worker_hijack_root_logger = False
logger = logging.getLogger()
logging.getLogger('pika').setLevel(logging.WARNING)
grpc_client = grpcclient.InferenceServerClient(url=GPI_MODEL_URL)
category = "product_image"


@celery_app.task
def batch_generate_product(batch_request_data):
    batch_size = len(batch_request_data['batch_data_list'])
    logger.info(f"batch_generate_product batch_request_data:{json.dumps(batch_request_data, indent=4)}")
    batch_tasks_id = batch_request_data['batch_tasks_id']
    user_id = batch_request_data['user_id']
    result_data_list = []

    for i, data in enumerate(batch_request_data['batch_data_list']):
        tasks_id = data['tasks_id']
        image = pre_processing_image(data['image_url'])
        image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
        images = [image.astype(np.uint8)] * 1
        prompts = [data['prompt']] * 1
        if data['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(data['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(data['image_strength'], dtype=np.float32).reshape((1))
        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]

        try:
            if data['product_type'] == "single":
                result = grpc_client.infer(model_name=GPI_MODEL_NAME_SINGLE, inputs=inputs, priority=100)
                image = result.as_numpy("generated_cnet_image")
            else:
                result = grpc_client.infer(model_name=GPI_MODEL_NAME_OVERALL, inputs=inputs, priority=100)
                image = result.as_numpy("generated_inpaint_image")
            image_result = Image.fromarray(np.squeeze(image.astype(np.uint8)))
        except Exception as e:
            if 'mask_list' in str(e):
                e_text_obj = np.array(prompts, dtype="object").reshape((-1, 1))
                e_image_obj = np.array(images, dtype=np.uint8).reshape((-1, 768, 512, 3))
                e_image_strength_obj = np.array(data['image_strength'], dtype=np.float32).reshape((-1, 1))

                e_input_text = grpcclient.InferInput("prompt", e_text_obj.shape, np_to_triton_dtype(e_text_obj.dtype))
                e_input_image = grpcclient.InferInput("input_image", e_image_obj.shape, "UINT8")
                e_input_image_strength = grpcclient.InferInput("image_strength", e_image_strength_obj.shape, np_to_triton_dtype(e_image_strength_obj.dtype))

                e_input_text.set_data_from_numpy(e_text_obj)
                e_input_image.set_data_from_numpy(e_image_obj)
                e_input_image_strength.set_data_from_numpy(e_image_strength_obj)

                result = grpc_client.infer(model_name=GPI_MODEL_NAME_SINGLE, inputs=[e_input_text, e_input_image, e_input_image_strength], priority=100)
                image = result.as_numpy("generated_cnet_image")
                image_result = Image.fromarray(np.squeeze(image.astype(np.uint8)))
            else:
                image_result = str(e)
                logger.error(image_result)

        if isinstance(image_result, Image.Image):
            image_url = upload_SDXL_image(image_result, user_id=user_id, category=f"{category}", file_name=f"{tasks_id}-batch-{i}.png")
            data['product_img'] = image_url
            result_data_list.append(data)
        else:
            image_url = image_result
            data['product_img'] = image_url
            result_data_list.append(data)

        # 发送每条结果
        if DEBUG:
            logger.info(f" [x]Queue : {BATCH_GPI_RABBITMQ_QUEUES} | tasks_id：{tasks_id} | progress：{i + 1}/{batch_size} | result_data：{data}")
            print(f" [x]Queue : {BATCH_GPI_RABBITMQ_QUEUES} | tasks_id：{tasks_id} | progress：{i + 1}/{batch_size} | result_data：{data}")
        else:
            publish_status(tasks_id, f"{i + 1}/{batch_size}", data)
            logger.info(f" [x]Queue : {BATCH_GPI_RABBITMQ_QUEUES} | tasks_id：{tasks_id} | progress：{i + 1}/{batch_size} | result_data：{data}")

    # 任务完成，发送所有数据结果
    if DEBUG:
        print(result_data_list)
        logger.info(f" [x]Queue : {BATCH_GPI_RABBITMQ_QUEUES} | batch_tasks_id：{batch_tasks_id} | progress：OK | result_data_list：{result_data_list}")
        print(f" [x]Queue : {BATCH_GPI_RABBITMQ_QUEUES} | batch_tasks_id：{batch_tasks_id} | progress：OK | result_data_list：{result_data_list}")
    else:
        publish_status(batch_tasks_id, f"OK", result_data_list)
        logger.info(f" [x]Queue : {BATCH_GPI_RABBITMQ_QUEUES} | batch_tasks_id：{batch_tasks_id} | progress：OK | result_data_list：{result_data_list}")


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


def publish_status(task_id, progress, result):
    connection = pika.BlockingConnection(pika.ConnectionParameters(**RABBITMQ_PARAMS))
    channel = connection.channel()
    channel.queue_declare(queue=BATCH_GPI_RABBITMQ_QUEUES, durable=True)
    message = {'task_id': task_id, 'progress': progress, "result": result}
    channel.basic_publish(exchange='',
                          routing_key=BATCH_GPI_RABBITMQ_QUEUES,
                          body=json.dumps(message),
                          properties=pika.BasicProperties(
                              delivery_mode=2,
                          ))
    connection.close()


if __name__ == '__main__':
    # rd = BatchGenerateProductImageModel(
    #     tasks_id="123-15-51-89",
    #     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="overall",
    #     batch_size=20
    # )
    # batch_generate_product(rd.dict())
    # rd = {
    #     "user_id": "89",
    #     "batch_data_list": [
    #         {
    #             "tasks_id": "A-123-15-51-89",
    #             "image_strength": 0.7,
    #             "prompt": " The best quality, ma123sterpiece, real image.Outwear,high quality clothing details,8K realistic,HDR",
    #             "image_url": "aida-results/result_40b1a2fe-e220-11ef-9bfa-0242ac150003.png",
    #             "product_type": "overall",
    #         },
    #         {
    #             "tasks_id": "B-123-15-51-89",
    #             "image_strength": 0.7,
    #             "prompt": " The best quality, masterpiece, real image.Outwear123,high quality clothing details,8K realistic,HDR",
    #             "image_url": "aida-results/result_40b1a2fe-e220-11ef-9bfa-0242ac150003.png",
    #             "product_type": "overall",
    #         }
    #     ]
    # }
    rd = BatchGenerateProductImageModel(
        batch_tasks_id="abcd",
        user_id="89",
        batch_data_list=[
            ProductItemModel(
                tasks_id="123-5464",
                image_strength=0.7,
                product_type="overall",
                image_url="aida-results/result_40b1a2fe-e220-11ef-9bfa-0242ac150003.png",
                prompt="123"
            ),
            ProductItemModel(
                tasks_id="123-5464123",
                image_strength=0.7,
                product_type="overall",
                image_url="aida-results/result_40b1a2fe-e220-11ef-9bfa-0242ac150003.png",
                prompt="123"
            )
        ]
    )
    batch_generate_product(rd.dict())
