rerun 0.35.0

mod info;

use std::collections::BTreeSet;
use std::fs::File;
use std::io::BufWriter;

use clap::Subcommand;
use clap::builder::TypedValueParser as _;
use re_log_encoding::Encoder;
use re_log_types::{Duration, LogMsg, RecordingId, TimeType, Timestamp};
use re_mcap::{DecoderIdentifier, SelectedDecoders, TopicFilter};
use re_sdk::external::re_importer::{McapImporter, supported_mcap_decoder_identifiers};
use re_sdk::{ApplicationId, ImportedData, Importer, ImporterSettings};

use info::InfoCommand;

fn possible_timeline_types() -> impl clap::builder::TypedValueParser {
    clap::builder::PossibleValuesParser::new(["timestamp", "duration"]).map(|value: String| {
        match value.as_str() {
            "timestamp" => TimeType::TimestampNs,
            "duration" => TimeType::DurationNs,
            _ => unreachable!("PossibleValuesParser already validated the input"),
        }
    })
}

fn possible_decoders() -> clap::builder::PossibleValuesParser {
    static DECODER_IDS: std::sync::LazyLock<Vec<String>> = std::sync::LazyLock::new(|| {
        supported_mcap_decoder_identifiers(true)
            .into_iter()
            .map(|identifier| identifier.to_string())
            .collect()
    });
    clap::builder::PossibleValuesParser::new(
        DECODER_IDS.iter().map(String::as_str).collect::<Vec<_>>(),
    )
}

#[derive(Debug, Clone, clap::Parser)]
pub struct ConvertCommand {
    /// Paths to read from. Reads from standard input if none are specified.
    path_to_input_mcap: String,

    /// Path to write to. Writes to standard output if unspecified.
    #[arg(short = 'o', long = "output", value_name = "dst.rrd")]
    path_to_output_rrd: Option<String>,

    /// If set, specifies the application id of the output.
    #[clap(long = "application-id")]
    application_id: Option<String>,

    /// Specifies which decoders to apply during conversion.
    #[clap(short = 'd', long = "decoder", value_parser = possible_decoders())]
    selected_decoders: Vec<String>,

    /// Disable using the raw decoder as a fallback for unsupported channels.
    /// By default, channels that cannot be handled by semantic decoders (protobuf, ROS2)
    /// will be processed by the raw decoder.
    #[clap(long = "disable-raw-fallback")]
    disable_raw_fallback: bool,

    /// If set, specifies the recording id of the output.
    ///
    /// When this flag is set and multiple input .rdd files are specified,
    /// blueprint activation commands will be dropped from the resulting
    /// output.
    #[clap(long = "recording-id")]
    recording_id: Option<String>,

    /// If set, an offset in nanoseconds to add to all timestamp timelines.
    ///
    /// This can be used to shift all timestamps of the MCAP file if they are not yet
    /// relative to the UNIX epoch.
    ///
    /// Duration and sequence timelines are not affected by this offset.
    #[clap(long = "timestamp-offset-ns")]
    timestamp_offset_ns: Option<i64>,

    /// The timeline type to use for timestamp timelines.
    ///
    /// "timestamp" (default) creates `TimestampNs` timelines (nanoseconds since Unix epoch).
    /// "duration" creates `DurationNs` timelines (nanosecond durations).
    #[clap(long = "timeline-type", value_parser = possible_timeline_types(), default_value = "timestamp")]
    timeline_type: TimeType,

    /// Include only topics matching this regex (RE2 syntax). Repeatable.
    ///
    /// If omitted, all topics are included. Patterns are not implicitly anchored;
    /// use `^` / `$` if you need anchoring.
    ///
    /// Example: `-y "^/tf.*" -n ".*depth.*" -y "^/camera/(compressed|camera_info)$"`
    #[clap(short = 'y', long = "include-topic-regex")]
    include_topic_regex: Vec<String>,

