re_log_encoding 0.36.0

use std::collections::BTreeMap;

use re_async::AsyncReadAt;
use re_log_types::{ApplicationId, LogMsg, StoreId};

use crate::rrd::{
    CodecError, Decodable as _, DecoderEntrypoint as _, MessageHeader, MessageKind, StreamFooter,
    StreamHeader,
};
use crate::{CachingApplicationIdInjector, RrdFooter, ToApplication as _};

/// Read the full RRD footer using positional I/O.
///
/// Returns `Ok(None)` if the data is a valid RRD but has no footer (legacy RRD).
/// Returns `Err` if the data is not a valid RRD or is corrupted.
///
/// The returned [`RrdFooter`] contains manifests for ALL stores in the file.
/// Caller is responsible for selecting the desired store.
pub async fn read_rrd_footer<R: AsyncReadAt>(reader: &R) -> Result<Option<RrdFooter>, CodecError> {
    let Some(payload) = read_rrd_footer_payload(reader).await? else {
        return Ok(None);
    };

    let transport_footer = re_protos::log_msg::v1alpha1::RrdFooter::from_rrd_bytes(&payload)?;
    Ok(Some(transport_footer.to_application(())?))
}

pub(super) async fn read_rrd_footer_payload<R: AsyncReadAt>(
    reader: &R,
) -> Result<Option<bytes::Bytes>, CodecError> {
    let file_len = reader.size().await?;

    // 1. Validate the StreamHeader to confirm this is actually an RRD file.
    if file_len < StreamHeader::ENCODED_SIZE_BYTES as u64 {
        return Err(CodecError::FrameDecoding(
            "file too small to be an RRD".to_owned(),
        ));
    }
    let header_buf = reader
        .read_exact_at(0, StreamHeader::ENCODED_SIZE_BYTES)
        .await?;
    StreamHeader::from_rrd_bytes(&header_buf)?; // validates FourCC + version

    // 2. Read the StreamFooter from the end of the file.
    if file_len < StreamFooter::ENCODED_SIZE_BYTES as u64 {
        return Ok(None); // File too small to have a footer.
    }
    // SAFETY: The preceding size check prevents the subtraction from underflowing.
    let footer_buf = reader
        .read_exact_at(
            file_len - StreamFooter::ENCODED_SIZE_BYTES as u64,
            StreamFooter::ENCODED_SIZE_BYTES,
        )
        .await?;

    let Ok(stream_footer) = StreamFooter::from_rrd_bytes(&footer_buf) else {
        return Ok(None); // Valid RRD, but no footer (legacy).
    };

    // 2. For each entry, read and validate the RrdFooter payload.
    //    In practice there is always exactly one entry.
    let Some(entry) = stream_footer.entries.first() else {
        return Ok(None);
    };

    let span = &entry.rrd_footer_byte_span_from_start_excluding_header;
    let payload_len = usize::try_from(span.len)?;

    // Sanity check: payload must fit within the file.
    if span.start + span.len > file_len {
        return Err(CodecError::FrameDecoding(format!(
            "RrdFooter payload span ({start}..{end}) exceeds file size ({file_len})",
            start = span.start,
            end = span.start + span.len,
        )));
    }

    // 3. Read the RrdFooter payload.
    let payload_buf = reader.read_exact_at(span.start, payload_len).await?;

    // 4. Validate CRC.
    let actual_crc = StreamFooter::compute_crc(&payload_buf);
    if actual_crc != entry.crc_excluding_header {
        return Err(CodecError::CrcMismatch {
            expected: entry.crc_excluding_header,
            got: actual_crc,
        });
    }

    Ok(Some(payload_buf))
}

