import socket
import json
import os
import math
import numpy as np
import time

from typing import Literal
from pprint import pprint

import threading
import pybullet as p
import pybullet_industrial as pi

from logger import logger


class UrdfManager:
    urdf_filename = None
    urdf_filename_tool = "welding_gun"
    physics_client = None
    body_id = None

    def __init__(self, robot_start_position, application_path):
        self.robot_start_position = robot_start_position
        self.application_path = application_path

    def start_loop(self, urdf):
        self.urdf_filename = urdf

        p.resetSimulation()
        self.load_models()

    def run_pybullet(self, type="DIRECT"):
        self.physics_client = p.connect(getattr(p, type))
        logger.info(f"Connect to {self.physics_client} by {type}")

    def create_cube(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,
            ],  # halfExtents задаются радиусами по осям x, y и z
        )

        # Создание визуальной формы для куба
        visual_shape_id = p.createVisualShape(
            shapeType=p.GEOM_BOX,
            halfExtents=[size[0] / 2.0, size[1] / 2.0, size[2] / 2.0],
            rgbaColor=[0, 0, 1, 1],  # Синий цвет
        )

        # Создание многосоставного тела с формой коллизии и визуальной формой
        body_id = p.createMultiBody(
            baseMass=1.0,  # Масса тела
            baseCollisionShapeIndex=collision_shape_id,
            baseVisualShapeIndex=visual_shape_id,
            basePosition=position,
        )

        return body_id

    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,  # Это статичный объект, поэтому масса 0
            baseCollisionShapeIndex=collision_shape_id,
            baseVisualShapeIndex=visual_shape_id,
            basePosition=position,
        )

        return conveyor_id

    def load_models(self):
        # p.loadURDF("urdf/plane.urdf", physicsClientId=self.physics_client)

        urdf_path = os.path.join(
            self.application_path, "urdf", f"{self.urdf_filename}.urdf"
        )
        self.body_id = p.loadURDF(
            urdf_path, self.robot_start_position, useFixedBase=True
        )

    def load_palletizing(self):
        # Путь к URDF модели палеты
        urdf_path_pallet = os.path.join(
            self.application_path,
            "urdf",
            "europallet.urdf",  # Убедитесь, что имя файла соответствует
        )

        # Загрузка палеты рядом с роботом
        # Предположим, что вы хотите расположить палету на расстоянии от робота
        pallet_offset_x = 1.5  # Расстояние в метрах по оси X
        pallet_offset_y = 0.0  # Расстояние по оси Y
        pallet_height = 0.145  # Высота палеты (подгонка под высоту робота)

        pallet_position = [
            self.robot_start_position[0] + pallet_offset_x,
            self.robot_start_position[1] + pallet_offset_y,
            pallet_height
            / 2.0,  # Разместите палету выше, чтобы она не находилась на земле
        ]
        pallet_orientation = p.getQuaternionFromEuler([0, 0, np.deg2rad(90)])
        self.pallet_id = p.loadURDF(
            urdf_path_pallet, pallet_position, pallet_orientation, useFixedBase=True
        )

        # Создание коробок на палете
        box_size = [0.2, 0.2, 0.2]  # Размеры коробки: [ширина, высота, глубина]
        boxes_per_row = 3
        boxes_per_column = 2

        for row in range(boxes_per_column):
            for col in range(boxes_per_row):
                box_x_offset = col * (box_size[0] + 0.05)
                box_y_offset = row * (box_size[1] + 0.05)
                box_z_position = pallet_height + (box_size[2] / 2)

                box_position = [
                    pallet_position[0]
                    + box_x_offset
                    - (boxes_per_row * (box_size[0] + 0.05) / 2),
                    pallet_position[1]
                    + box_y_offset
                    - (boxes_per_column * (box_size[1] + 0.05) / 2),
                    box_z_position,
                ]

                # Создаем кубик
                self.create_cube(box_size, box_position)

        # Создание конвейера
        conveyor_size = [3.0, 0.5, 0.1]  # Длина, ширина и высота
        conveyor_position = [
            0.5,
            1.5,
            0.5,
        ]  # Позиция в пространстве (выше уровня земли)
        self.conveyor_id = self.create_conveyor(conveyor_size, conveyor_position)

        time.sleep(1)

    def get_pybullet_image(self):
        if self.physics_client is None:
            return

        width, height, rgb, _, _ = p.getCameraImage(
            width=500,
            height=500,
            viewMatrix=p.computeViewMatrix(
                cameraEyePosition=[4, -4, 1.5],
                # cameraEyePosition=[4,-4,4],
                cameraTargetPosition=[0, 0, 0],  # Центр фокусировки камеры
                cameraUpVector=[0, 0, 1],  # Направление вверх
            ),
            projectionMatrix=p.computeProjectionMatrixFOV(
                fov=60.0, aspect=1.0, nearVal=0.1, farVal=10.0
            ),
            # renderer=p.ER_TINY_RENDERER,
            physicsClientId=self.physics_client,
        )

        return (rgb, width, height)