    /// Exclude topics matching this regex (RE2 syntax). Repeatable.
    ///
    /// Applied after includes: a topic is kept only if it matches an include
    /// (or no includes are set) AND matches no exclude.
    #[clap(short = 'n', long = "exclude-topic-regex")]
    exclude_topic_regex: Vec<String>,

    /// Inclusive lower bound on the raw MCAP `log_time`.
    ///
    /// Accepts Unix timestamps with a unit suffix (`ns`, `ms`, `s`, …), or an RFC 3339 timestamp.
    /// Bare integers are interpreted as nanoseconds.
    ///
    /// If set, only data within this time range gets converted.
    #[clap(long = "start-time", value_name = "TIME", value_parser = parse_time)]
    start_time: Option<u64>,

    /// Exclusive upper bound on the raw MCAP `log_time`.
    ///
    /// Accepts Unix timestamps with a unit suffix (`ns`, `ms`, `s`, …), or an RFC 3339 timestamp.
    /// Bare integers are interpreted as nanoseconds.
    ///
    /// If set, only data within this time range gets converted.
    #[clap(long = "end-time", value_name = "TIME", value_parser = parse_time)]
    end_time: Option<u64>,

    /// Recover a missing or invalid MCAP summary in memory.
    ///
    /// This allows conversion of MCAP files that lack a footer (e.g. corrupted recordings).
    #[clap(long = "recover")]
    recover: bool,
}

fn compile_topic_filter(include: &[String], exclude: &[String]) -> anyhow::Result<TopicFilter> {
    for pattern in include {
        TopicFilter::default()
            .with_include_patterns(std::slice::from_ref(pattern))
            .map_err(|err| anyhow::anyhow!("Invalid include topic regex {pattern:?}: {err}"))?;
    }
    for pattern in exclude {
        TopicFilter::default()
            .with_exclude_patterns(std::slice::from_ref(pattern))
            .map_err(|err| anyhow::anyhow!("Invalid exclude topic regex {pattern:?}: {err}"))?;
    }

    TopicFilter::default()
        .with_include_patterns(include)
        .and_then(|filter| filter.with_exclude_patterns(exclude))
        .map_err(|err| anyhow::anyhow!("Invalid topic regex in include/exclude filters: {err}"))
}

fn parse_time(value: &str) -> Result<u64, String> {
    if let Ok(nanos) = value.parse::<u64>() {
        return Ok(nanos);
    }

    // `Duration` only provides unit-aware parsing here; the result remains an absolute offset from
    // the Unix epoch, not an offset relative to the start of the MCAP file.
    if let Ok(duration) = value.parse::<Duration>() {
        return u64::try_from(duration.as_nanos())
            .map_err(|_err| "Time cannot be negative".to_owned());
    }

    if let Ok(timestamp) = value.parse::<Timestamp>() {
        return u64::try_from(timestamp.nanos_since_epoch())
            .map_err(|_err| "Time cannot be before the Unix epoch".to_owned());
    }

    Err(format!(
        "invalid time {value:?}; expected nanoseconds, a Unix timestamp with a unit suffix, or an RFC 3339 timestamp"
    ))
}

fn compile_time_range(
    start_time: Option<u64>,
    end_time: Option<u64>,
) -> anyhow::Result<Option<(u64, u64)>> {
    if start_time.is_none() && end_time.is_none() {
        return Ok(None);
    }

    let start = start_time.unwrap_or(0);
    let end = end_time.unwrap_or(u64::MAX);
    anyhow::ensure!(
        start < end,
        "start-time ({start}) must be less than end-time ({end}); the range is half-open [start, end)"
    );

    Ok(Some((start, end)))
}

