From f0cce348c46fb3cb08ac3412d69d325c95329da5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20Rebelo?= <joserebelo@outlook.com>
Date: Wed, 16 Oct 2024 22:46:19 +0100
Subject: [PATCH] Xiaomi-protobuf: Fix naps

---
 .../devices/xiaomi/XiaomiSampleProvider.java    | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/app/src/main/java/nodomain/freeyourgadget/gadgetbridge/devices/xiaomi/XiaomiSampleProvider.java b/app/src/main/java/nodomain/freeyourgadget/gadgetbridge/devices/xiaomi/XiaomiSampleProvider.java
index 8b5167703..c76068829 100644
--- a/app/src/main/java/nodomain/freeyourgadget/gadgetbridge/devices/xiaomi/XiaomiSampleProvider.java
+++ b/app/src/main/java/nodomain/freeyourgadget/gadgetbridge/devices/xiaomi/XiaomiSampleProvider.java
@@ -123,10 +123,20 @@ public class XiaomiSampleProvider extends AbstractSampleProvider<XiaomiActivityS
      * found.
      */
     public void overlaySleep(final List<XiaomiActivitySample> samples, final int timestamp_from, final int timestamp_to) {
-        final RangeMap<Long, ActivityKind> stagesMap = new RangeMap<>();
+        final RangeMap<Long, ActivityKind> stagesMap = new RangeMap<>(RangeMap.Mode.LOWER_BOUND);
 
+        final XiaomiSleepTimeSampleProvider sleepTimeSampleProvider = new XiaomiSleepTimeSampleProvider(getDevice(), getSession());
         final XiaomiSleepStageSampleProvider sleepStagesSampleProvider = new XiaomiSleepStageSampleProvider(getDevice(), getSession());
 
+        // First populate all samples within this range
+        final List<XiaomiSleepTimeSample> sleepTimesWithinRange = sleepTimeSampleProvider.getAllSamples(timestamp_from * 1000L, timestamp_to * 1000L);
+        LOG.trace("Found {} sleep samples between {} and {}", sleepTimesWithinRange.size(), timestamp_from, timestamp_to);
+
+        for (final XiaomiSleepTimeSample sleepTimeSample : sleepTimesWithinRange) {
+            stagesMap.put(sleepTimeSample.getWakeupTime(), ActivityKind.UNKNOWN);
+            stagesMap.put(sleepTimeSample.getTimestamp(), ActivityKind.LIGHT_SLEEP);
+        }
+
         // Retrieve the last stage before this time range, as the user could have been asleep during
         // the range transition
         final XiaomiSleepStageSample lastSleepStageBeforeRange = sleepStagesSampleProvider.getLastSampleBefore(timestamp_from * 1000L);
@@ -151,13 +161,14 @@ public class XiaomiSampleProvider extends AbstractSampleProvider<XiaomiActivityS
             }
         }
 
-        final XiaomiSleepTimeSampleProvider sleepTimeSampleProvider = new XiaomiSleepTimeSampleProvider(getDevice(), getSession());
 
         // Find last sleep sample before the requested range, as the recorded wake up time may be
         // in the current range
         final XiaomiSleepTimeSample lastSleepTimesBeforeRange = sleepTimeSampleProvider.getLastSampleBefore(timestamp_from * 1000L);
 
         if (lastSleepTimesBeforeRange != null) {
+            LOG.debug("Last sleep time before range: ts={}, stage={}", lastSleepTimesBeforeRange.getTimestamp(), lastSleepTimesBeforeRange);
+
             stagesMap.put(lastSleepTimesBeforeRange.getWakeupTime(), ActivityKind.UNKNOWN);
             stagesMap.put(lastSleepTimesBeforeRange.getTimestamp(), ActivityKind.LIGHT_SLEEP);
         }
@@ -182,7 +193,7 @@ public class XiaomiSampleProvider extends AbstractSampleProvider<XiaomiActivityS
         }
 
         if (!stagesMap.isEmpty()) {
-            LOG.debug("Found {} sleep samples between {} and {}", stagesMap.size(), timestamp_from, timestamp_to);
+            LOG.debug("Found {} sleep stage samples between {} and {}", stagesMap.size(), timestamp_from, timestamp_to);
             // FIXME if no samples were retrieved from the database that were generated by other
             //       activity files, the stages will not get overlayed/inserted and the sleep charts
             //       will stay empty.