/// Enumerate all [`StoreId`]s present in an RRD file, without reading chunk data.
///
/// - **With footer** (modern RRDs): reads the footer and returns the keys of its manifests map.
///   Cheap: a few small reads, no chunk data read. All stores are visible regardless of message ordering.
/// - **Without footer** (legacy RRDs): walks message frames, decoding only `SetStoreInfo`
///   payloads and skipping past everything else. All `SetStoreInfo`s are discovered regardless
///   of how they interleave with `ArrowMsg`s.
///
/// The returned list is sorted by [`StoreId`]'s natural order for determinism.
pub async fn enumerate_rrd_stores<R: AsyncReadAt>(reader: &R) -> Result<Vec<StoreId>, CodecError> {
    // Try footer first (cheap: a few small reads, no chunk data read).
    if let Some(footer) = read_rrd_footer(reader).await? {
        let mut store_ids: Vec<StoreId> = footer.manifests.into_keys().collect();
        store_ids.sort();
        return Ok(store_ids);
    }

    enumerate_legacy_metadata(reader)
        .await
        .map(|metadata| metadata.store_ids)
}

/// Metadata collected from RRD message frames without decoding chunk data.
pub struct RrdMetadata {
    /// All store ids announced by `SetStoreInfo` messages, sorted and deduplicated.
    pub store_ids: Vec<StoreId>,

    /// The last default-blueprint activation for each application id.
    pub default_blueprint_by_app_id: BTreeMap<ApplicationId, StoreId>,
}

/// Enumerate the metadata available through legacy RRD message frames.
///
/// Reads each 16-byte [`MessageHeader`], decodes only `SetStoreInfo` and
/// `BlueprintActivationCommand` payloads, and skips past `ArrowMsg` payloads. No Arrow IPC decoding
/// happens.
///
/// This does not inspect the footer. Store ids therefore come only from `SetStoreInfo` messages and
/// may be incomplete for footer-backed RRDs. The same `StoreId` can appear in multiple
/// `SetStoreInfo` messages, so the returned list is sorted and deduplicated.
pub async fn enumerate_legacy_metadata<R: AsyncReadAt>(
    reader: &R,
) -> Result<RrdMetadata, CodecError> {
    // TODO(grtlr): Profile this scan without holding a Puffin scope across `.await` points.
    // `ProfilerScope` must end on its starting thread and therefore makes the future non-`Send`.

    // Read and validate the StreamHeader; it also carries the crate version we need when
    // decoding individual message payloads (for version-dependent migrations).
    let stream_header_buf = reader
        .read_exact_at(0, StreamHeader::ENCODED_SIZE_BYTES)
        .await?;
    let stream_header = StreamHeader::from_rrd_bytes(&stream_header_buf)?;
    let (version, _options) = stream_header.to_version_and_options()?;

    let mut store_ids = Vec::new();
    let mut default_blueprint_by_app_id = BTreeMap::new();
    let mut app_id_cache = CachingApplicationIdInjector::default();

    let mut offset = StreamHeader::ENCODED_SIZE_BYTES as u64;
    loop {
        let msg_header_buf = match reader
            .read_exact_at(offset, MessageHeader::ENCODED_SIZE_BYTES)
            .await
        {
            Ok(buf) => buf,
            Err(err) if err.kind() == std::io::ErrorKind::UnexpectedEof => break,
            Err(err) => return Err(CodecError::Io(err)),
        };
        offset += MessageHeader::ENCODED_SIZE_BYTES as u64;
        let header = MessageHeader::from_rrd_bytes(&msg_header_buf)?;

        match header.kind {
            MessageKind::End => break,

            MessageKind::SetStoreInfo | MessageKind::BlueprintActivationCommand => {
                let payload_len = usize::try_from(header.len).map_err(CodecError::Overflow)?;
                let payload = reader.read_exact_at(offset, payload_len).await?;
                offset += header.len;
                let byte_span = re_chunk::Span::from_start_len(0, header.len);
                match LogMsg::decode(
                    payload,
                    byte_span,
                    header.kind,
                    &mut app_id_cache,
                    Some(version),
                )? {
                    Some(LogMsg::SetStoreInfo(set_store_info)) => {
                        store_ids.push(set_store_info.info.store_id);
                    }
                    Some(LogMsg::BlueprintActivationCommand(command)) if command.make_default => {
                        default_blueprint_by_app_id.insert(
                            command.blueprint_id.application_id().clone(),
                            command.blueprint_id,
                        );
                    }
                    _ => {}
                }
            }

            MessageKind::ArrowMsg => {
                offset += header.len;
            }
        }
    }

    store_ids.sort();
    store_ids.dedup();

    // From a server-perspective, it does not make sense to keep
    // former blueprint activation commands, if they are not part
    // of this RRD.
    default_blueprint_by_app_id.retain(|_, store_id| store_ids.contains(store_id));

    Ok(RrdMetadata {
        store_ids,
        default_blueprint_by_app_id,
    })
}

