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_and_vector(image, detections, follow_vector):
    for detection in detections:
        bbox = detection['bbox']
        class_name = detection['class_name']
        mean_depth = detection['mean_depth']
        x1, y1, x2, y2 = bbox
        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
        label = f"{class_name}: {mean_depth:.2f} cm"
        cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)

    img_center = (image.shape[1] // 2, image.shape[0] // 2)
    vector_endpoint = (int(img_center[0] + follow_vector[0]), int(img_center[1] + follow_vector[1]))
    cv2.arrowedLine(image, img_center, vector_endpoint, (0, 0, 255), 2, tipLength=0.3)
    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", [])
        follow_vector = response.get("follow_vector", [0, 0, 0])
        color_matrix_with_boxes = draw_boxes_and_vector(color_matrix.copy(), detections_info, follow_vector)
        save_results(color_matrix_with_boxes, depth_matrix)