"""
ml/SimManager.py - PyBullet Simulation & Multi-Body Manager
"""
from typing import Dict, List, Tuple
import pybullet as p
import pybullet_data
import DataTypes as dt


class SimManager:
    """Manages PyBullet simulation lifecycle and multi-robot physics."""

    def __init__(self, use_gui: bool = True):
        self.use_gui = use_gui
        self.physics_client = None
        self.robot_joints: Dict[int, List[int]] = {}

    def connect(self):
        """Connects to PyBullet and hides side GUI panels."""
        if self.physics_client is not None and p.isConnected(self.physics_client):
            return

        flags = p.GUI if self.use_gui else p.DIRECT
        self.physics_client = p.connect(flags)

        p.setAdditionalSearchPath(pybullet_data.getDataPath())
        p.setGravity(0, 0, -9.81, physicsClientId=self.physics_client)

        if self.use_gui:
            # Disable PyBullet side panel and preview windows
            p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0, physicsClientId=self.physics_client)
            p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0, physicsClientId=self.physics_client)
            p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client)
            p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 0, physicsClientId=self.physics_client)

    def load_scene(
        self, urdf_path: str, num_robots: int, robot_spacing: float, base_pos_fn
    ) -> Tuple[int, List[int], List[List[int]]]:
        """Loads plane and hexapod bodies into the simulation scene."""
        plane_id = p.loadURDF("plane.urdf", physicsClientId=self.physics_client)
        robots = []
        robot_joint_indices = []
        self.robot_joints.clear()

        for r_id in range(num_robots):
            base_pos = base_pos_fn(r_id, num_robots, robot_spacing)
            robot = p.loadURDF(
                urdf_path, 
                basePosition=base_pos, 
                useFixedBase=False, 
                physicsClientId=self.physics_client
            )
            robots.append(robot)

            joint_indices = [
                i for i in range(p.getNumJoints(robot, physicsClientId=self.physics_client))
                if p.getJointInfo(robot, i, physicsClientId=self.physics_client)[2] == p.JOINT_REVOLUTE
            ]
            robot_joint_indices.append(joint_indices)
            self.robot_joints[robot] = joint_indices

        return plane_id, robots, robot_joint_indices

    def updatePosForBody(self, body_id: int, current_rad: dt.RadArray):
        """Sets joint motor position targets on individual robot bodies."""
        if body_id not in self.robot_joints:
            return
            
        joint_indices = self.robot_joints[body_id]
        radflat = current_rad.data.flatten()
        
        for joint_index, target_angle in zip(joint_indices, radflat):
            p.setJointMotorControl2(
                bodyIndex=body_id,
                jointIndex=joint_index,
                controlMode=p.POSITION_CONTROL,
                targetPosition=float(target_angle),
                force=250,
                physicsClientId=self.physics_client
            )

    def step(self):
        p.stepSimulation(physicsClientId=self.physics_client)

    def disconnect(self):
        if self.physics_client is not None and p.isConnected(self.physics_client):
            p.disconnect(self.physics_client)
            self.physics_client = None