# Copyright 2022 Open Source Robotics Foundation, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from argparse import FileType

from rclpy.qos import InvalidQoSProfileException
from ros2bag.api import add_standard_reader_args
from ros2bag.api import check_not_negative_int
from ros2bag.api import check_positive_float
from ros2bag.api import convert_yaml_to_qos_profile
from ros2bag.api import print_error
from ros2bag.verb import VerbExtension
from ros2cli.node import NODE_NAME_PREFIX
from rosbag2_py import Player
from rosbag2_py import PlayOptions
from rosbag2_py import StorageOptions
import yaml


class BurstVerb(VerbExtension):
    """Burst data from a bag."""

    def add_arguments(self, parser, cli_name):  # noqa: D102
        add_standard_reader_args(parser)

        parser.add_argument(
            '--read-ahead-queue-size', type=int, default=1000,
            help='size of message queue rosbag tries to hold in memory to help deterministic '
                 'playback. Larger size will result in larger memory needs but might prevent '
                 'delay of message playback.')
        parser.add_argument(
            '--topics', type=str, default=[], nargs='+',
            help='topics to replay, separated by space. If none specified, all topics will be '
                 'replayed.')
        parser.add_argument(
            '--qos-profile-overrides-path', type=FileType('r'),
            help='Path to a yaml file defining overrides of the QoS profile for specific topics.')
        parser.add_argument(
            '--remap', '-m', default='', nargs='+',
            help='list of topics to be remapped: in the form '
                 '"old_topic1:=new_topic1 old_topic2:=new_topic2 etc." ')
        parser.add_argument(
            '--storage-config-file', type=FileType('r'),
            help='Path to a yaml file defining storage specific configurations. '
                 'See storage plugin documentation for the format of this file.')
        parser.add_argument(
            '--start-offset', type=check_positive_float, default=0.0,
            help='Start the playback player this many seconds into the bag file.')
        parser.add_argument(
            '-n', '--num-messages', type=check_not_negative_int, default=0,
            help='Burst the specified number of messages, then pause.')

    def main(self, *, args):  # noqa: D102
        qos_profile_overrides = {}  # Specify a valid default
        if args.qos_profile_overrides_path:
            qos_profile_dict = yaml.safe_load(args.qos_profile_overrides_path)
            try:
                qos_profile_overrides = convert_yaml_to_qos_profile(
                    qos_profile_dict)
            except (InvalidQoSProfileException, ValueError) as e:
                return print_error(str(e))

        storage_config_file = ''
        if args.storage_config_file:
            storage_config_file = args.storage_config_file.name

        topic_remapping = ['--ros-args']
        for remap_rule in args.remap:
            topic_remapping.append('--remap')
            topic_remapping.append(remap_rule)

        storage_options = StorageOptions(
            uri=args.bag_path,
            storage_id=args.storage,
            storage_config_uri=storage_config_file,
        )
        play_options = PlayOptions()
        play_options.read_ahead_queue_size = args.read_ahead_queue_size
        play_options.node_prefix = NODE_NAME_PREFIX
        play_options.rate = 1.0
        play_options.topics_to_filter = args.topics
        play_options.topic_qos_profile_overrides = qos_profile_overrides
        play_options.loop = False
        play_options.topic_remapping_options = topic_remapping
        play_options.clock_publish_frequency = 0
        play_options.delay = 0.0
        play_options.disable_keyboard_controls = False  # Give the user control
        play_options.start_paused = True  # Important for allowing the burst
        play_options.start_offset = args.start_offset
        play_options.wait_acked_timeout = -1

        player = Player()
        player.burst(storage_options, play_options, args.num_messages)
