mecha10-behavior-runtime 0.8.0

//! Action node that captures a picture of a uniquely-found target and publishes it for the
//! platform API to pick up and store for user review. Watches a vision detections topic and
//! treats a target as found only when exactly one detection is present.

use crate::{BehaviorNode, Context, NodeStatus};
use async_trait::async_trait;
use mecha10_core::messages::{Detection, Image, Message};
use mecha10_core::topics::Topic;
use mecha10_core::Receiver;
use serde::{Deserialize, Serialize};
use serde_json::Value;

fn default_image_topic() -> String {
    "/camera/rgb".to_string()
}

fn default_detections_topic() -> String {
    "vision/detections".to_string()
}

fn default_capture_topic() -> String {
    "/capture-target/captures".to_string()
}

/// Published once per newly-unique target - the payload a downstream service relays to the
/// platform API for user review.
#[derive(Debug, Clone, Serialize, Deserialize)]
struct CaptureTargetRequest {
    image: Image,
    detection: Detection,
}

impl Message for CaptureTargetRequest {}

pub struct CaptureTargetNode {
    image_topic: String,
    detections_topic: String,
    capture_topic: String,
    image_receiver: Option<Receiver<Value>>,
    detections_receiver: Option<Receiver<Value>>,
    latest_image: Option<Image>,
    target: Option<Detection>,
    captured_for_current_target: bool,
}

impl std::fmt::Debug for CaptureTargetNode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CaptureTargetNode")
            .field("image_topic", &self.image_topic)
            .field("detections_topic", &self.detections_topic)
            .field("capture_topic", &self.capture_topic)
            .field("has_image", &self.latest_image.is_some())
            .field("target", &self.target)
            .field("captured_for_current_target", &self.captured_for_current_target)
            .finish()
    }
}

#[derive(Debug, Serialize, Deserialize)]
struct CaptureTargetConfig {
    #[serde(default = "default_image_topic")]
    image_topic: String,
    #[serde(default = "default_detections_topic")]
    detections_topic: String,
    #[serde(default = "default_capture_topic")]
    capture_topic: String,
}

impl CaptureTargetNode {
    /// Create a new CaptureTargetNode from JSON configuration. See `SetRlTargetNode::from_json`'s
    /// doc comment for why `null` is normalized to `{}` here.
    pub fn from_json(config: Value) -> anyhow::Result<Self> {
        let config = if config.is_null() {
            serde_json::json!({})
        } else {
            config
        };
        let cfg: CaptureTargetConfig = serde_json::from_value(config)?;
        Ok(Self {
            image_topic: cfg.image_topic,
            detections_topic: cfg.detections_topic,
            capture_topic: cfg.capture_topic,
            image_receiver: None,
            detections_receiver: None,
            latest_image: None,
            target: None,
            captured_for_current_target: false,
        })
    }

    /// Drains every pending frame off `receiver`, keeping only the most recent - mirrors
    /// `SetInstanceTrackTargetNode::poll_latest_request`'s "only the latest matters" semantics,
    /// applied to camera frames instead of target requests.
    fn poll_latest_image(receiver: &mut Receiver<Value>) -> Option<Image> {
        let mut latest = None;
        while let Ok(value) = receiver.try_recv() {
            match serde_json::from_value::<Image>(value) {
                Ok(image) => latest = Some(image),
                Err(e) => tracing::warn!("CaptureTargetNode: dropping malformed camera frame: {}", e),
            }
        }
        latest
    }

    /// Drains every pending detections batch off `receiver`, returning the last one seen (if
    /// any) - only the most recent batch matters, same "drain then keep latest" shape as the
    /// image poll above.
    fn poll_latest_detections(receiver: &mut Receiver<Value>) -> Option<Vec<Detection>> {
        let mut latest = None;
        while let Ok(value) = receiver.try_recv() {
            match serde_json::from_value::<Vec<Detection>>(value) {
                Ok(detections) => latest = Some(detections),
                Err(e) => tracing::warn!("CaptureTargetNode: dropping malformed detections batch: {}", e),
            }
        }
        latest
    }

    async fn publish_capture(&self, ctx: &Context, image: Image, detection: Detection) -> anyhow::Result<()> {
        let request = CaptureTargetRequest { image, detection };

        let topic_static: &'static str = Box::leak(self.capture_topic.clone().into_boxed_str());
        let topic = Topic::<CaptureTargetRequest>::new(topic_static);

        ctx.publish_to(topic, &request).await?;
        tracing::debug!(
            "CaptureTargetNode: published capture for class {:?}",
            request.detection.class_name
        );
        Ok(())
    }
}

