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

import config as cfg
import DataTypes as dt


class Simulation:
    def __init__(self):
        self.physics_client = p.connect(p.GUI)
        self.robot = p.loadURDF(cfg.urdf_path, useFixedBase=True)

        self.revolute_joints = [
            i
            for i in range(p.getNumJoints(self.robot))
            if p.getJointInfo(self.robot, i)[2] == p.JOINT_REVOLUTE
        ]

        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 range(p.getNumJoints(self.robot)):
            joint_info = p.getJointInfo(self.robot, joint_index)
            joint_type = joint_info[2]
            if joint_type == p.JOINT_REVOLUTE:
                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):
        p.disconnect(self.physics_client)


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=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 (Idle -> Walking -> Target Step -> Physics Step)
            robot.tick()
            time.sleep(1.0 / 60.0)
    except KeyboardInterrupt:
        sim_instance.disconnect()