import asyncio
from datetime import datetime
from pathlib import Path
from typing import Dict

import rerun as rr
import wandelbots_api_client as wb
from loguru import logger
from nova import MotionGroup
from nova.actions import Action, CombinedActions
from nova.api import models
from nova.core.nova import Nova
from wandelbots_api_client.models import (
    FeedbackOutOfWorkspace,
    PlanTrajectoryFailedResponseErrorFeedback,
)

from nova_rerun_bridge import colors
from nova_rerun_bridge.blueprint import send_blueprint
from nova_rerun_bridge.collision_scene import log_collision_scenes
from nova_rerun_bridge.consts import RECORDING_INTERVAL, TIME_INTERVAL_NAME
from nova_rerun_bridge.helper_scripts.download_models import get_project_root
from nova_rerun_bridge.stream_state import stream_motion_group
from nova_rerun_bridge.trajectory import TimingMode, log_motion


class NovaRerunBridge:
    """Bridge between Nova and Rerun for visualization.

    This class provides functionality to visualize Nova data in Rerun.
    It handles trajectoy, collision scenes, blueprints and proper cleanup of resources.

    Example:
        ```python
        from nova.core.nova import Nova
        from nova_rerun_bridge import NovaRerunBridge

        async def main():
            nova = Nova()
            async with NovaRerunBridge(nova) as bridge:
                await bridge.setup_blueprint()
                await bridge.log_collision_scenes()

        ```

    Args:
        nova (Nova): Instance of Nova client
        spawn (bool, optional): Whether to spawn Rerun viewer. Defaults to True.
    """

    def __init__(self, nova: Nova, spawn: bool = True) -> None:
        self._ensure_models_exist()
        self.nova = nova
        self._streaming_tasks = {}
        if spawn:
            recording_id = f"nova_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
            rr.init(application_id="nova", recording_id=recording_id, spawn=True)
        logger.add(sink=rr.LoggingHandler("logs/handler"))

    def _ensure_models_exist(self):
        """Ensure robot models are downloaded"""
        models_dir = Path(get_project_root()) / "models"
        if not models_dir.exists() or not list(models_dir.glob("*.glb")):
            print("Models not found, run update_robot_models() or poetry run download-models")

    async def setup_blueprint(self) -> None:
        """Configure and send blueprint configuration to Rerun.

        Fetches motion groups from Nova and configures visualization layout.
        """
        cell = self.nova.cell()

        controllers = await cell.controllers()
        motion_groups = []

        if not controllers:
            logger.warning("No controllers found")
            return

        for controller in controllers:
            async with controller[0] as motion_group:
                motion_groups.append(motion_group.motion_group_id)

        send_blueprint(motion_groups)

    async def fetch_and_log_collision_scenes(self) -> Dict[str, models.CollisionScene]:
        """Fetch and log all collision scenes from Nova to Rerun."""
        collision_scenes = (
            await self.nova._api_client.store_collision_scenes_api.list_stored_collision_scenes(
                cell=self.nova.cell()._cell_id
            )
        )
        log_collision_scenes(collision_scenes)
        return collision_scenes

    async def fetch_and_log_collision_scene(
        self, scene_id: str
    ) -> Dict[str, models.CollisionScene]:
        """Log a specific collision scene by its ID.

        Args:
            scene_id (str): The ID of the collision scene to log

        Raises:
            ValueError: If scene_id is not found in stored collision scenes
        """
        collision_scenes = (
            await self.nova._api_client.store_collision_scenes_api.list_stored_collision_scenes(
                cell=self.nova.cell()._cell_id
            )
        )

        if scene_id not in collision_scenes:
            raise ValueError(f"Collision scene with ID {scene_id} not found")

        log_collision_scenes({scene_id: collision_scenes[scene_id]})
        return {scene_id: collision_scenes[scene_id]}

    def log_collision_scene(self, collision_scenes: Dict[str, models.CollisionScene]) -> None:
        log_collision_scenes(collision_scenes=collision_scenes)

    async def log_motion(
        self, motion_id: str, timing_mode=TimingMode.CONTINUE, time_offset: float = 0
    ) -> None:
        # Fetch motion details from api
        motion = await self.nova._api_client.motion_api.get_planned_motion(
            self.nova.cell()._cell_id, motion_id
        )
        optimizer_config = (
            await self.nova._api_client.motion_group_infos_api.get_optimizer_configuration(
                self.nova.cell()._cell_id, motion.motion_group
            )
        )
        trajectory = await self.nova._api_client.motion_api.get_motion_trajectory(
            self.nova.cell()._cell_id, motion_id, int(RECORDING_INTERVAL * 1000)
        )

        motion_groups = await self.nova._api_client.motion_group_api.list_motion_groups(
            self.nova.cell()._cell_id
        )
        motion_motion_group = next(
            (mg for mg in motion_groups.instances if mg.motion_group == motion.motion_group), None
        )

        collision_scenes = (
            await self.nova._api_client.store_collision_scenes_api.list_stored_collision_scenes(
                cell=self.nova.cell()._cell_id
            )
        )

        log_motion(
            motion_id=motion_id,
            model_from_controller=motion_motion_group.model_from_controller,
            motion_group=motion.motion_group,
            optimizer_config=optimizer_config,
            trajectory=trajectory.trajectory,
            collision_scenes=collision_scenes,
            time_offset=time_offset,
            timing_mode=timing_mode,
        )

    async def log_trajectory(
        self,
        joint_trajectory: models.JointTrajectory,
        tcp: str,
        motion_group: MotionGroup,
        timing_mode=TimingMode.CONTINUE,
        time_offset: float = 0,
    ) -> None:
        load_plan_response = await motion_group._load_planned_motion(joint_trajectory, tcp)
        await self.log_motion(
            load_plan_response.motion, timing_mode=timing_mode, time_offset=time_offset
        )

    async def log_error_feedback(
        self, error_feedback: PlanTrajectoryFailedResponseErrorFeedback
    ) -> None:
        if isinstance(error_feedback.actual_instance, FeedbackOutOfWorkspace):
            rr.log(
                "motion/errors/FeedbackOutOfWorkspace",
                rr.Points3D(
                    [
                        error_feedback.actual_instance.invalid_tcp_pose.position[0],
                        error_feedback.actual_instance.invalid_tcp_pose.position[1],
                        error_feedback.actual_instance.invalid_tcp_pose.position[2],
                    ],
                    radii=rr.Radius.ui_points([5.0]),
                    colors=[(255, 0, 0, 255)],
                    labels=["Out of Workspace"],
                ),
                timeless=True,
            )

    async def start_streaming(self, motion_group: MotionGroup) -> None:
        """Start streaming real-time robot state to Rerun viewer."""
        if motion_group in self._streaming_tasks:
            return

        task = asyncio.create_task(stream_motion_group(self, motion_group=motion_group))
        self._streaming_tasks[motion_group] = task

    async def stop_streaming(self) -> None:
        """Stop all streaming tasks."""
        for task in self._streaming_tasks.values():
            task.cancel()
        self._streaming_tasks.clear()

    async def log_actions(self, actions: list[Action] | Action) -> wb.models.JointTrajectory:
        from nova_rerun_bridge import trajectory

        rr.set_time_seconds(TIME_INTERVAL_NAME, trajectory._last_end_time)

        if not isinstance(actions, list):
            actions = [actions]

        if len(actions) == 0:
            raise ValueError("No actions provided")

        poses = CombinedActions(items=tuple(actions)).poses()
        positions = []

        # Collect all positions
        for pose in poses:
            logger.debug(f"Pose: {pose}")
            positions.append([pose.position.x, pose.position.y, pose.position.z])

        # Log all positions at once
        rr.log(
            "motion/actions",
            rr.Points3D(
                positions,
                colors=[colors.colors[0]] * len(positions),  # Green points
                radii=rr.Radius.ui_points([1.0]) * len(positions),
            ),
            timeless=True,
            static=True,
        )

        rr.log(
            "motion/actions/connection", rr.LineStrips3D([positions], colors=[155, 155, 155, 50])
        )

    async def __aenter__(self) -> "NovaRerunBridge":
        """Context manager entry point.

        Returns:
            NovaRerunBridge: Self reference for context manager usage.
        """
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        """Context manager exit point, ensures cleanup."""
        await self.cleanup()

    async def cleanup(self) -> None:
        """Cleanup resources and close Nova API client connection."""
        if hasattr(self.nova, "_api_client"):
            await self.nova._api_client.close()
