from typing import Dict, Any, List, Tuple
import pybullet as p
import pybullet_data
import numpy as np


class SimManager:
    """
    Manages the PyBullet simulation lifecycle and live HUD overlays.
    """

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

    def connect(self):
        """Connects to PyBullet and sets up the basic physics world."""
        if self.physics_client is not None:
            p.disconnect(self.physics_client)

        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)

        if self.use_gui:
            # Disable unnecessary side panels for a clean UI
            p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0)
            p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0)
            p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0)
            p.configureDebugVisualizer(p.COV_ENABLE_RGB_OUTPUT_PREVIEW, 0)

    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 robots into the active PyBullet simulation."""
        plane_id = p.loadURDF("plane.urdf")
        robots = []
        robot_joint_indices = []

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

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

        return plane_id, robots, robot_joint_indices

    def disconnect(self):
        """Safely disconnects from the simulation."""
        if self.physics_client is not None:
            p.disconnect(self.physics_client)
            self.physics_client = None
            self.debug_text_ids.clear()

    def update_hud(self, stats: Dict[str, Any]):
        """
        Renders live training metrics onto the PyBullet 3D viewport.
        Uses replaceItemUniqueId to prevent flicker.
        """
        if not self.use_gui or self.physics_client is None:
            return

        # Fixed position near the top-left of the origin in 3D world coordinates
        x_pos, y_pos, z_start = -1.2, -1.2, 1.6
        line_height = 0.10

        for i, (label, value) in enumerate(stats.items()):
            if isinstance(value, float):
                display_text = f"{label}: {value:.3f}"
            elif isinstance(value, (list, np.ndarray)):
                formatted_vals = ", ".join(f"{v:.2f}" for v in np.atleast_1d(value))
                display_text = f"{label}: [{formatted_vals}]"
            else:
                display_text = f"{label}: {value}"

            color = [0, 0, 0]  # Black text for clear visibility against the light plane

            if label in self.debug_text_ids:
                p.addUserDebugText(
                    display_text,
                    [x_pos, y_pos, z_start - i * line_height],
                    textColorRGB=color,
                    textSize=1.1,
                    replaceItemUniqueId=self.debug_text_ids[label],
                )
            else:
                self.debug_text_ids[label] = p.addUserDebugText(
                    display_text,
                    [x_pos, y_pos, z_start - i * line_height],
                    textColorRGB=color,
                    textSize=1.1,
                )

    def reset_hud(self):
        """Clears debug text tracking."""
        self.debug_text_ids.clear()