re_view_spatial 0.37.1

use std::sync::Arc;

use re_chunk_store::ChunkTrackingMode;
use re_log_types::EntityPath;
use re_renderer::external::re_video::VideoLoadError;
use re_renderer::video::Video;
use re_sdk_types::Archetype as _;
use re_sdk_types::archetypes::{AssetVideo, VideoFrameReference, VideoStream};
use re_sdk_types::components::{Blob, MediaType, Opacity, VideoTimestamp};
use re_video::player::VideoSliceSource;
use re_viewer_context::{
    IdentifiedViewSystem, SharablePlayableVideoStream, VideoAssetCache, VideoStoreSource,
    VideoStreamCache, VideoStreamProcessingError, ViewClass as _, ViewContext,
    ViewContextCollection, ViewQuery, ViewSystemExecutionError, ViewerContext,
    VisualizerExecutionOutput, VisualizerQueryInfo, VisualizerSystem, typed_fallback_for,
};

use crate::PickableTexturedRect;
use crate::contexts::SpatialSceneVisualizerInstructionContext;
use crate::visualizers::SpatialViewVisualizerData;
use crate::visualizers::entity_iterator::process_archetype;
use crate::visualizers::video::{
    AT_TIME_CURSOR_SALT, VideoFrameRenderInfo, VideoPlaybackIssue, VideoPlaybackIssueSeverity,
    show_video_frame, video_stream_id, video_stream_processing_issue,
};

#[derive(Default)]
pub struct VideoFrameReferenceVisualizer;

impl IdentifiedViewSystem for VideoFrameReferenceVisualizer {
    fn identifier() -> re_viewer_context::ViewSystemIdentifier {
        re_viewer_context::external::re_string_interner::intern_static!(
            re_viewer_context::ViewSystemIdentifier,
            "VideoFrameReference"
        )
    }
}

impl VisualizerSystem for VideoFrameReferenceVisualizer {
    fn visualizer_query_info(
        &self,
        _app_options: &re_viewer_context::AppOptions,
    ) -> VisualizerQueryInfo {
        VisualizerQueryInfo::single_required_component::<VideoTimestamp>(
            &VideoFrameReference::descriptor_timestamp(),
            &VideoFrameReference::all_components(),
        )
    }

    fn affinity(&self) -> Option<re_sdk_types::ViewClassIdentifier> {
        Some(crate::SpatialView2D::identifier())
    }

    fn execute(
        &self,
        ctx: &ViewContext<'_>,
        view_query: &ViewQuery<'_>,
        context_systems: &ViewContextCollection,
    ) -> Result<VisualizerExecutionOutput, ViewSystemExecutionError> {
        re_tracing::profile_function!();

        let mut data = SpatialViewVisualizerData::default();
        let output = VisualizerExecutionOutput::default();

        process_archetype::<VideoFrameReference, _, _>(
            ctx,
            view_query,
            context_systems,
            &output,
            self,
            |ctx, spatial_ctx, results| {
                // TODO(andreas): Should ignore range queries here and only do latest-at.
                // Not only would this simplify the code here quite a bit, it would also avoid lots of overhead.
                // Same is true for the image visualizers in general - there seems to be no practical reason to do range queries
                // for visualization here.

                let entity_path = ctx.target_entity_path;

                let all_video_timestamps =
                    results.iter_required(VideoFrameReference::descriptor_timestamp().component);
                if all_video_timestamps.is_empty() {
                    return Ok(());
                }
                let all_video_references = results
                    .iter_optional(VideoFrameReference::descriptor_video_reference().component);
                let all_opacities =
                    results.iter_optional(VideoFrameReference::descriptor_opacity().component);

                for (_index, video_timestamps, video_references, opacity) in re_query::range_zip_1x2(
                    all_video_timestamps.component_slow(),
                    all_video_references.slice::<String>(),
                    all_opacities.slice::<f32>(),
                ) {
                    let Some(video_timestamp): Option<&VideoTimestamp> = video_timestamps.first()
                    else {
                        continue;
                    };

                    Self::process_video_frame(
                        &mut data,
                        ctx,
                        spatial_ctx,
                        video_timestamp,
                        video_references,
                        opacity
                            .and_then(|slice| slice.first())
                            .copied()
                            .map(Opacity::from)
                            .unwrap_or_else(|| {
                                typed_fallback_for(
                                    ctx,
                                    VideoFrameReference::descriptor_opacity().component,
                                )
                            }),
                        entity_path,
                    );
                }

                Ok(())
            },
        )?;

        Ok(output
            .with_draw_data([PickableTexturedRect::to_draw_data(
                ctx.viewer_ctx.render_ctx(),
                &data.pickable_rects,
            )?])
            .with_visualizer_data(data))
    }
}

