"""
simulation.py - PyBullet Simulation Interface & Standalone Runner
"""
import time
import math
import numpy as np
import pybullet as p
import pybullet_data  # Added missing import

from config import cfg
import DataTypes as dt


class Simulation:
    def __init__(self, urdf_path: str = cfg.urdf_path):
        self.urdf_path = urdf_path
        
        # Connect to PyBullet GUI
        self.physicsClient = p.connect(p.GUI)
        p.setAdditionalSearchPath(pybullet_data.getDataPath())
        p.setGravity(0, 0, -9.81)
        
        # Load plane and robot URDF
        self.planeId = p.loadURDF("plane.urdf")
        self.robot = p.loadURDF(self.urdf_path, [0, 0, 0.2])

        # Discover revolute joint indices dynamically
        self.revolute_joints = []
        for j in range(p.getNumJoints(self.robot)):
            joint_info = p.getJointInfo(self.robot, j)
            if joint_info[2] == p.JOINT_REVOLUTE:
                self.revolute_joints.append(j)

        self.set_all_joints_to_90()
        p.resetDebugVisualizerCamera(
            cameraDistance=1.0,
            cameraYaw=50,
            cameraPitch=-35,
            cameraTargetPosition=[0, 0, 0],
        )

    def set_all_joints_to_90(self):
        for joint_index in self.revolute_joints:
            p.resetJointState(self.robot, joint_index, math.radians(90))

    def updatePos(self, current_rad: dt.RadArray):
        radflat = current_rad.data.flatten()
        for joint_index, target_angle in zip(self.revolute_joints, radflat):
            p.setJointMotorControl2(
                bodyIndex=self.robot,
                jointIndex=joint_index,
                controlMode=p.POSITION_CONTROL,
                targetPosition=target_angle,
                force=500,
            )

    def step(self):
        p.stepSimulation()

    def disconnect(self):
        if p.isConnected(self.physicsClient):
            p.disconnect(self.physicsClient)

    def close(self):
        """Cleanup wrapper for Robot backend compatibility."""
        self.disconnect()


if __name__ == "__main__":
    from Robot import Robot, PyBulletBackend

    # 1. Initialize PyBullet simulation environment
    sim_instance = Simulation()
    backend = PyBulletBackend(sim_instance)

    # 2. Instantiate Robot with simulation backend
    robot = Robot(backend_type=backend)
    robot.reset_to_init()

    # 3. Command forward movement [vx, vy, omega]
    robot.vector_dirmov = [1.0, 0.0, 0.0]

    # 4. Main test execution loop
    try:
        while True:
            # Executes state machine logic
            robot.tick()
            time.sleep(1.0 / 60.0)
    except KeyboardInterrupt:
        sim_instance.disconnect()