impl ConvertCommand {
    fn run(&self) -> anyhow::Result<()> {
        let Self {
            path_to_input_mcap,
            path_to_output_rrd,
            application_id,
            recording_id,
            selected_decoders,
            disable_raw_fallback,
            timestamp_offset_ns,
            timeline_type,
            include_topic_regex,
            exclude_topic_regex,
            start_time,
            end_time,
            recover,
        } = self;

        let topic_filter = compile_topic_filter(include_topic_regex, exclude_topic_regex)?;
        let time_range = compile_time_range(*start_time, *end_time)?;

        let start_time = std::time::Instant::now();

        let application_id = application_id
            .to_owned()
            .map(ApplicationId::from)
            .unwrap_or_else(|| ApplicationId::from(path_to_input_mcap.clone()));

        let recording_id = recording_id
            .to_owned()
            .map(RecordingId::from)
            .unwrap_or_else(RecordingId::random);

        let selected_decoders = if selected_decoders.is_empty() {
            SelectedDecoders::All
        } else {
            SelectedDecoders::Subset(
                selected_decoders
                    .iter()
                    .cloned()
                    .map(DecoderIdentifier::from)
                    .collect(),
            )
        };

        let importer: &dyn Importer = &McapImporter::new(&selected_decoders)
            .with_raw_fallback(!*disable_raw_fallback)
            .with_topic_filter(topic_filter)
            .with_time_range(time_range)
            .with_recover(*recover);

        // TODO(#10862): This currently loads the entire file into memory.
        let (tx, rx) = crossbeam::channel::bounded::<ImportedData>(1024);
        importer.import_from_path(
            &ImporterSettings {
                application_id: Some(application_id),
                timestamp_offset_ns: *timestamp_offset_ns,
                timeline_type: *timeline_type,
                ..ImporterSettings::recommended(recording_id)
            },
            path_to_input_mcap.into(),
            tx,
        )?;

        if let Some(path) = path_to_output_rrd {
            let writer = BufWriter::new(File::create(path)?);
            process_mcap(writer, &rx)?;
        } else {
            let stdout = std::io::stdout();
            let lock = stdout.lock();
            let writer = BufWriter::new(lock);
            process_mcap(writer, &rx)?;
        }

        re_log::info!("Processing took {}s", start_time.elapsed().as_secs());

        Ok(())
    }
}

/// Manipulate the contents of .mcap files.
#[derive(Debug, Clone, Subcommand)]
pub enum McapCommands {
    /// Convert an .mcap file to an .rrd
    Convert(ConvertCommand),

    /// Print timeline / sortedness diagnostics for an .mcap file
    Info(InfoCommand),
}

impl McapCommands {
    pub fn run(&self) -> anyhow::Result<()> {
        match self {
            Self::Convert(cmd) => cmd.run(),
            Self::Info(cmd) => cmd.run(),
        }
    }
}

fn process_mcap<W: std::io::Write>(
    writer: W,
    receiver: &crossbeam::channel::Receiver<ImportedData>,
) -> anyhow::Result<()> {
    let mut num_total_msgs = 0;
    let mut topics = BTreeSet::new();
    let options = re_log_encoding::rrd::EncodingOptions::PROTOBUF_COMPRESSED;
    let version = re_build_info::CrateVersion::LOCAL;
    let mut encoder = Encoder::new_eager(version, options, writer)?;

    while let Ok(res) = receiver.recv() {
        num_total_msgs += 1;

        let log_msg = match res {
            ImportedData::LogMsg(_, log_msg) => log_msg,
            ImportedData::Chunk(_, store_id, chunk) => {
                topics.insert(chunk.entity_path().clone());
                let arrow_msg = chunk.to_arrow_msg()?;
                LogMsg::ArrowMsg(store_id, arrow_msg)
            }
            ImportedData::ArrowMsg(_, store_id, arrow_msg) => LogMsg::ArrowMsg(store_id, arrow_msg),
        };
        encoder.append(&log_msg)?;
    }

    re_log::info_once!("Processed {num_total_msgs} messages.");
    re_log::info_once!("Entities: {topics:#?}");
    Ok(())
}