import cv2
import numpy as np
from Interfaces.RealSense import RealSenseController
from Server.ServerProvider import RealSenseObjectDetectionServer

def get_shot(controller):
    depth_matrix = controller.acquisition.get_depth_image()
    color_matrix = controller.acquisition.get_color_image()
    print(f"RS Data: {depth_matrix.shape}")
    return depth_matrix, color_matrix

def draw_boxes_on_image(image, detections, target, follow_vector):
    for detection in detections:
        bbox = detection['bbox']
        class_name = detection['class_name']
        mean_depth = detection['mean_depth']
        x1, y1, x2, y2 = bbox
        color = (0, 0, 255) if detection == target else (0, 255, 0)
        cv2.rectangle(image, (x1, y1), (x2, y2), color, 2)
        label = f"{class_name}: {mean_depth:.2f} cm"
        cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

    follow_text = f"Follow Vector: {follow_vector}"
    cv2.putText(image, follow_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
    return image

def save_results(image, depth_matrix, filename="output_image.png", depth_filename="depth_matrix.npy"):
    cv2.imwrite(filename, image)
    print(f"Image saved as {filename}")
    np.save(depth_filename, depth_matrix)
    print(f"Depth matrix saved as {depth_filename}")

if __name__ == "__main__":
    controller = RealSenseController()
    interface = RealSenseObjectDetectionServer()

    depth_matrix, color_matrix = get_shot(controller)
    response = interface.send_images(color_matrix, depth_matrix)

    if response:
        print("Ответ сервера:", response)
        detections_info = response.get("detections_info", [])
        target = response.get("target", {})
        follow_vector = response.get("follow_vector", [])
        color_matrix_with_boxes = draw_boxes_on_image(color_matrix.copy(), detections_info, target, follow_vector)
        save_results(color_matrix_with_boxes, depth_matrix)