#!/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 io
import json
import logging
from io import BytesIO

import imageio
import numpy as np
import tritonclient.grpc as grpcclient
from PIL import Image
from celery import Celery
from minio import Minio
from tritonclient.utils import np_to_triton_dtype

from app.core.config import *
from app.schemas.pose_transform import BatchPoseTransformModel
from app.service.generate_image.utils.pose_transform_upload import upload_gif, upload_video
from app.service.utils.new_oss_client import oss_upload_image
from app.service.utils.oss_client import oss_get_image

minio_client = Minio(MINIO_URL, access_key=MINIO_ACCESS, secret_key=MINIO_SECRET, secure=MINIO_SECURE)

logger = logging.getLogger()
celery_app = Celery('post_transform_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_post_transform'
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=PT_MODEL_URL)
category = "pose_transform"


def upload_first_image(image, user_id, category, file_name):
    try:
        image_data = io.BytesIO()
        image.save(image_data, format='PNG')
        image_data.seek(0)
        image_bytes = image_data.read()
        object_name = f'{user_id}/{category}/{file_name}'
        req = oss_upload_image(oss_client=minio_client, bucket=GI_MINIO_BUCKET, object_name=object_name, image_bytes=image_bytes)
        image_url = f"aida-users/{object_name}"
        return image_url
    except Exception as e:
        logging.warning(f"upload_png_mask runtime exception : {e}")


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))
    result_image = result_image.convert("RGB")
    image = np.array(result_image)

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

    return image


@celery_app.task
def batch_generate_pose_transform(batch_request_data):
    logger.info(f"batch_generate_pose_transform batch_request_data: {json.dumps(batch_request_data, indent=4)}")
    batch_size = batch_request_data['batch_size']
    image_url = batch_request_data['image_url']
    image = pre_processing_image(image_url)
    pose_num = batch_request_data['pose_id']
    tasks_id = batch_request_data['tasks_id']
    user_id = tasks_id.rsplit('-', 1)[1]

    pose_num = [pose_num] * 1
    pose_num_obj = np.array(pose_num, dtype="object").reshape((-1, 1))
    input_pose_num = grpcclient.InferInput("pose_num", pose_num_obj.shape, np_to_triton_dtype(pose_num_obj.dtype))
    input_pose_num.set_data_from_numpy(pose_num_obj)

    image_files = [image.astype(np.uint8)] * 1
    image_files_obj = np.array(image_files, dtype=np.uint8).reshape((-1, 768, 512, 3))
    input_image_files = grpcclient.InferInput("image_file", image_files_obj.shape, "UINT8")
    input_image_files.set_data_from_numpy(image_files_obj)

    result_url_list = []
    for i in range(batch_size):
        try:
            result = grpc_client.infer(model_name="animatex_1", inputs=[input_pose_num, input_image_files], client_timeout=60000, priority=100)
            result_data = np.squeeze(result.as_numpy("generated_image_list").astype(np.uint8))[:, :, :, ::-1]
            # 第一帧图像
            first_image = Image.fromarray(result_data[0])
            first_image_url = upload_first_image(first_image, user_id=user_id, category=f"{category}_first_img", file_name=f"{tasks_id}_batch_{i}.png")

            # 上传GIF
            gif_buffer = BytesIO()
            imageio.mimsave(gif_buffer, result_data, format='GIF', fps=5)
            gif_buffer.seek(0)
            gif_url = upload_gif(gif_buffer=gif_buffer, user_id=user_id, category=f"{category}_gif", file_name=f"{tasks_id}_batch_{i}.gif")

            # 上传video
            video_url = upload_video(frames=result_data, user_id=user_id, category=f"{category}_video", file_name=f"{tasks_id}_batch_{i}.mp4")
            data = {
                "gif_url": gif_url,
                "video_url": video_url,
                "first_image_url": first_image_url,
            }
        except Exception as e:
            print(e)
            data = {}
        result_url_list.append(data)
        if DEBUG is False:
            if i + 1 < batch_size:
                publish_status(tasks_id, f"{i + 1}/{batch_size}", data)
                logger.info(f" [x]Queue : {BATCH_PS_RABBITMQ_QUEUES} | tasks_id：{tasks_id} | progress：{i + 1}/{batch_size} | image_url：{image_url}")
                # print(f" [x]Queue : {BATCH_GRI_RABBITMQ_QUEUES} | tasks_id：{tasks_id} | progress：{i + 1}/{batch_size} | image_url：{image_url}")
            else:
                publish_status(tasks_id, f"OK", result_url_list)
                logger.info(f" [x]Queue : {BATCH_PS_RABBITMQ_QUEUES} | tasks_id：{tasks_id} | progress：OK | image_url：{image_url}")
                # print(f" [x]Queue : {BATCH_PS_RABBITMQ_QUEUES} | tasks_id：{tasks_id} | progress：OK | image_url：{image_url}")


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


if __name__ == '__main__':
    rd = BatchPoseTransformModel(
        tasks_id="123-89",
        image_url='aida-results/result_0000b606-1902-11ef-9424-0242ac180002.png',
        pose_id="1",
        batch_size=10
    )
    batch_generate_pose_transform(rd.dict())