#[cfg(test)]
#[cfg(not(target_arch = "wasm32"))]
mod tests {
    use std::fs::File;

    use super::*;
    use crate::rrd::test_util::{encode_test_rrd, encode_test_rrd_to_file, make_test_chunks};

    fn set_store_info(store_id: StoreId) -> LogMsg {
        LogMsg::SetStoreInfo(re_log_types::SetStoreInfo {
            row_id: *re_chunk::RowId::ZERO,
            info: re_log_types::StoreInfo::new(store_id, re_log_types::StoreSource::Unknown),
        })
    }

    #[test]
    fn test_read_footer_roundtrip() {
        let chunks = make_test_chunks(5);
        let (file, _store_id) = encode_test_rrd(&chunks);

        let file = File::open(file.path()).unwrap();
        let footer = futures::executor::block_on(read_rrd_footer(&file)).unwrap();
        assert!(footer.is_some(), "Footer should be present");
        let footer = footer.unwrap();
        assert!(
            !footer.manifests.is_empty(),
            "Should have at least one manifest"
        );
    }

    #[test]
    fn test_read_footer_no_footer() {
        // Needs with_footer: false, so uses the lower-level helper.
        let file = tempfile::NamedTempFile::new().unwrap();
        let chunks = make_test_chunks(3);
        encode_test_rrd_to_file(file.path(), &chunks, false);

        let file = File::open(file.path()).unwrap();
        let footer = futures::executor::block_on(read_rrd_footer(&file)).unwrap();
        assert!(footer.is_none(), "Legacy RRD should have no footer");
    }

    #[test]
    fn test_enumerate_rrd_stores_footer_path() {
        let chunks = make_test_chunks(5);
        let (file, store_id) = encode_test_rrd(&chunks);

        let file = File::open(file.path()).unwrap();
        let ids = futures::executor::block_on(enumerate_rrd_stores(&file)).unwrap();
        assert_eq!(ids, vec![store_id]);
    }

    #[test]
    fn test_enumerate_rrd_stores_legacy_path() {
        let file = tempfile::NamedTempFile::new().unwrap();
        let chunks = make_test_chunks(3);
        encode_test_rrd_to_file(file.path(), &chunks, false);

        let file = File::open(file.path()).unwrap();
        let ids = futures::executor::block_on(enumerate_rrd_stores(&file)).unwrap();
        assert_eq!(
            ids.len(),
            1,
            "Legacy RRD should have its single store discovered"
        );
    }

    /// Cross-check: the legacy (frame-scan) and modern (footer) paths must return identical
    /// results on the same logical RRD content.
    #[test]
    fn test_enumerate_rrd_stores_footer_vs_legacy_agree() {
        let chunks = make_test_chunks(5);
        let store_id =
            re_log_types::StoreId::random(re_log_types::StoreKind::Recording, "cross_check");

        let with_footer = tempfile::NamedTempFile::new().unwrap();
        let without_footer = tempfile::NamedTempFile::new().unwrap();
        crate::rrd::test_util::encode_test_rrd_to_file_with_options(
            with_footer.path(),
            &chunks,
            &store_id,
            true,
            crate::EncodingOptions::PROTOBUF_COMPRESSED,
        );
        crate::rrd::test_util::encode_test_rrd_to_file_with_options(
            without_footer.path(),
            &chunks,
            &store_id,
            false,
            crate::EncodingOptions::PROTOBUF_COMPRESSED,
        );

        let with_footer_file = File::open(with_footer.path()).unwrap();
        let ids_footer =
            futures::executor::block_on(enumerate_rrd_stores(&with_footer_file)).unwrap();
        let without_footer_file = File::open(without_footer.path()).unwrap();
        let ids_legacy =
            futures::executor::block_on(enumerate_rrd_stores(&without_footer_file)).unwrap();

        assert_eq!(ids_footer, vec![store_id]);
        assert_eq!(
            ids_footer, ids_legacy,
            "footer path and legacy path must agree on the same content"
        );
    }

