#!/usr/bin/env python3
"""Record, inspect, and replay six supplied ROS 2 messages; no robot required."""

import argparse
import json
from pathlib import Path
import signal
import subprocess
import tempfile
import time

import rclpy
from rclpy.node import Node
from rclpy.parameter import Parameter
from rclpy.serialization import deserialize_message
from geometry_msgs.msg import PointStamped
from rosgraph_msgs.msg import Clock
import rosbag2_py
from rosbag2_interfaces.srv import Resume

TOPIC = "/bag_demo/point"
TYPE = "geometry_msgs/msg/PointStamped"
NS = 1_000_000_000
COUNT = 6
TIMEOUT = 15.0


def require(condition, message):
    if not condition:
        raise AssertionError(message)


def sample(index):
    message = PointStamped()
    message.header.stamp.sec = 100 + index
    message.header.frame_id = "map"
    message.point.x = 0.5 * index
    message.point.y = -0.25 * index
    return message


def fields(message):
    return {
        "header_ns": message.header.stamp.sec * NS + message.header.stamp.nanosec,
        "frame": message.header.frame_id,
        "point": [message.point.x, message.point.y, message.point.z],
    }


EXPECTED = [fields(sample(index)) for index in range(COUNT)]


def spin_until(node, predicate, description, timeout=TIMEOUT):
    deadline = time.monotonic() + timeout
    while not predicate():
        require(time.monotonic() < deadline, f"Timed out: {description}")
        rclpy.spin_once(node, timeout_sec=0.02)


def settle(node, seconds):
    """Process callbacks using a wall-time bound even when ROS time is zero."""
    deadline = time.monotonic() + seconds
    while time.monotonic() < deadline:
        rclpy.spin_once(node, timeout_sec=0.02)


def start(command, path):
    log = path.open("w")
    try:
        process = subprocess.Popen(command, stdout=log, stderr=subprocess.STDOUT)
    except BaseException:
        log.close()
        raise
    return process, log


def stop(process, log):
    if process.poll() is None:
        process.send_signal(signal.SIGINT)
        try:
            process.wait(timeout=10)
        except subprocess.TimeoutExpired:
            process.kill()
            process.wait(timeout=5)
    log.close()


class Fixture(Node):
    def __init__(self):
        super().__init__("bag_lesson_fixture")
        self.point_publisher = self.create_publisher(PointStamped, TOPIC, 10)
        self.clock_publisher = self.create_publisher(Clock, "/clock", 10)
        self.clock_seconds = 100
        # This node uses the default system clock. Its timer can run before any
        # simulation clock exists, and repeats the current supplied clock value.
        self.timer = self.create_timer(0.02, self.publish_clock)

    def publish_clock(self):
        clock = Clock()
        clock.clock.sec = self.clock_seconds
        self.clock_publisher.publish(clock)


def record(work):
    bag = work / "fixture_bag"
    print("$ ros2 bag record --storage mcap --output fixture_bag --use-sim-time "
          "--topics /bag_demo/point --disable-keyboard-controls", flush=True)
    source = Fixture()
    recorder, log = start([
        "ros2", "bag", "record", "--storage", "mcap", "--output", str(bag),
        "--use-sim-time", "--topics", TOPIC, "--disable-keyboard-controls",
    ], work / "record.log")
    try:
        # The recorder waits for its first /clock before creating its data
        # subscription. Keep spinning the fixture's clock timer during discovery.
        spin_until(source, lambda: source.point_publisher.get_subscription_count() == 1,
                   "recorder subscription (while publishing /clock)")
        for index in range(COUNT):
            source.clock_seconds = 100 + index
            # Give the separate recorder process time to receive this clock
            # update. This is not a measured deadline or latency assertion.
            # The saved timestamps are independently checked after recording.
            settle(source, 0.25)
            source.point_publisher.publish(sample(index))
            require(source.point_publisher.wait_for_all_acked(
                rclpy.duration.Duration(seconds=5)), "Recorder did not acknowledge sample")
            settle(source, 0.1)
    finally:
        stop(recorder, log)
        source.destroy_node()
    require(recorder.returncode == 0, f"Recorder exited {recorder.returncode}")
    print("RECORDED: six supplied PointStamped messages; /clock was not recorded")
    return bag


def inspect(bag, work):
    print("$ ros2 bag info fixture_bag", flush=True)
    result = subprocess.run(["ros2", "bag", "info", str(bag)], text=True,
                            capture_output=True, timeout=15, check=True)
    (work / "bag-info.txt").write_text(result.stdout)
    # Show the actual info command's stable fields; the full output is retained
    # with --output, including file size, start time, end time, and duration.
    for line in result.stdout.splitlines():
        if line.lstrip().startswith(("Storage id:", "Messages:", "Topic information:")):
            print(line.strip())
    metadata = rosbag2_py.Info().read_metadata(str(bag), "mcap")
    require(metadata.message_count == COUNT, "Bag metadata count differs from six")
    require(len(metadata.topics_with_message_count) == 1, "Unexpected topic in bag")
    topic = metadata.topics_with_message_count[0]
    require(topic.topic_metadata.name == TOPIC and topic.topic_metadata.type == TYPE,
            "Bag metadata topic/type differs from fixture")
    require(topic.message_count == COUNT, "Per-topic count differs from six")
    reader = rosbag2_py.SequentialReader()
    reader.open(rosbag2_py.StorageOptions(uri=str(bag), storage_id="mcap"),
                rosbag2_py.ConverterOptions("", ""))
    rows = []
    while reader.has_next():
        name, data, received_ns = reader.read_next()
        require(name == TOPIC, "Unexpected recorded topic")
        row = fields(deserialize_message(data, PointStamped))
        row["recorded_ns"] = received_ns
        rows.append(row)
    require([{key: row[key] for key in EXPECTED[0]} for row in rows] == EXPECTED,
            "Stored sequence, message contents, or header stamps changed")
    require([row["recorded_ns"] for row in rows] == [row["header_ns"] for row in EXPECTED],
            "Recorder did not receive each supplied /clock before its point")
    print("STORED: count=6; frames=map; header seconds=[100, 101, 102, 103, 104, 105]")
    print("STORED: recorder seconds=[100, 101, 102, 103, 104, 105]")
    print("STORED: points=[(0, 0, 0), (0.5, -0.25, 0), (1, -0.5, 0), "
          "(1.5, -0.75, 0), (2, -1, 0), (2.5, -1.25, 0)]")
    return rows


