#!/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 BatchGenerateRelightImageModel, RelightItemModel
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()
celery_app = Celery('relight_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_relight'
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
logging.getLogger('pika').setLevel(logging.WARNING)
grpc_client = grpcclient.InferenceServerClient(url=GRI_MODEL_URL)
category = "relight_image"


@celery_app.task
def batch_generate_relight(batch_request_data):
    batch_size = len(batch_request_data['batch_data_list'])
    logger.info(f"batch_generate_relight 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 = []
    negative_prompt = 'lowres, bad anatomy, bad hands, cropped, worst quality'
    seed = "1"

    for i, data in enumerate(batch_request_data['batch_data_list']):
        direction = data['direction']

        prompt = data['prompt']
        product_type = data['product_type']
        image_url = data['image_url']
        image = pre_processing_image(image_url)
        tasks_id = data['tasks_id']

        prompts = [prompt] * 1
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = cv2.resize(image, (512, 768))
        images = [image.astype(np.uint8)] * 1
        seeds = [seed] * 1
        nagetive_prompts = [negative_prompt] * 1
        directions = [direction] * 1

        if 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))
            na_text_obj = np.array(nagetive_prompts, dtype="object").reshape((-1, 1))
            seed_obj = np.array(seeds, dtype="object").reshape((-1, 1))
            direction_obj = np.array(directions, dtype="object").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))
            na_text_obj = np.array(nagetive_prompts, dtype="object").reshape((1))
            seed_obj = np.array(seeds, dtype="object").reshape((1))
            direction_obj = np.array(directions, dtype="object").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_natext = grpcclient.InferInput("negative_prompt", na_text_obj.shape, np_to_triton_dtype(na_text_obj.dtype))
        input_seed = grpcclient.InferInput("seed", seed_obj.shape, np_to_triton_dtype(seed_obj.dtype))
        input_direction = grpcclient.InferInput("direction", direction_obj.shape, np_to_triton_dtype(direction_obj.dtype))

        input_text.set_data_from_numpy(text_obj)
        input_image.set_data_from_numpy(image_obj)
        input_natext.set_data_from_numpy(na_text_obj)
        input_seed.set_data_from_numpy(seed_obj)
        input_direction.set_data_from_numpy(direction_obj)

        inputs = [input_text, input_natext, input_image, input_seed, input_direction]
        try:
            if data['product_type'] == "single":
                result = grpc_client.infer(model_name=GRI_MODEL_NAME_SINGLE, inputs=inputs, priority=100)
                image = result.as_numpy("generated_relight_image")
            else:
                result = grpc_client.infer(model_name=GRI_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:
            print(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_na_text_obj = np.array(nagetive_prompts, dtype="object").reshape((-1, 1))
                e_seed_obj = np.array(seeds, dtype="object").reshape((-1, 1))
                e_direction_obj = np.array(directions, dtype="object").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_natext = grpcclient.InferInput("negative_prompt", e_na_text_obj.shape, np_to_triton_dtype(e_na_text_obj.dtype))
                e_input_seed = grpcclient.InferInput("seed", e_seed_obj.shape, np_to_triton_dtype(e_seed_obj.dtype))
                e_input_direction = grpcclient.InferInput("direction", e_direction_obj.shape, np_to_triton_dtype(e_direction_obj.dtype))

                e_input_text.set_data_from_numpy(e_text_obj)
                e_input_image.set_data_from_numpy(e_image_obj)
                e_input_natext.set_data_from_numpy(e_na_text_obj)
                e_input_seed.set_data_from_numpy(e_seed_obj)
                e_input_direction.set_data_from_numpy(e_direction_obj)

                e_inputs = [e_input_text, e_input_natext, e_input_image, e_input_seed, e_input_direction]

                result = grpc_client.infer(model_name=GRI_MODEL_NAME_SINGLE, inputs=e_inputs, priority=100)
                image = result.as_numpy("generated_relight_image")
                image_result = Image.fromarray(np.squeeze(image.astype(np.uint8)))
            else:
                image_result = str(e)
                logger.error(e)
        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['relight_img'] = image_url

            result_data_list.append(data)
        else:
            image_url = image_result
            data['relight_img'] = image_url
            result_data_list.append(data)

        # 发送每条结果
        if DEBUG:
            logger.info(f" [x]Queue : {BATCH_GRI_RABBITMQ_QUEUES} | tasks_id：{tasks_id} | progress：{i + 1}/{batch_size} | result_data：{data}")
            print(f" [x]Queue : {BATCH_GRI_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_GRI_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_GRI_RABBITMQ_QUEUES} | batch_tasks_id：{batch_tasks_id} | progress：OK | result_data_list：{result_data_list}")
        print(f" [x]Queue : {BATCH_GRI_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_GRI_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 publish_status(task_id, progress, result):
    connection = pika.BlockingConnection(pika.ConnectionParameters(**RABBITMQ_PARAMS))
    channel = connection.channel()
    channel.queue_declare(queue=BATCH_GRI_RABBITMQ_QUEUES, durable=True)
    message = {'task_id': task_id, 'progress': progress, "result": result}
    channel.basic_publish(exchange='',
                          routing_key=BATCH_GRI_RABBITMQ_QUEUES,
                          body=json.dumps(message),
                          properties=pika.BasicProperties(
                              delivery_mode=2,
                          ))
    connection.close()


if __name__ == '__main__':
    rd = BatchGenerateRelightImageModel(
        batch_tasks_id="abcd",
        user_id="89",
        batch_data_list=[
            RelightItemModel(
                tasks_id="123-5464",
                product_type="overall",
                image_url="test/703190759.png",
                prompt="Colorful black",
                direction="Right Light",
            ),
            RelightItemModel(
                tasks_id="123-5464123",
                product_type="overall",
                image_url="test/703190759.png",
                direction="Right Light",
                prompt="Colorful black",
            )
        ]
    )
    batch_generate_relight(rd.dict())
    # X = {
    #     "batch_tasks_id": "abcd",
    #     "user_id": "89",
    #     "batch_data_list": [
    #         {
    #             "tasks_id": "123-5464",
    #             "product_type": "overall",
    #             "image_url": "aida-users/89/product_image/02894523-19b5-46eb-a9c6-2f512f5fec84-0-89.png",
    #             "prompt": "Colorful black",
    #             "direction": "Right Light",
    #         },
    #         {
    #             "tasks_id": "123-5464",
    #             "product_type": "overall",
    #             "image_url": "aida-users/89/product_image/02894523-19b5-46eb-a9c6-2f512f5fec84-0-89.png",
    #             "prompt": "Colorful black",
    #             "direction": "Right Light",
    #         }
    #
    #     ]
    # }
