import cv2
import mediapipe as mp
import numpy as np
from minio import Minio

from app.core.config import MINIO_URL, MINIO_ACCESS, MINIO_SECRET, MINIO_SECURE
from app.schemas.mannequin_edit import MannequinModel
from app.service.utils.new_oss_client import oss_get_image, oss_upload_image

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


class MannequinEditService():
    def __init__(self, request_data):
        self.scale = request_data.scale
        self.image = oss_get_image(oss_client=minio_client, bucket=request_data.mannequins.split('/')[0], object_name=request_data.mannequins[request_data.mannequins.find('/') + 1:], data_type="cv2")
        self.mannequin_name = request_data.mannequin_name
        self.bucket_name = request_data.bucket_name
        if self.image.shape[2] == 4:
            self.bgr = self.image[:, :, :3]
            self.alpha = self.image[:, :, 3]
            self.bgr = cv2.bitwise_and(self.bgr, self.bgr, mask=cv2.normalize(self.alpha, None, 0, 1, cv2.NORM_MINMAX))
            self.h, self.w, _ = self.bgr.shape
        else:
            self.bgr = self.image
            self.h, self.w, _ = self.bgr.shape
            self.alpha = None

    def __call__(self, *args, **kwargs):
        leg_top, leg_bottom = self.attitude_detection()
        if leg_top and leg_bottom:
            new_mannequin = self.resize_leg(leg_top, leg_bottom)
            _, encoded_image = cv2.imencode('.png', new_mannequin)
            image_bytes = encoded_image.tobytes()
            req = oss_upload_image(oss_client=minio_client, bucket=self.bucket_name, object_name=f"{self.mannequin_name}.png", image_bytes=image_bytes)
            return req.bucket_name + "/" + req.object_name
        else:
            return "No leg detected"

    def attitude_detection(self):
        mp_pose = mp.solutions.pose
        pose = mp_pose.Pose()

        # 将 BGR 图像转换为 RGB 格式
        image_rgb = cv2.cvtColor(self.bgr, cv2.COLOR_BGR2RGB)
        leg_top, leg_bottom = None, None
        # 进行姿态检测
        results = pose.process(image_rgb)
        if results.pose_landmarks:
            # 获取腿部关键点
            landmarks = results.pose_landmarks.landmark

            # 找到腿部上边界和下边界
            leg_top = int(landmarks[mp_pose.PoseLandmark.LEFT_HIP].y * self.h)
            leg_bottom = int(max(landmarks[mp_pose.PoseLandmark.LEFT_ANKLE].y,
                                 landmarks[mp_pose.PoseLandmark.RIGHT_ANKLE].y) * self.h)

        return leg_top, leg_bottom

    def resize_leg(self, leg_top, leg_bottom):
        # 上半身
        top_part_bgr = self.bgr[:leg_top, :]
        top_part_bgr_alpha = self.alpha[:leg_top, :]

        # 小腿
        part_thigh = self.bgr[leg_top:leg_bottom, :]
        part_thigh_alpha = self.alpha[leg_top:leg_bottom, :]

        # 大腿
        part_calf = self.bgr[leg_bottom:, :]
        part_calf_alpha = self.alpha[leg_bottom:, :]

        new_thigh_height = int((leg_bottom - leg_top) * self.scale[0])
        new_calf_height = int((self.h - leg_bottom) * self.scale[1])

        resized_thigh = cv2.resize(part_thigh, (self.w, new_thigh_height), interpolation=cv2.INTER_LINEAR)
        resized_thigh_alpha = cv2.resize(part_thigh_alpha, (self.w, new_thigh_height), interpolation=cv2.INTER_LINEAR)
        resized_calf = cv2.resize(part_calf, (self.w, new_calf_height), interpolation=cv2.INTER_LINEAR)
        resized_calf_alpha = cv2.resize(part_calf_alpha, (self.w, new_calf_height), interpolation=cv2.INTER_LINEAR)

        new_bgr = np.vstack((top_part_bgr, resized_thigh, resized_calf))
        new_bgr_alpha = np.vstack((top_part_bgr_alpha, resized_thigh_alpha, resized_calf_alpha))

        if self.alpha is not None:
            # 拼接 alpha 通道
            # 合并 BGR 通道和 alpha 通道
            new_image = np.dstack((new_bgr, new_bgr_alpha))
        else:
            new_image = new_bgr
        return new_image


if __name__ == '__main__':
    request_data = MannequinModel(
        mannequins="aida-sys-image/models/male/dc36ce58-46c3-4b6f-8787-5ca7d6fc26e6.png",
        scale=[0.75, 0.75],
        bucket_name="test",
        mannequin_name="mannequin_name"
    )
    service = MannequinEditService(request_data)
    print(service())