#[async_trait]
impl BehaviorNode for CaptureTargetNode {
    async fn tick(&mut self, ctx: &Context) -> anyhow::Result<NodeStatus> {
        if self.image_receiver.is_none() {
            self.image_receiver = Some(ctx.subscribe_raw::<Value>(&self.image_topic).await?);
        }
        if self.detections_receiver.is_none() {
            self.detections_receiver = Some(ctx.subscribe_raw::<Value>(&self.detections_topic).await?);
        }

        if let Some(image) = Self::poll_latest_image(self.image_receiver.as_mut().expect("just subscribed above")) {
            self.latest_image = Some(image);
        }

        if let Some(detections) =
            Self::poll_latest_detections(self.detections_receiver.as_mut().expect("just subscribed above"))
        {
            let unique = match detections.len() {
                1 => detections.into_iter().next(),
                _ => None,
            };
            if unique.is_none() {
                self.captured_for_current_target = false;
            }
            self.target = unique;
        }

        if let Some(detection) = self.target.clone() {
            if !self.captured_for_current_target {
                if let Some(image) = self.latest_image.clone() {
                    self.publish_capture(ctx, image, detection).await?;
                    self.captured_for_current_target = true;
                }
            }
            Ok(NodeStatus::Running)
        } else {
            Ok(NodeStatus::Failure)
        }
    }

    fn name(&self) -> &str {
        "CaptureTargetNode"
    }