    /// Legacy RRD where the *same* `StoreId` is announced by several `SetStoreInfo`
    /// messages (can happen e.g. after a flush/reconnect). Enumeration must dedup.
    #[test]
    fn test_enumerate_rrd_stores_legacy_duplicate_set_store_info() {
        use re_log_types::{LogMsg, SetStoreInfo, StoreId, StoreInfo, StoreKind, StoreSource};

        let chunks = make_test_chunks(2);
        let store_id = StoreId::random(StoreKind::Recording, "dup_test");

        let file = tempfile::NamedTempFile::new().unwrap();
        {
            let mut out = std::fs::File::create(file.path()).unwrap();
            let mut encoder = crate::Encoder::new_eager(
                re_build_info::CrateVersion::LOCAL,
                crate::EncodingOptions::PROTOBUF_COMPRESSED,
                &mut out,
            )
            .unwrap();
            encoder.do_not_emit_footer();

            let info = StoreInfo::new(store_id.clone(), StoreSource::Unknown);
            // Three SetStoreInfos for the same store, interleaved with chunks.
            encoder
                .append(&LogMsg::SetStoreInfo(SetStoreInfo {
                    row_id: *re_chunk::RowId::ZERO,
                    info: info.clone(),
                }))
                .unwrap();
            let arrow_msg_0 = chunks[0].to_arrow_msg().unwrap();
            encoder
                .append(&LogMsg::ArrowMsg(store_id.clone(), arrow_msg_0))
                .unwrap();
            encoder
                .append(&LogMsg::SetStoreInfo(SetStoreInfo {
                    row_id: *re_chunk::RowId::ZERO,
                    info: info.clone(),
                }))
                .unwrap();
            let arrow_msg_1 = chunks[1].to_arrow_msg().unwrap();
            encoder
                .append(&LogMsg::ArrowMsg(store_id.clone(), arrow_msg_1))
                .unwrap();
            encoder
                .append(&LogMsg::SetStoreInfo(SetStoreInfo {
                    row_id: *re_chunk::RowId::ZERO,
                    info,
                }))
                .unwrap();
            encoder.finish().unwrap();
        }

        let file = File::open(file.path()).unwrap();
        let ids = futures::executor::block_on(enumerate_rrd_stores(&file)).unwrap();
        assert_eq!(
            ids,
            vec![store_id],
            "duplicate SetStoreInfos for the same StoreId must be deduped"
        );
    }

