import rerun as rr
import numpy as np
from typing import Dict, Any, List
from utils.geometry import vector3_to_numpy, euler_to_quaternion
from visualization.mesh import create_subject_mesh


class SimulationLogger:
    def __init__(self):
        rr.init("camera_simulation")
        rr.log("world", rr.ViewCoordinates.RIGHT_HAND_Y_UP, timeless=True)

    def log_metadata(self, instructions: List[Dict[str, Any]]) -> None:
        if not instructions:
            return

        rr.log("metadata/instructions", rr.TextDocument(
            "\n".join([
                f"Instruction {i+1}:\n" +
                f"  Movement: {inst.get('cameraMovement', 'N/A')}\n" +
                f"  Easing: {inst.get('movementEasing', 'N/A')}\n" +
                f"  Frames: {inst.get('frameCount', 'N/A')}\n" +
                f"  Camera Angle: {inst.get('initialCameraAngle', 'N/A')}\n" +
                f"  Shot Type: {inst.get('initialShotType', 'N/A')}\n" +
                f"  Subject Index: {inst.get('subjectIndex', 'N/A')}"
                for i, inst in enumerate(instructions)
            ])
        ), timeless=True)

    def log_subjects(self, subjects: List[Dict[str, Any]], selected_subject: Optional[int] = None) -> None:
        if not subjects:
            return

        for idx, subject in enumerate(subjects):
            try:
                vertices, faces = create_subject_mesh(subject)
                subject_color = [0.8, 0.2, 0.2, 1.0] if idx == selected_subject else [
                    0.8, 0.8, 0.8, 1.0]

                rr.log(
                    f"world/subject_{idx}",
                    rr.Mesh3D(
                        vertex_positions=vertices,
                        vertex_indices=faces,  # Changed from indices to vertex_indices
                        colors=np.tile(subject_color, (len(vertices), 1))
                    ),
                    timeless=True
                )

                rr.log(f"world/subject_{idx}/class",
                       rr.TextDocument(subject.get('objectClass', 'Unknown')),
                       timeless=True)
            except Exception as e:
                print(f"Error creating mesh for subject {idx}: {str(e)}")

    def log_camera_trajectory(self, camera_frames: List[Dict[str, Any]]) -> None:
        if not camera_frames:
            return

        try:
            camera_positions = np.array([
                vector3_to_numpy(frame['position']) for frame in camera_frames
            ])
            rr.log(
                "world/camera_trajectory",
                rr.Points3D(
                    camera_positions,
                    colors=np.full((len(camera_positions), 4),
                                   [0.0, 0.8, 0.8, 1.0])
                ),
                timeless=True
            )
        except Exception as e:
            print(f"Error logging camera trajectory: {str(e)}")

    def log_camera_frames(self, camera_frames: List[Dict[str, Any]]) -> None:
        if not camera_frames:
            return

        for frame_idx, camera_frame in enumerate(camera_frames):
            try:
                rr.set_time_sequence("frame", frame_idx)

                position = vector3_to_numpy(camera_frame['position'])
                rotation_q = euler_to_quaternion(camera_frame['angle'])

                rr.log(
                    "world/camera",
                    rr.Transform3D(
                        translation=position,
                        rotation=rr.Quaternion(xyzw=rotation_q)
                    )
                )

                rr.log(
                    "world/camera/view",
                    rr.Pinhole(
                        focal_length=camera_frame.get(
                            'focalLength', 50),  # Added default
                        width=1920,
                        height=1080
                    )
                )

                rr.log(
                    "metadata/current_frame",
                    rr.TextDocument(
                        f"Frame: {frame_idx + 1}/{len(camera_frames)}"),
                )
            except Exception as e:
                print(f"Error logging camera frame {frame_idx}: {str(e)}")