"""
ml/MetricsOverlay.py - Camera-Facing (Billboard) 3D Floating Text Overlay
"""
from typing import List, Tuple, Optional, Dict
import numpy as np
import pybullet as p


class MetricsHUD:
    """Renders real-time telemetry as black text floating in 3D, always facing the active camera."""

    def __init__(self, physics_client_id: int = 0):
        self.client_id = physics_client_id
        self._text_id: Optional[int] = None

    def _get_camera_facing_orientation(self) -> List[float]:
        """Calculates a quaternion that points the text towards the current GUI camera."""
        try:
            cam_info = p.getDebugVisualizerCamera(physicsClientId=self.client_id)
            # cam_info index 8: yaw, index 9: pitch
            yaw = cam_info[8]
            pitch = cam_info[9]
            
            pitch_rad = np.radians(pitch + 90.0)
            yaw_rad = np.radians(yaw)
            
            text_orientation = p.getQuaternionFromEuler(
                [pitch_rad, 0.0, yaw_rad], 
                physicsClientId=self.client_id
            )
            return text_orientation
        except Exception:
            return [0.0, 0.0, 0.0, 1.0]

    def update(
        self,
        episode: int,
        step: int,
        robot_rewards: List[float],
        cmd_vel: np.ndarray,
        fps: float = 0.0,
        height: float = 0.0,
        roll_pitch: Tuple[float, float] = (0.0, 0.0),
        mode: str = "direct",
        phase: str = "STAND_ONLY",
        distance: float = 0.0,
        status: str = "ALIVE",
        reward_components: Optional[Dict[str, float]] = None,
        ep_step: int = 0,
    ) -> None:
        """Updates floating text block in 3D space with expanded telemetry."""
        sorted_rewards = sorted(robot_rewards, reverse=True)
        top1 = f"{sorted_rewards[0]:+.2f}" if len(sorted_rewards) > 0 else "0.00"

        vx = cmd_vel[0] if len(cmd_vel) > 0 else 0.0
        vy = cmd_vel[1] if len(cmd_vel) > 1 else 0.0
        omega = cmd_vel[2] if len(cmd_vel) > 2 else 0.0

        lines = [
            "=== JACKBOT TELEMETRY ===",
            f"Mode:        {mode.upper()}",
            f"Curriculum:  {phase}",
            f"Status:      {status}",
            f"Episode:     {episode} (Step {ep_step})",
            f"Global Step: {step}",
            f"FPS:         {fps:.1f}",
            "-------------------------",
            f"Episode Rew: {top1}",
            f"Cmd (X,Y,W): [{vx:+.2f}, {vy:+.2f}, {omega:+.2f}]",
            f"Height:      {height:.3f} m",
            f"Roll/Pitch:  {roll_pitch[0]:+.1f}° / {roll_pitch[1]:+.1f}°",
            f"Max Dist:    {distance:.2f} m",
        ]

        if reward_components:
            lin_v = reward_components.get("lin_vel", 0.0)
            stab = reward_components.get("stability", 0.0)
            h_rew = reward_components.get("height", 0.0)
            jit = reward_components.get("jitter_penalty", 0.0)
            lines.append("--- Reward Components ---")
            lines.append(f"LinVel: {lin_v:.2f} | Stab: {stab:.2f}")
            lines.append(f"Height: {h_rew:.2f} | Jitter: {jit:+.3f}")

        hud_text = "\n".join(lines)

        # Position above origin in simulation world
        text_position = [-0.8, -0.8, 1.2]
        text_color = [0, 0, 0]  # Pure black

        # Calculate dynamic orientation to align text flat against camera plane
        text_orientation = self._get_camera_facing_orientation()

        # Safely remove old text to prevent PyBullet ghosting/overlapping
        if self._text_id is not None:
            try:
                p.removeUserDebugItem(self._text_id, physicsClientId=self.client_id)
            except Exception:
                pass
                
        # Draw fresh text
        self._text_id = p.addUserDebugText(
            text=hud_text,
            textPosition=text_position,
            textColorRGB=text_color,
            textSize=0.085,
            textOrientation=text_orientation,
            physicsClientId=self.client_id
        )

    def reset(self) -> None:
        """Removes the active debug text item so a new episode starts with a clean overlay."""
        if self._text_id is not None:
            try:
                p.removeUserDebugItem(self._text_id, physicsClientId=self.client_id)
            except Exception:
                pass
        self._text_id = None