import os
import cv2
import numpy as np
import supervision as sv
from PIL import Image, ImageDraw, ImageFont
from supervision.annotators.utils import resolve_color
# visualization toos based on supervision
BOUNDING_BOX_ANNOTATOR = sv.BoundingBoxAnnotator(thickness=2)

class LabelAnnotator(sv.LabelAnnotator):

    @staticmethod
    def resolve_text_background_xyxy(
        center_coordinates,
        text_wh,
        position,
    ):
        center_x, center_y = center_coordinates
        text_w, text_h = text_wh
        return center_x, center_y, center_x + text_w, center_y + text_h

    def _draw_labels(
        self,
        scene: np.ndarray,
        labels: list[str],
        label_properties: np.ndarray,
        detections,
        custom_color_lookup,
    ) -> None:
        assert len(labels) == len(label_properties) == len(detections), (
            f"Number of label properties ({len(label_properties)}), "
            f"labels ({len(labels)}) and detections ({len(detections)}) "
            "do not match."
        )

        color_lookup = (
            custom_color_lookup
            if custom_color_lookup is not None
            else self.color_lookup
        )

        font = ImageFont.truetype("simhei.ttf", int(30 * self.text_scale))
        
        for idx, label_property in enumerate(label_properties):
            background_color = resolve_color(
                color=self.color,
                detections=detections,
                detection_idx=idx,
                color_lookup=color_lookup,
            )
            text_color = resolve_color(
                color=self.text_color,
                detections=detections,
                detection_idx=idx,
                color_lookup=color_lookup,
            )

            box_xyxy = label_property[:4]
            text_height_padded = label_property[4]
            self.draw_rounded_rectangle(
                scene=scene,
                xyxy=box_xyxy,
                color=background_color.as_bgr(),
                border_radius=self.border_radius,
            )

            text_x = box_xyxy[0] + self.text_padding
            text_y = box_xyxy[1]
            
            scene_pil = Image.fromarray(cv2.cvtColor(scene, cv2.COLOR_BGR2RGB))
            draw = ImageDraw.Draw(scene_pil)
            draw.text(
                (text_x, text_y),
                labels[idx],
                font=font,
                fill=(text_color.r, text_color.g, text_color.b),
            )
            scene[:] = cv2.cvtColor(np.array(scene_pil), cv2.COLOR_RGB2BGR)


LABEL_ANNOTATOR = LabelAnnotator(text_padding=4,
                                 text_scale=0.5,
                                 text_thickness=1)


POINT_ANNOTATOR = sv.DotAnnotator(radius=6)

def draw_boxes_points_with_labels(
    cv2_image,
    boxes=None,
    points=None,
    classes=None,
    output_path=None,
):
    annotated_image = cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB)

    if boxes is not None and boxes.size:
        detections = sv.Detections(
            xyxy=boxes,
            class_id=np.arange(len(boxes)),
            confidence=np.ones(len(boxes))
        )
        annotated_image = BOUNDING_BOX_ANNOTATOR.annotate(
            annotated_image, detections)
    if points is not None and points.size:
        points = np.concatenate([points, points], axis=1)
        detections = sv.Detections(
            xyxy=points,
            class_id=np.arange(len(points)),
            confidence=np.ones(len(points))
        )
        annotated_image = POINT_ANNOTATOR.annotate(
            annotated_image, detections,
        )
    if classes:
        annotated_image = LABEL_ANNOTATOR.annotate(
            annotated_image, detections, labels=classes
        )

    if output_path:
        cv2.imwrite(
            output_path,
            cv2.cvtColor(annotated_image, cv2.COLOR_RGB2BGR)
        )
    
    return annotated_image