    fn reset_self(&mut self) {
        self.target = None;
        self.captured_for_current_target = false;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_helpers::create_test_context;

    fn sample_detection(class_name: &str) -> Value {
        serde_json::json!({
            "timestamp": 1,
            "class_id": 0,
            "class_name": class_name,
            "confidence": 0.9,
            "bbox": { "x": 0.1, "y": 0.2, "width": 0.3, "height": 0.4 }
        })
    }

    fn sample_image() -> Value {
        serde_json::json!({
            "timestamp": 1,
            "width": 4,
            "height": 4,
            "encoding": "rgb8",
            "data": [0, 0, 0]
        })
    }

    #[test]
    fn from_json_reads_all_topics() {
        let node = CaptureTargetNode::from_json(serde_json::json!({
            "image_topic": "/camera/rgb",
            "detections_topic": "vision/detections",
            "capture_topic": "/capture-target/captures"
        }))
        .unwrap();

        assert_eq!(node.image_topic, "/camera/rgb");
        assert_eq!(node.detections_topic, "vision/detections");
        assert_eq!(node.capture_topic, "/capture-target/captures");
    }

    #[test]
    fn from_json_defaults_all_topics() {
        let node = CaptureTargetNode::from_json(serde_json::json!({})).unwrap();

        assert_eq!(node.image_topic, "/camera/rgb");
        assert_eq!(node.detections_topic, "vision/detections");
        assert_eq!(node.capture_topic, "/capture-target/captures");
    }

    #[test]
    fn from_json_accepts_null_the_same_as_an_empty_object() {
        let node = CaptureTargetNode::from_json(Value::Null).unwrap();

        assert_eq!(node.image_topic, "/camera/rgb");
        assert_eq!(node.detections_topic, "vision/detections");
        assert_eq!(node.capture_topic, "/capture-target/captures");
    }

    #[test]
    fn poll_latest_image_returns_none_when_empty() {
        let (_tx, rx) = tokio::sync::mpsc::unbounded_channel::<Value>();
        let mut receiver = Receiver::from(rx);
        assert!(CaptureTargetNode::poll_latest_image(&mut receiver).is_none());
    }

    #[test]
    fn poll_latest_image_keeps_only_the_most_recent_frame() {
        let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Value>();
        let mut receiver = Receiver::from(rx);
        tx.send(serde_json::json!({ "timestamp": 1, "width": 1, "height": 1, "encoding": "rgb8", "data": [] }))
            .unwrap();
        tx.send(serde_json::json!({ "timestamp": 2, "width": 2, "height": 2, "encoding": "rgb8", "data": [] }))
            .unwrap();

        let image = CaptureTargetNode::poll_latest_image(&mut receiver).unwrap();
        assert_eq!(image.timestamp, 2);
        assert_eq!(image.width, 2);
    }

    #[test]
    fn poll_latest_image_drops_malformed_frames_and_keeps_polling() {
        let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Value>();
        let mut receiver = Receiver::from(rx);
        tx.send(serde_json::json!("not an image")).unwrap();
        tx.send(sample_image()).unwrap();

        let image = CaptureTargetNode::poll_latest_image(&mut receiver).unwrap();
        assert_eq!(image.width, 4);
    }

    #[test]
    fn poll_latest_detections_returns_none_when_empty() {
        let (_tx, rx) = tokio::sync::mpsc::unbounded_channel::<Value>();
        let mut receiver = Receiver::from(rx);
        assert!(CaptureTargetNode::poll_latest_detections(&mut receiver).is_none());
    }

    #[test]
    fn poll_latest_detections_keeps_only_the_most_recent_batch() {
        let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Value>();
        let mut receiver = Receiver::from(rx);
        tx.send(serde_json::json!([sample_detection("cat")])).unwrap();
        tx.send(serde_json::json!([sample_detection("dog"), sample_detection("cat")]))
            .unwrap();

        let detections = CaptureTargetNode::poll_latest_detections(&mut receiver).unwrap();
        assert_eq!(detections.len(), 2);
    }

    #[test]
    fn poll_latest_detections_drops_malformed_batches_and_keeps_polling() {
        let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<Value>();
        let mut receiver = Receiver::from(rx);
        tx.send(serde_json::json!("not a detections batch")).unwrap();
        tx.send(serde_json::json!([sample_detection("cat")])).unwrap();

        let detections = CaptureTargetNode::poll_latest_detections(&mut receiver).unwrap();
        assert_eq!(detections.len(), 1);
        assert_eq!(detections[0].class_name, "cat");
    }

    #[tokio::test]
    #[ignore]
    async fn tick_captures_once_when_exactly_one_detection_is_present() {
        let ctx = create_test_context("test_capture_target_unique").await;

        let mut node = CaptureTargetNode::from_json(serde_json::json!({
            "image_topic": "/camera/rgb",
            "detections_topic": "vision/detections",
            "capture_topic": "/capture-target/captures"
        }))
        .unwrap();

        let mut capture_receiver = ctx.subscribe_raw::<Value>("/capture-target/captures").await.unwrap();

        node.tick(&ctx).await.unwrap();
        ctx.publish_raw("/camera/rgb", &sample_image()).await.unwrap();
        ctx.publish_raw("vision/detections", &serde_json::json!([sample_detection("cat")]))
            .await
            .unwrap();
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;

        let status = node.tick(&ctx).await.unwrap();
        assert_eq!(status, NodeStatus::Running);

        let received = tokio::time::timeout(std::time::Duration::from_millis(500), capture_receiver.recv())
            .await
            .expect("timed out waiting for published capture")
            .expect("channel closed without a message");

        assert_eq!(received["detection"]["class_name"], "cat");
        assert_eq!(received["image"]["width"], 4);

        let status = node.tick(&ctx).await.unwrap();
        assert_eq!(status, NodeStatus::Running);
        let timed_out = tokio::time::timeout(std::time::Duration::from_millis(200), capture_receiver.recv()).await;
        assert!(
            timed_out.is_err(),
            "expected no second publish for the same unique target"
        );
    }

    #[tokio::test]
    #[ignore = "Requires Redis infrastructure"]
    async fn tick_returns_failure_with_zero_or_multiple_detections() {
        let ctx = create_test_context("test_capture_target_ambiguous").await;

        let mut node = CaptureTargetNode::from_json(serde_json::json!({})).unwrap();

        let status = node.tick(&ctx).await.unwrap();
        assert_eq!(
            status,
            NodeStatus::Failure,
            "no detections yet means nothing to capture"
        );

        ctx.publish_raw(
            "vision/detections",
            &serde_json::json!([sample_detection("cat"), sample_detection("dog")]),
        )
        .await
        .unwrap();
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;

        let status = node.tick(&ctx).await.unwrap();
        assert_eq!(
            status,
            NodeStatus::Failure,
            "ambiguous detections are not a unique target"
        );
    }

    #[tokio::test]
    #[ignore = "Requires Redis infrastructure"]
    async fn tick_recaptures_after_the_target_is_lost_and_refound() {
        let ctx = create_test_context("test_capture_target_refound").await;

        let mut node = CaptureTargetNode::from_json(serde_json::json!({})).unwrap();
        let mut capture_receiver = ctx.subscribe_raw::<Value>("/capture-target/captures").await.unwrap();

        node.tick(&ctx).await.unwrap();
        ctx.publish_raw("/camera/rgb", &sample_image()).await.unwrap();
        ctx.publish_raw("vision/detections", &serde_json::json!([sample_detection("cat")]))
            .await
            .unwrap();
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
        node.tick(&ctx).await.unwrap();
        tokio::time::timeout(std::time::Duration::from_millis(500), capture_receiver.recv())
            .await
            .expect("timed out waiting for first capture")
            .expect("channel closed without a message");

        ctx.publish_raw("vision/detections", &serde_json::json!([]))
            .await
            .unwrap();
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
        let status = node.tick(&ctx).await.unwrap();
        assert_eq!(status, NodeStatus::Failure, "target lost");

        ctx.publish_raw("vision/detections", &serde_json::json!([sample_detection("cat")]))
            .await
            .unwrap();
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
        node.tick(&ctx).await.unwrap();

        tokio::time::timeout(std::time::Duration::from_millis(500), capture_receiver.recv())
            .await
            .expect("timed out waiting for second capture after refind")
            .expect("channel closed without a message");
    }
}