"""ROS bag loader for ROS 1 (.bag) and ROS 2 (.mcap) files."""

from __future__ import annotations

import logging
from pathlib import Path

import numpy as np
import pandas as pd

from orbit.analyzer.dataset_loader import DatasetInfo
from orbit.loaders.base import BaseLoader, normalize_dataframe, sample_episode_indices

logger = logging.getLogger(__name__)


def _check_mcap():
    try:
        import mcap  # noqa: F401

        return True
    except ImportError:
        return False


def _check_rosbag():
    try:
        import rosbag  # noqa: F401

        return True
    except ImportError:
        return False


class RosBagLoader(BaseLoader):
    """Load ROS 1 .bag and ROS 2 .mcap files.

    Configuration via keyword args or CLI flags:
        action_topic: Topic containing action commands (default: auto-detect)
        state_topic: Topic containing robot state (default: auto-detect)
        episode_gap: Seconds of silence to split episodes (default: 2.0)
        fps: Target FPS for resampling (default: infer from data)
    """

    # Common topic patterns for auto-detection
    ACTION_TOPIC_PATTERNS = [
        "/joint_commands",
        "/action",
        "/cmd_vel",
        "/joint_group_position_controller/command",
        "/target_joint_positions",
    ]
    STATE_TOPIC_PATTERNS = [
        "/joint_states",
        "/robot_state",
        "/state",
        "/joint_group_position_controller/state",
    ]

    def __init__(
        self,
        action_topic: str | None = None,
        state_topic: str | None = None,
        episode_gap: float = 2.0,
        fps: float | None = None,
    ):
        self.action_topic = action_topic
        self.state_topic = state_topic
        self.episode_gap = episode_gap
        self.target_fps = fps

    def can_load(self, source: str) -> bool:
        """True if source is a .bag or .mcap file."""
        p = Path(source).expanduser()
        return p.is_file() and p.suffix in (".bag", ".mcap")

    def load_info(self, source: str) -> DatasetInfo:
        """Load metadata from ROS bag."""
        p = Path(source).expanduser().resolve()

        if p.suffix == ".mcap":
            return self._load_info_mcap(p)
        elif p.suffix == ".bag":
            return self._load_info_rosbag(p)
        else:
            raise ValueError(f"Unsupported ROS bag format: {p.suffix}")

    def load_episodes(
        self,
        source: str,
        episode_indices: list[int] | None = None,
        max_episodes: int | None = None,
    ) -> pd.DataFrame:
        """Load episode data from ROS bag."""
        if max_episodes is None:
            max_episodes = 100

        p = Path(source).expanduser().resolve()

        if p.suffix == ".mcap":
            raw = self._load_raw_mcap(p)
        elif p.suffix == ".bag":
            raw = self._load_raw_rosbag(p)
        else:
            raise ValueError(f"Unsupported ROS bag format: {p.suffix}")

        if not raw["timestamps"]:
            return normalize_dataframe(pd.DataFrame())

        # Segment into episodes based on timestamp gaps
        episodes = self._segment_episodes(raw)

        # Select episodes
        if episode_indices is not None:
            episodes = [
                ep for i, ep in enumerate(episodes) if i in set(episode_indices)
            ]
        elif len(episodes) > max_episodes:
            sampled = set(sample_episode_indices(len(episodes), max_episodes))
            episodes = [ep for i, ep in enumerate(episodes) if i in sampled]

        # Build DataFrame
        rows: list[dict] = []
        fps = self.target_fps or raw.get("inferred_fps", 30.0)

        for ep_idx, ep in enumerate(episodes):
            for frame_idx, (t, action, state) in enumerate(
                zip(ep["timestamps"], ep["actions"], ep["states"])
            ):
                rows.append(
                    {
                        "episode_index": ep_idx,
                        "frame_index": frame_idx,
                        "action": action,
                        "state": state,
                        "timestamp": t,
                    }
                )

        if not rows:
            return normalize_dataframe(pd.DataFrame())

        return normalize_dataframe(pd.DataFrame(rows))

    def _segment_episodes(self, raw: dict) -> list[dict]:
        """Split raw messages into episodes based on timestamp gaps."""
        timestamps = raw["timestamps"]
        actions = raw["actions"]
        states = raw["states"]

        if not timestamps:
            return []

        episodes: list[dict] = []
        current_ep: dict = {"timestamps": [], "actions": [], "states": []}

        for i in range(len(timestamps)):
            if i > 0 and (timestamps[i] - timestamps[i - 1]) > self.episode_gap:
                if current_ep["timestamps"]:
                    episodes.append(current_ep)
                current_ep = {"timestamps": [], "actions": [], "states": []}

            current_ep["timestamps"].append(timestamps[i])
            current_ep["actions"].append(actions[i])
            current_ep["states"].append(states[i])

        if current_ep["timestamps"]:
            episodes.append(current_ep)

        return episodes

    def _find_topic(self, available_topics: list[str], patterns: list[str]) -> str | None:
        """Find a matching topic from available topics."""
        for pattern in patterns:
            for topic in available_topics:
                if topic == pattern or topic.endswith(pattern):
                    return topic
        return None

    def _extract_numeric_from_msg(self, msg) -> list[float] | None:
        """Extract numeric array from a ROS message."""
        # JointState message
        if hasattr(msg, "position"):
            pos = list(msg.position)
            if pos:
                return pos
        # Float64MultiArray / similar
        if hasattr(msg, "data"):
            data = list(msg.data)
            if data and all(isinstance(x, (int, float)) for x in data):
                return data
        # Twist message
        if hasattr(msg, "linear") and hasattr(msg, "angular"):
            return [
                msg.linear.x,
                msg.linear.y,
                msg.linear.z,
                msg.angular.x,
                msg.angular.y,
                msg.angular.z,
            ]
        return None

    # --- MCAP implementation ---

    def _load_info_mcap(self, p: Path) -> DatasetInfo:
        if not _check_mcap():
            raise ImportError(
                "MCAP format detected but 'mcap' is not installed.\n"
                "  Install with: pip install orbit-robotics[rosbag]"
            )

        from mcap.reader import make_reader

        with open(p, "rb") as f:
            reader = make_reader(f)
            summary = reader.get_summary()

        if summary is None:
            return DatasetInfo(
                repo_id=str(p),
                num_episodes=1,
                num_frames=0,
                fps=30.0,
                robot_type="unknown",
                camera_names=[],
                video_keys=[],
                shapes={},
            )

        total_messages = sum(
            ch.message_count for ch in (summary.channels or {}).values()
        )

        return DatasetInfo(
            repo_id=str(p),
            num_episodes=1,  # Will be determined during load
            num_frames=total_messages,
            fps=30.0,  # Will be inferred during load
            robot_type="unknown",
            camera_names=[],
            video_keys=[],
            shapes={},
        )

    def _load_raw_mcap(self, p: Path) -> dict:
        if not _check_mcap():
            raise ImportError(
                "MCAP format detected but 'mcap' is not installed.\n"
                "  Install with: pip install orbit-robotics[rosbag]"
            )

        from mcap.reader import make_reader

        try:
            from mcap_ros2.decoder import DecoderFactory

            decoder_factories = [DecoderFactory()]
        except ImportError:
            decoder_factories = []

        timestamps: list[float] = []
        actions: list[list[float]] = []
        states: list[list[float]] = []

        with open(p, "rb") as f:
            reader = make_reader(f, decoder_factories=decoder_factories)

            # Discover topics
            summary = reader.get_summary()
            available_topics = []
            if summary and summary.channels:
                available_topics = [ch.topic for ch in summary.channels.values()]

            action_topic = self.action_topic or self._find_topic(
                available_topics, self.ACTION_TOPIC_PATTERNS
            )
            state_topic = self.state_topic or self._find_topic(
                available_topics, self.STATE_TOPIC_PATTERNS
            )

            if not action_topic and not state_topic:
                logger.warning(
                    "No action or state topics found in %s. Available topics: %s",
                    p,
                    available_topics,
                )
                return {
                    "timestamps": [],
                    "actions": [],
                    "states": [],
                    "inferred_fps": 30.0,
                }

            # Read messages
            topics_to_read = [t for t in [action_topic, state_topic] if t]

            action_msgs: list[tuple[float, list[float]]] = []
            state_msgs: list[tuple[float, list[float]]] = []

            for schema, channel, message, decoded in reader.iter_decoded_messages(
                topics=topics_to_read
            ):
                t = message.log_time / 1e9  # nanoseconds to seconds
                values = self._extract_numeric_from_msg(decoded)
                if values is None:
                    continue

                if channel.topic == action_topic:
                    action_msgs.append((t, values))
                elif channel.topic == state_topic:
                    state_msgs.append((t, values))

        # Synchronize by nearest timestamp
        return self._synchronize_messages(action_msgs, state_msgs)

    # --- ROS 1 bag implementation ---

    def _load_info_rosbag(self, p: Path) -> DatasetInfo:
        if not _check_rosbag():
            raise ImportError(
                "ROS 1 bag format detected but 'rosbag' is not installed.\n"
                "  Install with: pip install orbit-robotics[rosbag-ros1]"
            )

        import rosbag

        bag = rosbag.Bag(str(p))
        info = bag.get_type_and_topic_info()

        total_messages = sum(
            topic_info.message_count
            for topic_info in info.topics.values()
        )

        bag.close()

        return DatasetInfo(
            repo_id=str(p),
            num_episodes=1,
            num_frames=total_messages,
            fps=30.0,
            robot_type="unknown",
            camera_names=[],
            video_keys=[],
            shapes={},
        )

    def _load_raw_rosbag(self, p: Path) -> dict:
        if not _check_rosbag():
            raise ImportError(
                "ROS 1 bag format detected but 'rosbag' is not installed.\n"
                "  Install with: pip install orbit-robotics[rosbag-ros1]"
            )

        import rosbag

        bag = rosbag.Bag(str(p))
        info = bag.get_type_and_topic_info()
        available_topics = list(info.topics.keys())

        action_topic = self.action_topic or self._find_topic(
            available_topics, self.ACTION_TOPIC_PATTERNS
        )
        state_topic = self.state_topic or self._find_topic(
            available_topics, self.STATE_TOPIC_PATTERNS
        )

        if not action_topic and not state_topic:
            logger.warning(
                "No action or state topics found in %s. Available topics: %s",
                p,
                available_topics,
            )
            bag.close()
            return {"timestamps": [], "actions": [], "states": [], "inferred_fps": 30.0}

        topics_to_read = [t for t in [action_topic, state_topic] if t]

        action_msgs: list[tuple[float, list[float]]] = []
        state_msgs: list[tuple[float, list[float]]] = []

        for topic, msg, t in bag.read_messages(topics=topics_to_read):
            timestamp = t.to_sec()
            values = self._extract_numeric_from_msg(msg)
            if values is None:
                continue

            if topic == action_topic:
                action_msgs.append((timestamp, values))
            elif topic == state_topic:
                state_msgs.append((timestamp, values))

        bag.close()

        return self._synchronize_messages(action_msgs, state_msgs)

    def _synchronize_messages(
        self,
        action_msgs: list[tuple[float, list[float]]],
        state_msgs: list[tuple[float, list[float]]],
    ) -> dict:
        """Synchronize action and state messages by nearest timestamp."""
        timestamps: list[float] = []
        actions: list[list[float]] = []
        states: list[list[float]] = []

        # Use action timestamps as the primary timeline
        primary = action_msgs if action_msgs else state_msgs
        secondary = state_msgs if action_msgs else action_msgs
        is_action_primary = bool(action_msgs)

        if not primary:
            return {"timestamps": [], "actions": [], "states": [], "inferred_fps": 30.0}

        # Build sorted secondary timestamps for nearest-neighbor lookup
        sec_times = np.array([t for t, _ in secondary]) if secondary else np.array([])
        sec_values = [v for _, v in secondary]

        for t, values in primary:
            if is_action_primary:
                action = values
                if sec_times.size > 0:
                    idx = int(np.argmin(np.abs(sec_times - t)))
                    state = sec_values[idx]
                else:
                    state = None
            else:
                state = values
                if sec_times.size > 0:
                    idx = int(np.argmin(np.abs(sec_times - t)))
                    action = sec_values[idx]
                else:
                    action = None

            timestamps.append(t)
            actions.append(action)
            states.append(state)

        # Infer FPS from median message rate
        inferred_fps = 30.0
        if len(timestamps) > 1:
            dts = np.diff(timestamps)
            median_dt = float(np.median(dts[dts > 0])) if np.any(dts > 0) else 1 / 30
            inferred_fps = round(1.0 / median_dt, 1)

        return {
            "timestamps": timestamps,
            "actions": actions,
            "states": states,
            "inferred_fps": inferred_fps,
        }