impl VideoFrameReferenceVisualizer {
    fn process_video_frame(
        data: &mut SpatialViewVisualizerData,
        ctx: &re_viewer_context::QueryContext<'_>,
        spatial_ctx: &SpatialSceneVisualizerInstructionContext<'_>,
        video_timestamp: &VideoTimestamp,
        video_references: Option<Vec<re_sdk_types::ArrowString>>,
        opacity: Opacity,
        entity_path: &EntityPath,
    ) {
        re_tracing::profile_function!();

        let player_stream_id = video_stream_id(
            entity_path,
            VideoFrameReference::descriptor_video_reference().component,
            AT_TIME_CURSOR_SALT,
        );

        // The referenced entity can hold either an asset or a stream.
        let video_reference: EntityPath = video_references
            .and_then(|v| v.first().map(|e| e.as_str().into()))
            .unwrap_or_else(|| ctx.target_entity_path.clone());

        let world_from_entity = spatial_ctx
            .transform_info
            .single_transform_required_for_entity(
                ctx.target_entity_path,
                VideoFrameReference::name(),
            )
            .as_affine3a();

        // Note that we may or may not know the video size independently of error occurrence.
        // (if it's just a decoding error we may still know the size from the container!)
        // In case we haven error we want to center the message in the middle, so we need some area.
        // Note that this area is also used for the bounding box which is important for the 2D view to determine default bounds.
        let mut video_resolution = glam::vec2(1280.0, 720.0);

        let frame = if let Some((video, video_buffer)) =
            latest_at_query_video_from_datastore(ctx.viewer_ctx(), &video_reference)
        {
            match video.as_ref() {
                Ok(video) => {
                    if let Some([w, h]) = video.dimensions() {
                        video_resolution = glam::vec2(w as _, h as _);
                    }

                    let video_time = re_viewer_context::video_timestamp_component_to_video_time(
                        Some(ctx.viewer_ctx().time_ctrl),
                        *video_timestamp,
                        video.data_descr().timescale,
                    );

                    let frame_output = video.frame_at(
                        ctx.render_ctx(),
                        player_stream_id,
                        video_time,
                        &VideoSliceSource(&video_buffer),
                    );

                    let bit_depth = video
                        .data_descr()
                        .encoding_details
                        .as_ref()
                        .and_then(|d| d.bit_depth);
                    Ok((frame_output, bit_depth))
                }
                Err(err) => Err(VideoPlaybackIssue::custom(
                    err.to_string(),
                    VideoPlaybackIssueSeverity::Error,
                )),
            }
        } else {
            match latest_at_query_video_stream_from_datastore(
                ctx.viewer_ctx(),
                &video_reference,
                &ctx.query,
            ) {
                None => Err(VideoPlaybackIssue::custom(
                    format!("No video asset or stream at {video_reference:?}"),
                    VideoPlaybackIssueSeverity::Informational,
                )),
                Some(Err(err)) => Err(video_stream_processing_issue(&video_reference, &err)),
                Some(Ok(video)) => {
                    let sample_component = VideoStream::descriptor_sample().component;
                    let video = video.read();
                    let video_description = video.video_descr();
                    match video_description.samples.front() {
                        None => Err(video_stream_processing_issue(
                            &video_reference,
                            &VideoStreamProcessingError::NoVideoSamplesFound,
                        )),
                        Some(first_stream_timestamp) => {
                            if let Some([w, h]) = video.video_renderer.dimensions() {
                                video_resolution = glam::vec2(w as _, h as _);
                            }

                            let video_time = first_stream_timestamp.decode_timestamp()
                                + re_viewer_context::video_timestamp_component_to_video_time(
                                    Some(ctx.viewer_ctx().time_ctrl),
                                    *video_timestamp,
                                    video_description.timescale,
                                );
                            let storage_engine = ctx.recording().storage_engine();
                            let frame_output = video.video_renderer.frame_at(
                                ctx.render_ctx(),
                                player_stream_id,
                                video_time,
                                &VideoStoreSource {
                                    store: storage_engine.store(),
                                    sample_component,
                                    indicate: true,
                                },
                            );
                            let bit_depth = video_description
                                .encoding_details
                                .as_ref()
                                .and_then(|details| details.bit_depth);
                            Ok((frame_output, bit_depth))
                        }
                    }
                }
            }
        };
        let (frame_output, bit_depth) = match frame {
            Ok(frame) => frame,
            Err(issue) => {
                show_video_frame(
                    ctx.view_ctx,
                    data,
                    entity_path,
                    world_from_entity,
                    spatial_ctx.highlight,
                    video_resolution,
                    spatial_ctx.visualizer_instruction,
                    None,
                    Some(issue),
                    None,
                    None,
                );
                return;
            }
        };

        #[expect(clippy::disallowed_methods)] // This is not a hard-coded color.
        let multiplicative_tint = re_renderer::Rgba::from_white_alpha(opacity.0.clamp(0.0, 1.0));

        show_video_frame(
            ctx.view_ctx,
            data,
            entity_path,
            world_from_entity,
            spatial_ctx.highlight,
            video_resolution,
            spatial_ctx.visualizer_instruction,
            frame_output.output.map(|texture| VideoFrameRenderInfo {
                texture,
                depth_offset: spatial_ctx.depth_offset,
                multiplicative_tint,
            }),
            frame_output.error.map(VideoPlaybackIssue::from),
            None,
            bit_depth,
        );
    }
}