class Consumer(Node):
    def __init__(self):
        super().__init__("bag_lesson_consumer", parameter_overrides=[
            Parameter("use_sim_time", value=True),
        ])
        self.rows = []
        self.subscription = self.create_subscription(PointStamped, TOPIC, self.receive, 10)

    def receive(self, message):
        row = fields(message)
        row["consumer_now_ns"] = self.get_clock().now().nanoseconds
        self.rows.append(row)


def replay(bag, work, clock):
    label = "with-clock" if clock else "without-clock"
    command = ["ros2", "bag", "play", str(bag), "--start-paused",
               "--disable-keyboard-controls", "--wait-for-all-acked", "3000"]
    if clock:
        command += ["--clock", "20"]
    print("$ " + " ".join("fixture_bag" if arg == str(bag) else arg for arg in command),
          flush=True)
    consumer = Consumer()
    require(consumer.get_clock().now().nanoseconds == 0, "Fresh ROS clock was not zero")
    player, log = start(command, work / f"play-{label}.log")
    try:
        resume = consumer.create_client(Resume, "/rosbag2_player/resume")
        spin_until(consumer, lambda: consumer.count_publishers(TOPIC) == 1
                   and resume.service_is_ready(), "paused player discovery")
        # Let peer discovery settle after observing the publisher and service.
        # This wait alone is not a delivery guarantee: the sequence is checked
        # after playback. --wait-for-all-acked lets reliable acknowledgments drain.
        settle(consumer, 0.5)
        if clock:
            spin_until(consumer, lambda: consumer.get_clock().now().nanoseconds > 0,
                       "paused player's initial /clock")
            require(consumer.get_clock().now().nanoseconds == 100 * NS,
                    "Paused playback did not initialize ROS time to the bag start")
        response = resume.call_async(Resume.Request())
        spin_until(consumer, response.done, "resume service response")
        require(response.exception() is None, "Resume service failed")
        spin_until(consumer, lambda: player.poll() is not None,
                   "finite player completion", timeout=20)
        settle(consumer, 0.25)
        require(player.returncode == 0, f"Player exited {player.returncode}")
        require([{key: row[key] for key in EXPECTED[0]} for row in consumer.rows] == EXPECTED,
                f"Replay {label}: sequence, count, content, frame, or stamp differs")
        times = [row["consumer_now_ns"] for row in consumer.rows]
        (work / f"replay-{label}.json").write_text(json.dumps(consumer.rows, indent=2) + "\n")
        if clock:
            require(all(value >= 100 * NS for value in times),
                    "Replayed ROS time preceded the bag start")
            require(times == sorted(times) and times[-1] > times[0],
                    "ROS clock did not advance monotonically during this single replay")
            print("REPAIR: count=6; all contents/header stamps unchanged; "
                  "consumer ROS time initializes at 100 seconds and advances")
        else:
            require(times == [0] * COUNT, "Clock advanced without any /clock publisher")
            print("FAILURE CASE: count=6; all contents/header stamps unchanged; "
                  "consumer ROS time=[0, 0, 0, 0, 0, 0]")
        return consumer.rows
    finally:
        stop(player, log)
        consumer.destroy_node()


def run(work):
    rclpy.init()
    try:
        bag = record(work)
        stored = inspect(bag, work)
        without = replay(bag, work, clock=False)
        with_clock = replay(bag, work, clock=True)
        (work / "report.json").write_text(json.dumps({
            "topic": TOPIC, "type": TYPE, "expected": EXPECTED, "stored": stored,
            "replay_without_clock": without, "replay_with_clock": with_clock,
        }, indent=2) + "\n")
        print("PASS: record, info, stored payloads, both replay counts, preserved stamps, "
              "missing-clock diagnosis, and clock repair")
    except BaseException:
        for path in sorted(work.glob("*.log")):
            print(f"--- {path.name} ---\n{path.read_text()}", flush=True)
        raise
    finally:
        rclpy.shutdown()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path,
                        help="New directory in which to retain bag, logs, and report.json")
    args = parser.parse_args()
    if args.output:
        # Refuse to reuse an existing path: never overwrite a user's recording.
        args.output.mkdir(parents=True, exist_ok=False)
        run(args.output)
        print(f"ARTIFACTS: {args.output}")
    else:
        with tempfile.TemporaryDirectory(prefix="bag-lesson-") as temporary:
            run(Path(temporary))


if __name__ == "__main__":
    main()