import os
import msgpack
import json


class Detection:
    def __init__(self, annotation_name, cls, x, y, w, h, confidence=None):
        self.annotation_name = annotation_name
        self.cls = cls
        self.x = x
        self.y = y
        self.w = w
        self.h = h
        self.confidence = confidence

    def __str__(self):
        return f'{self.cls}: {self.x:.2f} {self.y:.2f} {self.w:.2f} {self.h:.2f}, prob: {(self.confidence * 100):.1f}%'


class AnnotationCreatedMessage:
    def __init__(self, msgpack_bytes):
        unpacked_data = msgpack.unpackb(msgpack_bytes, raw=False)
        self.createdDate = unpacked_data.get(0)
        self.name = unpacked_data.get(1)
        self.originalMediaName = unpacked_data.get(2)
        self.time = unpacked_data.get(3)
        self.imageExtension = unpacked_data.get(4)
        detections_json_str = unpacked_data.get(5)
        self.detections = self._parse_detections(detections_json_str)
        self.image = unpacked_data.get(6)
        self.createdRole = unpacked_data.get(7)
        self.createdEmail = unpacked_data.get(8)
        self.source = unpacked_data.get(9)
        self.status = unpacked_data.get(10)

    def __str__(self):
        if self.detections:
            detections_str_list = [str(detection) for detection in self.detections]
            detections_str = ", ".join(detections_str_list)
            return f'{self.name}: [{detections_str}]'
        else:
            return f'{self.name}: [Empty]'

    @staticmethod
    def _parse_detections(detections_json_str):
        if detections_json_str:
            detections_list = json.loads(detections_json_str)
            return [Detection(
                d.get('AnnotationName'),
                d.get('ClassNumber'),
                d.get('CenterX'),
                d.get('CenterY'),
                d.get('Height'),
                d.get('Width'),
                d.get('Probability')
            ) for d in detections_list]
        return []

    def save_annotation(self, save_folder):
        image_folder = os.path.join(save_folder, 'images')
        labels_folder = os.path.join(save_folder, 'labels')

        os.makedirs(image_folder, exist_ok=True)
        os.makedirs(labels_folder, exist_ok=True)

        image_path = os.path.join(image_folder, f"{self.name}.{self.imageExtension}")
        label_path = os.path.join(labels_folder, f"{self.name}.txt")

        try:
            with open(image_path, 'wb') as image_file:
                image_file.write(self.image)
            print(f"Image saved to: {image_path}")
        except IOError as e:
            print(f"Error saving image: {e}")

        try:
            with open(label_path, 'w') as label_file:
                if self.detections:
                    label_file.writelines([
                        f'{detection.cls} {detection.x} {detection.y} {detection.w} {detection.h}'
                        for detection in self.detections
                    ])
                else:
                    label_file.write('')
            print(f'Label saved to: {label_path}')
        except IOError as e:
            print(f"Error saving label: {e}")