/// Queries a video from the datstore and caches it in the video cache.
///
/// Note that this does *NOT* check the blueprint store at all.
/// For this, we'd need a [`re_viewer_context::DataResult`] instead of merely a [`EntityPath`].
///
/// Returns `None` if there was no blob at the referenced path.
/// Returns `Some(Err(_))` if there was a blob but it failed to load for some reason.
/// Errors are cached as well so loading a failed video won't occur a high cost repeatedly.
fn latest_at_query_video_from_datastore(
    ctx: &ViewerContext<'_>,
    entity_path: &EntityPath,
) -> Option<(Arc<Result<Video, VideoLoadError>>, Blob)> {
    let query = ctx.current_query();

    let results = ctx.recording_engine().cache().latest_at(
        ChunkTrackingMode::Report,
        &query,
        entity_path,
        AssetVideo::all_component_identifiers(),
    );

    let blob_row_id = results.component_row_id(AssetVideo::descriptor_blob().component)?;
    let blob = results.component_instance::<Blob>(0, AssetVideo::descriptor_blob().component)?;
    let media_type =
        results.component_instance::<MediaType>(0, AssetVideo::descriptor_media_type().component);

    let video = ctx.store_context.memoizer(|c: &mut VideoAssetCache| {
        let debug_name = entity_path.to_string();
        c.entry(
            debug_name,
            blob_row_id,
            AssetVideo::descriptor_blob().component,
            &blob,
            media_type.as_ref(),
            ctx.app_options().video_decoder_settings(),
        )
    });
    Some((video, blob))
}

/// Queries a video stream from the datastore and caches it in the video stream cache.
///
/// Returns `None` if there was no video stream at the referenced path.
/// Returns `Some(Err(_))` if there was a video stream but it failed to load.
fn latest_at_query_video_stream_from_datastore(
    ctx: &ViewerContext<'_>,
    entity_path: &EntityPath,
    query: &re_chunk_store::LatestAtQuery,
) -> Option<Result<SharablePlayableVideoStream, VideoStreamProcessingError>> {
    let timeline = query.timeline()?;
    let codec_component = VideoStream::descriptor_codec().component;
    let results = ctx.recording_engine().cache().latest_at(
        ChunkTrackingMode::Report,
        query,
        entity_path,
        [codec_component],
    );
    results.component_row_id(codec_component)?;

    Some(ctx.store_context.memoizer(|cache: &mut VideoStreamCache| {
        cache.video_entry(
            ctx.recording(),
            entity_path,
            timeline,
            ctx.app_options().video_decoder_settings(),
            ChunkTrackingMode::Report,
        )
    }))
}