#!/usr/bin/env python3

import argparse
import sys
from pathlib import Path

import cv2
import numpy as np
import rosbag2_py
from rclpy.serialization import deserialize_message
from rosidl_runtime_py.utilities import get_message


def parse_args():
    parser = argparse.ArgumentParser(
        description="Export image data from a ROS 2 bag topic to frames and/or video."
    )
    parser.add_argument(
        "bag",
        help="Path to rosbag directory or single .mcap file.",
    )
    parser.add_argument(
        "--topic",
        required=True,
        help="Image topic to export, e.g. /fays/atrak/cam2/compressed.",
    )
    parser.add_argument(
        "--output-dir",
        default=None,
        help="Directory for extracted image files.",
    )
    parser.add_argument(
        "--video-output",
        default=None,
        help="Output video file path, e.g. /tmp/output.mp4.",
    )
    parser.add_argument(
        "--image-format",
        choices=("jpg", "png"),
        default="jpg",
        help="Extracted frame image format.",
    )
    parser.add_argument(
        "--fps",
        type=float,
        default=30.0,
        help="FPS used when encoding video output.",
    )
    parser.add_argument(
        "--start-index",
        type=int,
        default=0,
        help="Starting frame index for exported images.",
    )
    parser.add_argument(
        "--max-frames",
        type=int,
        default=None,
        help="Optional maximum number of frames to export.",
    )
    return parser.parse_args()


def resolve_bag_uri(bag_arg: str) -> str:
    bag_path = Path(bag_arg).expanduser().resolve()
    if bag_path.is_dir():
        return str(bag_path)
    if bag_path.is_file() and bag_path.suffix == ".mcap":
        return str(bag_path)
    raise FileNotFoundError(f"Bag path does not exist or is unsupported: {bag_path}")


def open_reader(bag_uri: str):
    storage_id = "mcap" if bag_uri.endswith(".mcap") else "sqlite3"
    if bag_uri.endswith(".mcap"):
        storage_id = "mcap"
    else:
        metadata_path = Path(bag_uri) / "metadata.yaml"
        if metadata_path.exists():
            metadata_text = metadata_path.read_text(encoding="utf-8")
            if "storage_identifier: mcap" in metadata_text:
                storage_id = "mcap"

    reader = rosbag2_py.SequentialReader()
    reader.open(
        rosbag2_py.StorageOptions(uri=bag_uri, storage_id=storage_id),
        rosbag2_py.ConverterOptions(
            input_serialization_format="cdr",
            output_serialization_format="cdr",
        ),
    )
    return reader


def get_topic_type(reader, topic_name: str) -> str:
    for topic_info in reader.get_all_topics_and_types():
        if topic_info.name == topic_name:
            return topic_info.type
    raise ValueError(f"Topic not found in bag: {topic_name}")


def decode_image(message, topic_type: str):
    if topic_type == "sensor_msgs/msg/CompressedImage":
        array = np.frombuffer(message.data, dtype=np.uint8)
        frame = cv2.imdecode(array, cv2.IMREAD_COLOR)
        if frame is None:
            raise ValueError("Failed to decode CompressedImage payload")
        return frame

    if topic_type == "sensor_msgs/msg/Image":
        if message.encoding not in ("bgr8", "rgb8", "mono8"):
            raise ValueError(f"Unsupported image encoding: {message.encoding}")

        frame = np.frombuffer(message.data, dtype=np.uint8)
        channels = 1 if message.encoding == "mono8" else 3
        frame = frame.reshape((message.height, message.width, channels))
        if message.encoding == "rgb8":
            frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        elif message.encoding == "mono8":
            frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
        return frame

    raise ValueError(f"Unsupported topic type: {topic_type}")


def build_video_writer(video_path: Path, fps: float, frame_shape):
    height, width = frame_shape[:2]
    video_path.parent.mkdir(parents=True, exist_ok=True)

    candidates = (
        ("avc1", ".mp4"),
        ("mp4v", ".mp4"),
        ("XVID", ".avi"),
    )
    suffix = video_path.suffix.lower()

    for fourcc_name, expected_suffix in candidates:
        if suffix and suffix != expected_suffix:
            continue
        writer = cv2.VideoWriter(
            str(video_path),
            cv2.VideoWriter_fourcc(*fourcc_name),
            fps,
            (width, height),
        )
        if writer.isOpened():
            return writer, fourcc_name
        writer.release()

    for fourcc_name, _ in candidates:
        writer = cv2.VideoWriter(
            str(video_path),
            cv2.VideoWriter_fourcc(*fourcc_name),
            fps,
            (width, height),
        )
        if writer.isOpened():
            return writer, fourcc_name
        writer.release()

    raise RuntimeError(f"Failed to create video writer for: {video_path}")


def main():
    args = parse_args()
    if args.output_dir is None and args.video_output is None:
        print("At least one of --output-dir or --video-output must be set.", file=sys.stderr)
        return 2

    bag_uri = resolve_bag_uri(args.bag)
    reader = open_reader(bag_uri)
    topic_type = get_topic_type(reader, args.topic)
    message_type = get_message(topic_type)

    output_dir = Path(args.output_dir).expanduser().resolve() if args.output_dir else None
    video_path = Path(args.video_output).expanduser().resolve() if args.video_output else None
    if output_dir is not None:
        output_dir.mkdir(parents=True, exist_ok=True)

    writer = None
    writer_codec = None
    frame_index = args.start_index
    exported_count = 0

    print(f"Bag: {bag_uri}")
    print(f"Topic: {args.topic}")
    print(f"Topic type: {topic_type}")
    if output_dir is not None:
        print(f"Image output dir: {output_dir}")
    if video_path is not None:
        print(f"Video output: {video_path}")

    try:
        while reader.has_next():
            current_topic, data, _ = reader.read_next()
            if current_topic != args.topic:
                continue

            msg = deserialize_message(data, message_type)
            frame = decode_image(msg, topic_type)

            if writer is None and video_path is not None:
                writer, writer_codec = build_video_writer(video_path, args.fps, frame.shape)
                print(f"Video codec: {writer_codec}")

            if output_dir is not None:
                image_path = output_dir / f"frame_{frame_index:06d}.{args.image_format}"
                if not cv2.imwrite(str(image_path), frame):
                    raise RuntimeError(f"Failed to write image: {image_path}")

            if writer is not None:
                writer.write(frame)

            frame_index += 1
            exported_count += 1

            if args.max_frames is not None and exported_count >= args.max_frames:
                break
    finally:
        if writer is not None:
            writer.release()

    print(f"Exported frames: {exported_count}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())