#!/usr/bin/env python3
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
# IMPORTUJEMY PROFIL QoS DLA SENSORÓW
from rclpy.qos import qos_profile_sensor_data 
import numpy as np
import cv2
from ultralytics import YOLO

class YoloGazeboDetector(Node):
    def __init__(self):
        super().__init__('yolo_gazebo_detector')
        
        self.get_logger().info("Ładowanie modelu YOLO...")
        self.model = YOLO('best.pt') 
        self.model.to('cpu') 
        
        self.window_name = "YOLOv8 - Gazebo Camera"
        cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
        cv2.resizeWindow(self.window_name, 640, 480)
        
        self.frame_count = 0
        self.camera_topic = '/arm_camera/image_raw' 
        
        # ZMIANA: Zastępujemy liczbę queue_size gotowym profilem qos_profile_sensor_data
        self.subscription = self.create_subscription(
            Image,
            self.camera_topic,
            self.camera_callback,
            qos_profile_sensor_data
        )
        
        self.window_timer = self.create_timer(0.03, self.refresh_blank_window)
        self.latest_annotated_frame = None
        
        self.get_logger().info(f"Node i Okno GUI gotowe (Zaktualizowano QoS). Słucham na: {self.camera_topic}")

    def refresh_blank_window(self):
        if self.latest_annotated_frame is None:
            loading_frame = np.zeros((480, 640, 3), dtype=np.uint8)
            cv2.putText(loading_frame, "Oczekiwanie na obrazy z Gazebo...", (80, 240),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
            cv2.imshow(self.window_name, loading_frame)
        else:
            cv2.imshow(self.window_name, self.latest_annotated_frame)
        cv2.waitKey(1)

    def camera_callback(self, msg):
        self.frame_count += 1
        
        # Wymuszony log dla każdej klatki, żeby potwierdzić wejście do funkcji
        self.get_logger().info(f"--- ODEBRANO WIADOMOŚĆ (Klatka {self.frame_count}) ---")
        self.get_logger().info(f"Wymiary z nagłówka: Width={msg.width}, Height={msg.height}, Encoding={msg.encoding}")
        self.get_logger().info(f"Rozmiar surowej tablicy danych (msg.data): {len(msg.data)} bajtów")

        if len(msg.data) == 0:
            self.get_logger().error("⚠️ Krytyczny błąd: Tablica msg.data jest całkowicie pusta!")
            return

        try:
            # Konwersja do formatu numpy
            flat_img_array = np.frombuffer(msg.data, dtype=np.uint8)
            
            # Dynamiczne obliczanie oczekiwanego rozmiaru na podstawie kanałów
            channels = 3
            if 'rgba' in msg.encoding or 'bgra' in msg.encoding:
                channels = 4
            elif 'mono' in msg.encoding:
                channels = 1
                
            expected_elements = msg.height * msg.width * channels
            if flat_img_array.size != expected_elements:
                self.get_logger().error(f"⚠️ Niezgodność rozmiarów! Otrzymano {flat_img_array.size} elementów, a oczekiwano {expected_elements}")
                # Próba awaryjnego dopasowania
                cv_image = flat_img_array.reshape((msg.height, msg.step // channels, channels))[:, :msg.width, :]
            else:
                cv_image = flat_img_array.reshape((msg.height, msg.width, channels))
            
            # Obsługa konwersji przestrzeni barw
            if 'rgb' in msg.encoding.lower():
                cv_image = cv2.cvtColor(cv_image, cv2.COLOR_RGB2BGR)
            elif channels == 4:
                cv_image = cv2.cvtColor(cv_image, cv2.COLOR_BGRA2BGR)
                
        except Exception as e:
            self.get_logger().error(f"💥 Błąd krytyczny podczas przetwarzania macierzy: {e}")
            return

        # Jeśli doszliśmy tutaj, obraz jest poprawny. Przekazujemy do YOLO
        try:
            results = self.model(cv_image, verbose=False, imgsz=320) 
            annotated_frame = cv_image.copy()
            for r in results:
                annotated_frame = r.plot()
                
            # Sukces! Podmieniamy klatkę dla okna wyświetlacza
            self.latest_annotated_frame = annotated_frame
        except Exception as e:
            self.get_logger().error(f"Błąd podczas wnioskowania YOLO: {e}")

def main(args=None):
    rclpy.init(args=args)
    node = YoloGazeboDetector()
    
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.get_logger().info("Zamykanie node'a...")
        cv2.destroyAllWindows()
        node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()

if __name__ == '__main__':
    main()