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 .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:
        rr.log("metadata/instructions", rr.TextDocument(
            "\n".join([
                f"Instruction {i+1}:\n" +
                f"  Movement: {inst['cameraMovement']}\n" +
                f"  Easing: {inst['movementEasing']}\n" +
                f"  Frames: {inst['frameCount']}\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: int = None) -> None:
        for idx, subject in enumerate(subjects):
            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,
                    indices=faces,
                    colors=np.tile(subject_color, (len(vertices), 1))
                ),
                timeless=True
            )

            rr.log(f"world/subject_{idx}/class",
                   rr.TextDocument(subject['objectClass']),
                   timeless=True)

    def log_camera_trajectory(self, camera_frames: List[Dict[str, Any]]) -> None:
        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
        )

    def log_camera_frames(self, camera_frames: List[Dict[str, Any]]) -> None:
        for frame_idx, camera_frame in enumerate(camera_frames):
            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['focalLength'],
                    width=1920,
                    height=1080
                )
            )

            rr.log(
                "metadata/current_frame",
                rr.TextDocument(
                    f"Frame: {frame_idx + 1}/{len(camera_frames)}"),
            )