import threading
import json

import cv2

import rclpy
from rclpy.node import Node
from cv_bridge import CvBridge

from interfaces.msg import ImgMsg


lock = threading.Lock()
cv_bridge = CvBridge()
sign = True


def main():
    rclpy.init()
    node = Node("get_image")
    sub = node.create_subscription(ImgMsg, "/img_msg", sub_callback, 10)
    node.get_logger().info("Waiting for image...")
    while sign:
        rclpy.spin_once(node)
    node.get_logger().info("Image received")
    node.destroy_node()
    rclpy.shutdown()


def sub_callback(msg):
    global sign
    camera_data = [
        msg.image_color,
        msg.image_depth,
        msg.karr,
        msg.darr
    ]

    cv2.imwrite("test/color_image.png", cv_bridge.imgmsg_to_cv2(camera_data[0]))
    cv2.imwrite("test/depth_image.png", cv_bridge.imgmsg_to_cv2(camera_data[1]))
    with open("test/K.json", "w") as f:
        f.write(json.dumps(camera_data[2].tolist()))
    with open("test/D.json", "w") as f:
        f.write(json.dumps(camera_data[3].tolist()))

    sign = False


if __name__ == '__main__':
    main()
