# model.py (멀티태스크 통합 버전)

import json
import numpy as np
import triton_python_backend_utils as pb_utils
from ultralytics import YOLO
import torch
from pathlib import Path


class TritonPythonModel:
    def initialize(self, args):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        model_dir = Path(args['model_repository']) / args['model_version']
        self.model = YOLO(model_dir / "model.engine")

        # 클래스 이름 로드
        with open(model_dir / "labels.json", 'r', encoding='utf-8') as f:
            self.class_names = {int(k): v for k, v in json.load(f).items()}

        # 모델 디렉토리명에서 task 추론
        model_name = Path(args['model_repository']).name.lower()
        if "pose" in model_name:
            self.task_type = "pose"
        elif "seg" in model_name:
            self.task_type = "seg"
        else:
            self.task_type = "detect"  # 기본: Detection (YOLO, RT-DETR)

        print(f"[INIT] Loaded model {model_name} as task type: {self.task_type}")

    def execute(self, requests):
        responses = []
        for request in requests:
            try:
                raw_image_tensor = pb_utils.get_input_tensor_by_name(request, "RAW_IMAGE")
                shape_tensor = pb_utils.get_input_tensor_by_name(request, "ORIGINAL_IMAGE_SHAPE")

                images_batch = raw_image_tensor.as_numpy()
                original_shapes_batch = shape_tensor.as_numpy()

                results_batch = self.model.predict(
                    list(images_batch), device=self.device, verbose=False
                )

                batch_output = []
                for i, results in enumerate(results_batch):
                    orig_h, orig_w = original_shapes_batch[i]

                    if self.task_type == "detect":
                        # Detection / RT-DETR
                        boxes = results.boxes
                        if boxes is not None and len(boxes) > 0:
                            abs_bboxes = boxes.xyxy.cpu().numpy()
                            rel_bboxes = (
                                abs_bboxes / np.array([orig_w, orig_h, orig_w, orig_h])
                                if orig_w > 0 and orig_h > 0 else np.zeros_like(abs_bboxes)
                            )
                            batch_output.append({
                                "bboxes": abs_bboxes.tolist(),
                                "relative_bboxes": rel_bboxes.tolist(),
                                "labels": [self.class_names.get(int(cls), "unknown") for cls in boxes.cls],
                                "scores": boxes.conf.cpu().numpy().tolist()
                            })
                        else:
                            batch_output.append({"bboxes": [], "relative_bboxes": [], "labels": [], "scores": []})

                    elif self.task_type == "pose":
                        # Pose (boxes + keypoints)
                        boxes = results.boxes
                        kpts = results.keypoints
                        item = {"bboxes": [], "relative_bboxes": [], "labels": [], "scores": [], "keypoints": []}

                        if boxes is not None and len(boxes) > 0:
                            abs_bboxes = boxes.xyxy.cpu().numpy()
                            rel_bboxes = (
                                abs_bboxes / np.array([orig_w, orig_h, orig_w, orig_h])
                                if orig_w > 0 and orig_h > 0 else np.zeros_like(abs_bboxes)
                            )
                            item["bboxes"] = abs_bboxes.tolist()
                            item["relative_bboxes"] = rel_bboxes.tolist()
                            item["labels"] = [self.class_names.get(int(cls), "unknown") for cls in boxes.cls]
                            item["scores"] = boxes.conf.cpu().numpy().tolist()

                        if kpts is not None and hasattr(kpts, "data"):
                            keypoint_data = kpts.data.cpu().numpy()
                            keypoints_xy = keypoint_data[:, :, :2].tolist()
                            keypoints_conf = keypoint_data[:, :, 2].tolist()
                            item["keypoints"] = {
                                "xy": keypoints_xy,
                                "conf": keypoints_conf
                            }

                        batch_output.append(item)

                    elif self.task_type == "seg":
                        # Segmentation (boxes + masks)
                        boxes = results.boxes
                        masks = results.masks
                        item = {"bboxes": [], "relative_bboxes": [], "labels": [], "scores": [], "masks": []}

                        if boxes is not None and len(boxes) > 0:
                            abs_bboxes = boxes.xyxy.cpu().numpy()
                            rel_bboxes = (
                                abs_bboxes / np.array([orig_w, orig_h, orig_w, orig_h])
                                if orig_w > 0 and orig_h > 0 else np.zeros_like(abs_bboxes)
                            )
                            item["bboxes"] = abs_bboxes.tolist()
                            item["relative_bboxes"] = rel_bboxes.tolist()
                            item["labels"] = [self.class_names.get(int(cls), "unknown") for cls in boxes.cls]
                            item["scores"] = boxes.conf.cpu().numpy().tolist()

                        if masks is not None:
                            # 저장 용량 줄이려면 polygon (masks.xy)만 저장 가능
                            item["masks"] = [poly.tolist() for poly in masks.xy]

                        batch_output.append(item)

                # 배치 결과 직렬화
                output_json = json.dumps(batch_output, ensure_ascii=False)
                output_tensor = pb_utils.Tensor("FINAL_RESULT", np.array([output_json], dtype=object))
                responses.append(pb_utils.InferenceResponse(output_tensors=[output_tensor]))

            except Exception as e:
                error = pb_utils.TritonError(f"Model execution failed: {e}")
                responses.append(pb_utils.InferenceResponse(output_tensors=[], error=error))
        return responses

    def finalize(self):
        print("Model unloaded.")