From 115b8556bae304d461a9a8eb2d5a3f5c5d88661f Mon Sep 17 00:00:00 2001
From: William Emfinger <waemfinger@gmail.com>
Date: Fri, 15 Aug 2025 22:20:08 -0500
Subject: [PATCH] feat(rtsp): Update jpeg frame to better handle out of order
 segments

---
 components/rtsp/include/jpeg_frame.hpp | 60 +++++++++++++++++++++++++-
 components/rtsp/src/rtsp_server.cpp    |  2 +-
 2 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/components/rtsp/include/jpeg_frame.hpp b/components/rtsp/include/jpeg_frame.hpp
index 19c62cc19..4be0f496a 100644
--- a/components/rtsp/include/jpeg_frame.hpp
+++ b/components/rtsp/include/jpeg_frame.hpp
@@ -72,9 +72,19 @@ class JpegFrame {
   ///       finalized, and no further scans can be added.
   /// @param packet The packet containing the scan to append.
   void add_scan(const espp::RtpJpegPacket &packet) {
-    add_scan(packet.get_jpeg_data());
+    add_scan(packet.get_jpeg_data(), packet.get_offset());
+    auto sequence_number = packet.get_sequence_number();
+    received_sequence_numbers_.push_back(sequence_number);
     if (packet.get_marker()) {
-      finalize();
+      // check to ensure we got all the packets (sequence numbers) we expected
+      // based on start and this packet. What this means is that there should be
+      // no gaps in the received sequence numbers.
+      if (!has_missing_sequence_numbers()) {
+        finalize();
+      } else {
+        // fmt::print("JpegFrame: Missing sequence numbers in received packets {}\n",
+        //            received_sequence_numbers_);
+      }
     }
   }
 
@@ -106,6 +116,51 @@ class JpegFrame {
     data_.insert(std::end(data_), std::begin(scan), std::end(scan));
   }
 
+  /// Append a JPEG scan to the frame at a specific offset.
+  /// This will add the JPEG data to the frame at the specified offset.
+  /// @param scan The jpeg scan to append.
+  /// @param frag_offset The offset at which to append the scan.
+  /// @note If the offset is greater than the current data size, the data will
+  ///       be resized to accommodate the new data.
+  void add_scan(std::span<const uint8_t> scan, size_t frag_offset) {
+    if (finalized_) {
+      // TODO: handle this error
+      return;
+    }
+    // if the required span position is greater than the current data size, we
+    // need to pad the data with zeros.
+    size_t required_size = frag_offset + scan.size();
+    if (required_size >= data_.size()) {
+      data_.resize(required_size, 0);
+    }
+    std::copy(scan.begin(), scan.end(), data_.begin() + frag_offset);
+  }
+
+  /// Check if there are any missing sequence numbers.
+  /// This will check if there are any missing sequence numbers in the received
+  /// sequence numbers.
+  /// @return True if there are any missing sequence numbers, false otherwise.
+  bool has_missing_sequence_numbers() {
+    if (received_sequence_numbers_.empty() || received_sequence_numbers_.size() == 1) {
+      return false;
+    }
+    // first let's sort the sequence numbers
+    std::sort(received_sequence_numbers_.begin(), received_sequence_numbers_.end());
+    // check for gaps in the sequence numbers. Since we've sorted the list, we
+    // can simply iterate and ensure that the difference between consecutive
+    // sequence numbers is 1.
+    int prev_seq = received_sequence_numbers_[0];
+    for (size_t i = 1; i < received_sequence_numbers_.size(); ++i) {
+      int curr_seq = received_sequence_numbers_[i];
+      if (curr_seq != (prev_seq + 1) % 65536) {
+        return true; // found a gap
+      }
+      prev_seq = curr_seq;
+    }
+    // if we made it here, there are no gaps
+    return false;
+  }
+
   /// Add the EOI marker to the frame.
   /// This will add the EOI marker to the frame. This must be called before
   /// calling get_data().
@@ -119,6 +174,7 @@ class JpegFrame {
     }
   }
 
+  std::vector<int> received_sequence_numbers_;
   std::vector<uint8_t> data_;
   JpegHeader header_;
   bool finalized_ = false;
diff --git a/components/rtsp/src/rtsp_server.cpp b/components/rtsp/src/rtsp_server.cpp
index 2166c0430..40fd27025 100644
--- a/components/rtsp/src/rtsp_server.cpp
+++ b/components/rtsp/src/rtsp_server.cpp
@@ -106,7 +106,7 @@ void RtspServer::send_frame(const espp::JpegFrame &frame) {
 
     static const int type_specific = 0;
     static const int fragment_type = 0;
-    int offset = i * max_data_size_;
+    int offset = start_index;
 
     std::unique_ptr<RtpJpegPacket> packet;
     // if this is the first packet, it has the quantization tables