import time

from PyQt5.QtWidgets import QWidget, QVBoxLayout, QHBoxLayout, QLabel, QSizePolicy
from PyQt5.QtGui import QImage, QPixmap
from PyQt5.QtCore import Qt, QTimer

from logger import logger


class Visualize(QWidget):
    def __init__(self, get_pybullet_image):
        super().__init__()
        self.get_pybullet_image = get_pybullet_image

        # Настройка компоновки
        self.layout = QVBoxLayout(self)
        self.layout.setContentsMargins(0, 0, 0, 0)  # Убираем отступы по краям

        self.h_layout = QHBoxLayout()
        self.h_layout.setContentsMargins(0, 0, 0, 0)  # Убираем отступы по краям

        self.label = QLabel(self)
        self.layout.addWidget(self.label)

        # Таймер для обновления изображения
        self.timer = QTimer()
        self.timer.timeout.connect(self.update_image)
        self.timer.start(int(500))  # Обновление каждые 100 мс

    def update_image(self):
        res = self.get_pybullet_image()
        
        if not res: return
        
        (rgb, width, height) = res

        image = QImage(rgb, width, height, QImage.Format_RGBA8888)
        image = image.convertToFormat(QImage.Format_RGB888)
        pixmap = QPixmap.fromImage(image)

        self.label.setPixmap(pixmap)
        self.label.setAlignment(Qt.AlignmentFlag.AlignHCenter)
        self.label.resize(image.size()) 
        self.label.repaint()

    def paintEvent(self, event):
        p = self.palette()
        p.setColor(self.backgroundRole(), Qt.magenta)
        self.setPalette(p)
        super().paintEvent(event)