import pybullet as p
import time
import threading
from logger import logger


class ConveyorBelt:
    def __init__(
        self, size, position, box_size, box_spacing, box_speed, palletizinFunc
    ):
        self.conveyor_size = size
        self.conveyor_position = position
        self.conveyor_id = self.create_conveyor(size, position)
        self.conveyor_stopped = False

        self.box_size = box_size
        self.box_spacing = box_spacing
        self.box_speed = box_speed
        self.boxes = []
        self.create_boxes()

        self.stop_thread = False
        self.thread = threading.Thread(target=self.run)
        self.thread.start()

        self.palletizingFunc = palletizinFunc

    def create_conveyor(self, size, position):
        collision_shape_id = p.createCollisionShape(
            shapeType=p.GEOM_BOX,
            halfExtents=[size[0] / 2.0, size[1] / 2.0, size[2] / 2.0],
        )
        visual_shape_id = p.createVisualShape(
            shapeType=p.GEOM_BOX,
            halfExtents=[size[0] / 2.0, size[1] / 2.0, size[2] / 2.0],
            rgbaColor=[0.5, 0.5, 0.5, 1],
        )
        conveyor_id = p.createMultiBody(
            baseMass=0.0,
            baseCollisionShapeIndex=collision_shape_id,
            baseVisualShapeIndex=visual_shape_id,
            basePosition=position,
        )
        return conveyor_id

    def box_position(self, i=0):
        x_position = (
            self.conveyor_position[0]
            - self.conveyor_size[0] * 0.5
            + i * (self.box_size[0] + self.box_spacing)
        )
        y_position = (
            self.conveyor_position[1] + self.conveyor_size[1] * 0.5 - self.box_size[1]
        )
        z_position = (
            self.conveyor_position[2] + self.conveyor_size[2] + self.box_size[2] * 0.5
        )
        return [x_position, y_position, z_position]

    def create_boxes(self):
        num_boxes = int(self.conveyor_size[0] / (self.box_size[0] + self.box_spacing))
        for i in range(num_boxes):
            position = self.box_position(i)
            self.boxes.append(self.create_box(position))

    def create_box(self, position):
        collision_shape_id = p.createCollisionShape(
            shapeType=p.GEOM_BOX,
            halfExtents=[
                self.box_size[0] / 2.0,
                self.box_size[1] / 2.0,
                self.box_size[2] / 2.0,
            ],
        )
        visual_shape_id = p.createVisualShape(
            shapeType=p.GEOM_BOX,
            halfExtents=[
                self.box_size[0] / 2.0,
                self.box_size[1] / 2.0,
                self.box_size[2] / 2.0,
            ],
            rgbaColor=[1, 0, 0, 1],
        )
        box_id = p.createMultiBody(
            baseMass=1.0,
            baseCollisionShapeIndex=collision_shape_id,
            baseVisualShapeIndex=visual_shape_id,
            basePosition=position,
            baseOrientation=p.getQuaternionFromEuler([0, 0, 0]),
        )
        return box_id

    def move_boxes(self):
        if not self.conveyor_stopped:
            for box_id in self.boxes[:]:
                pos, orn = p.getBasePositionAndOrientation(box_id)

                if (
                    pos[0] + self.box_size[0] * 0.5
                    > self.conveyor_position[0] + self.conveyor_size[0] * 0.5
                ):
                    # Если коробка выходит за конец конвейера, она падает вниз (физика продолжает действовать)
                    # Удаляем коробку с конвейера
                    self.boxes.remove(box_id)
                    self.conveyor_stopped = True
                    right_wall_position = [pos[0], pos[1] + self.box_size[1] * 0.5, pos[2]]
                    self.call_robot(right_wall_position, box_id)
                else:
                    # Двигаем коробку вдоль конвейера (по оси X)
                    new_pos = [pos[0] + self.box_speed, pos[1], pos[2]]
                    p.resetBasePositionAndOrientation(box_id, new_pos, orn)

        # Добавляем новые коробки, если на конвейере меньше, чем максимально возможное количество
        if len(self.boxes) < int(
            self.conveyor_size[0] / (self.box_size[0] + self.box_spacing)
        ):
            new_box = self.create_box(self.box_position())
            self.boxes.append(new_box)

    def call_robot(self, position, box_id):
        # Здесь можно вызвать робота для выполнения операции, например:
        logger.info("Конвейер остановлен, вызываем робота для работы с коробкой.")
        self.palletizingFunc(position, box_id)
        # time.sleep(5)
        # self.conveyor_stopped = False

    def run(self):
        while not self.stop_thread:
            self.move_boxes()
            time.sleep(1)

    def stop(self):
        self.stop_thread = True
        self.thread.join()