# Copyright 2020-2026 Ternaris
# SPDX-License-Identifier: Apache-2.0
"""Tools for reading all rosbag versions with unified api."""

from __future__ import annotations

import functools
import operator
from contextlib import suppress
from heapq import merge
from itertools import groupby
from typing import TYPE_CHECKING, cast

from rosbags.interfaces import MessageDefinition, MessageDefinitionFormat, TopicInfo
from rosbags.rosbag1 import (
    Reader as Reader1,
    ReaderError as ReaderError1,
)
from rosbags.rosbag2 import (
    Reader as Reader2,
    ReaderError as ReaderError2,
)
from rosbags.typesys import Stores, get_types_from_idl, get_types_from_msg, get_typestore

if TYPE_CHECKING:
    import sys
    from collections.abc import Generator, Iterable, Sequence
    from types import TracebackType
    from typing import Literal

    if sys.version_info >= (3, 11):
        from typing import Self
    else:
        from typing_extensions import Self

    from rosbags.interfaces import Connection
    from rosbags.interfaces.typing import RPath, Typesdict
    from rosbags.typesys.store import Typestore


class AnyReaderError(Exception):
    """Reader error."""


ReaderErrors = (ReaderError1, ReaderError2)


class AnyReader:
    """Unified rosbag1 and rosbag2 reader."""

    readers: Sequence[Reader1] | Sequence[Reader2]
    typestore: Typestore

    def __init__(
        self,
        paths: Sequence[RPath],
        *,
        default_typestore: Typestore | None = None,
    ) -> None:
        """Initialize RosbagReader.

        Opens one or multiple rosbag1 recordings or a single rosbag2 recording.

        Args:
            paths: Paths to multiple rosbag1 files or single rosbag2 directory.
            default_typestore: Typestore to deserialize messages if bag has
                no embedded message definions.

        Raises:
            AnyReaderError: If paths do not exist.

        """
        if not paths:
            msg = 'Must call with at least one path.'
            raise AnyReaderError(msg)

        if missing := [x for x in paths if not x.exists()]:
            msg = f'The following paths are missing: {missing!r}'
            raise FileNotFoundError(msg)

        self.paths = paths
        self.is2 = any(x.suffix != '.bag' for x in paths)
        self.isopen = False
        self.connections: list[Connection] = []
        self.default_typestore = default_typestore
        self.typestore = get_typestore(Stores.EMPTY)

        try:
            if self.is2:
                self.readers = [Reader2(x) for x in paths]
            else:
                self.readers = [Reader1(x) for x in paths]
        except ReaderErrors as err:
            raise AnyReaderError(*err.args) from err

    def _deser_ros1(self, rawdata: bytes, typ: str) -> object:
        """Deserialize ROS1 message."""
        return self.typestore.deserialize_ros1(rawdata, typ)

    def _deser_ros2(self, rawdata: bytes, typ: str) -> object:
        """Deserialize CDR message."""
        return self.typestore.deserialize_cdr(rawdata, typ)

    def deserialize(self, rawdata: bytes, typ: str) -> object:
        """Deserialize message with appropriate helper."""
        return self._deser_ros2(rawdata, typ) if self.is2 else self._deser_ros1(rawdata, typ)

    def open(self) -> None:
        """Open rosbags."""
        assert not self.isopen
        rollback: list[Reader1 | Reader2] = []
        try:
            for reader in self.readers:
                reader.open()
                rollback.append(reader)
        except ReaderErrors as err:
            for reader in rollback:
                with suppress(*ReaderErrors):
                    reader.close()
            raise AnyReaderError(*err.args) from err

        typs: Typesdict = {}
        self.connections = [y for x in self.readers for y in x.connections]
        if self.connections:
            connections = [
                x for x in self.connections if x.msgdef.format != MessageDefinitionFormat.NONE
            ]
            if connections:
                sep = '=' * 80 + '\n'
                for connection in connections:
                    if connection.msgdef.data.startswith(f'{sep}IDL: '):
                        for msgdef in connection.msgdef.data.split(sep)[1:]:
                            hdr, idl = msgdef.split('\n', 1)
                            assert hdr.startswith('IDL: ')
                            typs.update(get_types_from_idl(idl))
                    else:
                        typs.update(get_types_from_msg(connection.msgdef.data, connection.msgtype))

            elif self.default_typestore:
                typs.update(self.default_typestore.fielddefs)
            else:
                msg = (
                    'Bag contains no type definitions. '
                    'Instantiate AnyReader with a default_typestore argument.'
                )
                raise AnyReaderError(msg)
        self.typestore.register(typs)
        self.isopen = True

    def close(self) -> None:
        """Close rosbag."""
        assert self.isopen
        for reader in self.readers:
            with suppress(*ReaderErrors):
                reader.close()
        self.isopen = False

    def __enter__(self) -> Self:
        """Open rosbags when entering contextmanager."""
        self.open()
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> Literal[False]:
        """Close rosbags when exiting contextmanager."""
        self.close()
        return False

    @property
    def duration(self) -> int:
        """Duration in nanoseconds between earliest and latest messages."""
        return self.end_time - self.start_time

    @property
    def start_time(self) -> int:
        """Timestamp in nanoseconds of the earliest message."""
        return min(x.start_time for x in self.readers)

    @property
    def end_time(self) -> int:
        """Timestamp in nanoseconds after the latest message."""
        return max(x.end_time for x in self.readers)

    @property
    def message_count(self) -> int:
        """Total message count."""
        return sum(x.message_count for x in self.readers)

    @property
    def topics(self) -> dict[str, TopicInfo]:
        """Topics stored in the rosbags."""
        assert self.isopen

        def summarize(names_infos: Iterable[tuple[str, TopicInfo]]) -> TopicInfo:
            """Summarize topic infos."""
            infos = [x[1] for x in names_infos]
            return TopicInfo(
                msgtypes.pop() if len(msgtypes := {x.msgtype for x in infos}) == 1 else None,
                msgdefs.pop()
                if len(msgdefs := {x.msgdef for x in infos}) == 1
                else MessageDefinition(MessageDefinitionFormat.NONE, ''),
                sum(x.msgcount for x in infos),
                functools.reduce(operator.iadd, (x.connections for x in infos), []),
            )

        return {
            name: summarize(infos)
            for name, infos in groupby(
                sorted(
                    (x for reader in self.readers for x in reader.topics.items()),
                    key=lambda x: x[0],
                ),
                key=lambda x: x[0],
            )
        }

    def messages(
        self,
        connections: Iterable[Connection] = (),
        start: int | None = None,
        stop: int | None = None,
    ) -> Generator[tuple[Connection, int, bytes], None, None]:
        """Read messages from bags.

        Args:
            connections: Iterable with connections to filter for. An empty
                iterable disables filtering on connections.
            start: Yield only messages at or after this timestamp (ns).
            stop: Yield only messages before this timestamp (ns).

        Yields:
            Tuples of connection, timestamp (ns), and rawdata.

        """
        assert self.isopen

        def get_owner(connection: Connection) -> Reader1 | Reader2:
            return cast('Reader1 | Reader2', connection.owner)

        if connections:
            generators = [
                reader.messages(connections=list(conns), start=start, stop=stop)
                for reader, conns in groupby(
                    sorted(connections, key=lambda x: id(get_owner(x))), key=get_owner
                )
            ]
        else:
            generators = [reader.messages(start=start, stop=stop) for reader in self.readers]
        yield from merge(*generators, key=lambda x: x[1])