    /// Two-store RRD with `SetStoreInfo` messages interleaved with `ArrowMsg`s. Both the
    /// legacy (no-footer) path and the modern (footer) path must discover both stores and
    /// return the same result.
    #[test]
    fn test_enumerate_rrd_stores_interleaved_footer_vs_legacy() {
        use re_log_types::{LogMsg, SetStoreInfo, StoreId, StoreInfo, StoreKind, StoreSource};

        let chunks_a = make_test_chunks(2);
        let chunks_b = make_test_chunks(2);
        let store_a = StoreId::random(StoreKind::Recording, "test_a");
        let store_b = StoreId::random(StoreKind::Recording, "test_b");

        // Writes interleaved content: SetStoreInfo(A), chunks(A), SetStoreInfo(B), chunks(B).
        let write_interleaved = |path: &std::path::Path, with_footer: bool| {
            let mut out = std::fs::File::create(path).unwrap();
            let mut encoder = crate::Encoder::new_eager(
                re_build_info::CrateVersion::LOCAL,
                crate::EncodingOptions::PROTOBUF_COMPRESSED,
                &mut out,
            )
            .unwrap();
            if !with_footer {
                encoder.do_not_emit_footer();
            }

            let info_a = StoreInfo::new(store_a.clone(), StoreSource::Unknown);
            encoder
                .append(&LogMsg::SetStoreInfo(SetStoreInfo {
                    row_id: *re_chunk::RowId::ZERO,
                    info: info_a,
                }))
                .unwrap();
            for chunk in &chunks_a {
                let arrow_msg = chunk.to_arrow_msg().unwrap();
                encoder
                    .append(&LogMsg::ArrowMsg(store_a.clone(), arrow_msg))
                    .unwrap();
            }
            let info_b = StoreInfo::new(store_b.clone(), StoreSource::Unknown);
            encoder
                .append(&LogMsg::SetStoreInfo(SetStoreInfo {
                    row_id: *re_chunk::RowId::ZERO,
                    info: info_b,
                }))
                .unwrap();
            for chunk in &chunks_b {
                let arrow_msg = chunk.to_arrow_msg().unwrap();
                encoder
                    .append(&LogMsg::ArrowMsg(store_b.clone(), arrow_msg))
                    .unwrap();
            }
            encoder.finish().unwrap();
        };

        let with_footer = tempfile::NamedTempFile::new().unwrap();
        let without_footer = tempfile::NamedTempFile::new().unwrap();
        write_interleaved(with_footer.path(), true);
        write_interleaved(without_footer.path(), false);

        let with_footer_file = File::open(with_footer.path()).unwrap();
        let ids_footer =
            futures::executor::block_on(enumerate_rrd_stores(&with_footer_file)).unwrap();
        let without_footer_file = File::open(without_footer.path()).unwrap();
        let ids_legacy =
            futures::executor::block_on(enumerate_rrd_stores(&without_footer_file)).unwrap();

        let mut expected = vec![store_a, store_b];
        expected.sort();
        assert_eq!(ids_footer, expected, "footer path must find both stores");
        assert_eq!(
            ids_legacy, expected,
            "legacy path must find both stores despite interleaving"
        );
        assert_eq!(
            ids_footer, ids_legacy,
            "footer and legacy paths must agree on the same content"
        );
    }

    #[test]
    fn test_enumerate_legacy_metadata_default_blueprint_last_wins() {
        use re_log_types::{BlueprintActivationCommand, LogMsg, StoreId, StoreKind};

        let app_id = re_log_types::ApplicationId::from("metadata_app");
        let recording = StoreId::random(StoreKind::Recording, app_id.clone());
        let blueprint_a = StoreId::random(StoreKind::Blueprint, app_id.clone());
        let blueprint_b = StoreId::random(StoreKind::Blueprint, app_id);
        let file = tempfile::NamedTempFile::new().unwrap();
        {
            let mut out = std::fs::File::create(file.path()).unwrap();
            let mut encoder = crate::Encoder::new_eager(
                re_build_info::CrateVersion::LOCAL,
                crate::EncodingOptions::PROTOBUF_COMPRESSED,
                &mut out,
            )
            .unwrap();
            encoder.append(&set_store_info(recording.clone())).unwrap();
            encoder
                .append(&set_store_info(blueprint_a.clone()))
                .unwrap();
            encoder
                .append(&LogMsg::BlueprintActivationCommand(
                    BlueprintActivationCommand::make_default(blueprint_a),
                ))
                .unwrap();
            encoder
                .append(&set_store_info(blueprint_b.clone()))
                .unwrap();
            encoder
                .append(&LogMsg::BlueprintActivationCommand(
                    BlueprintActivationCommand::make_default(blueprint_b.clone()),
                ))
                .unwrap();
            encoder.finish().unwrap();
        }

        let file = File::open(file.path()).unwrap();
        let default_blueprints = futures::executor::block_on(enumerate_legacy_metadata(&file))
            .unwrap()
            .default_blueprint_by_app_id;
        assert_eq!(default_blueprints.len(), 1);
        assert_eq!(
            default_blueprints.get(recording.application_id()),
            Some(&blueprint_b)
        );
    }

    #[test]
    fn test_enumerate_legacy_metadata_ignores_non_default_blueprint_activation() {
        use re_log_types::{BlueprintActivationCommand, LogMsg, StoreId, StoreKind};

        let app_id = re_log_types::ApplicationId::from("metadata_app");
        let recording = StoreId::random(StoreKind::Recording, app_id.clone());
        let blueprint = StoreId::random(StoreKind::Blueprint, app_id);
        let file = tempfile::NamedTempFile::new().unwrap();
        {
            let mut out = std::fs::File::create(file.path()).unwrap();
            let mut encoder = crate::Encoder::new_eager(
                re_build_info::CrateVersion::LOCAL,
                crate::EncodingOptions::PROTOBUF_COMPRESSED,
                &mut out,
            )
            .unwrap();
            encoder.append(&set_store_info(recording)).unwrap();
            encoder.append(&set_store_info(blueprint.clone())).unwrap();
            encoder
                .append(&LogMsg::BlueprintActivationCommand(
                    BlueprintActivationCommand {
                        blueprint_id: blueprint,
                        make_active: false,
                        make_default: false,
                    },
                ))
                .unwrap();
            encoder.finish().unwrap();
        }

        let file = File::open(file.path()).unwrap();
        let default_blueprints = futures::executor::block_on(enumerate_legacy_metadata(&file))
            .unwrap()
            .default_blueprint_by_app_id;
        assert!(default_blueprints.is_empty());
    }

    #[test]
    fn test_enumerate_legacy_metadata_omits_missing_default_blueprint() {
        use re_log_types::{BlueprintActivationCommand, LogMsg, StoreId, StoreKind};

        let app_id = re_log_types::ApplicationId::from("metadata_app");
        let recording = StoreId::random(StoreKind::Recording, app_id.clone());
        let dangling_blueprint = StoreId::random(StoreKind::Blueprint, app_id);
        let file = tempfile::NamedTempFile::new().unwrap();
        {
            let mut out = std::fs::File::create(file.path()).unwrap();
            let mut encoder = crate::Encoder::new_eager(
                re_build_info::CrateVersion::LOCAL,
                crate::EncodingOptions::PROTOBUF_COMPRESSED,
                &mut out,
            )
            .unwrap();
            encoder.append(&set_store_info(recording.clone())).unwrap();
            encoder
                .append(&LogMsg::BlueprintActivationCommand(
                    BlueprintActivationCommand::make_default(dangling_blueprint.clone()),
                ))
                .unwrap();
            encoder.finish().unwrap();
        }

        let file = File::open(file.path()).unwrap();
        let default_blueprints = futures::executor::block_on(enumerate_legacy_metadata(&file))
            .unwrap()
            .default_blueprint_by_app_id;
        assert!(!default_blueprints.contains_key(recording.application_id()));
    }

    #[test]
    fn test_read_footer_not_an_rrd() {
        let file = tempfile::NamedTempFile::new().unwrap();
        std::fs::write(file.path(), b"this is not an rrd file at all").unwrap();

        let file = File::open(file.path()).unwrap();
        let result = futures::executor::block_on(read_rrd_footer(&file));
        assert!(result.is_err(), "Non-RRD file should return an error");
    }

    #[test]
    fn test_read_footer_too_small() {
        let file = tempfile::NamedTempFile::new().unwrap();
        std::fs::write(file.path(), b"tiny").unwrap();

        let file = File::open(file.path()).unwrap();
        let result = futures::executor::block_on(read_rrd_footer(&file));
        assert!(
            result.is_err(),
            "File too small for StreamHeader should error"
        );
    }

    #[test]
    fn test_read_footer_corrupted_crc() {
        let chunks = make_test_chunks(3);
        let (file, _store_id) = encode_test_rrd(&chunks);

        let mut data = std::fs::read(file.path()).unwrap();
        let file_len = data.len();

        let footer_bytes = &data[file_len - StreamFooter::ENCODED_SIZE_BYTES..];
        let stream_footer = StreamFooter::from_rrd_bytes(footer_bytes).unwrap();
        let entry = &stream_footer.entries[0];
        let payload_start = entry.rrd_footer_byte_span_from_start_excluding_header.start as usize;

        // Flip a byte in the payload.
        data[payload_start] ^= 0xFF;
        std::fs::write(file.path(), &data).unwrap();

        let file = File::open(file.path()).unwrap();
        let result = futures::executor::block_on(read_rrd_footer(&file));
        assert!(
            matches!(result, Err(CodecError::CrcMismatch { .. })),
            "Expected CRC mismatch, got: {result:?}"
        );
    }
}