From bc28e60b8c67488a66cbce46d7189a43076d37b3 Mon Sep 17 00:00:00 2001
From: Arjen Bos <1064589+arjenbos@users.noreply.github.com>
Date: Fri, 5 Jan 2024 11:49:50 +0100
Subject: [PATCH] Add calendar sensor

---
 custom_components/postnl/calendar.py | 93 ++++++++++++++++++++++++++++
 custom_components/postnl/const.py    |  5 +-
 custom_components/postnl/sensor.py   |  5 +-
 3 files changed, 99 insertions(+), 4 deletions(-)
 create mode 100644 custom_components/postnl/calendar.py

diff --git a/custom_components/postnl/calendar.py b/custom_components/postnl/calendar.py
new file mode 100644
index 0000000..249fd32
--- /dev/null
+++ b/custom_components/postnl/calendar.py
@@ -0,0 +1,93 @@
+import datetime
+import logging
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant, callback
+from homeassistant.components.calendar import CalendarEntity, CalendarEvent
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
+
+from .coordinator import PostNLCoordinator
+from .structs.package import Package
+
+_LOGGER = logging.getLogger(__name__)
+
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    async_add_entities: AddEntitiesCallback,
+) -> None:
+    """Set up the PostNL Calendar config entry."""
+
+    coordinator = PostNLCoordinator(hass)
+
+    async_add_entities(
+        [
+            PostNLCalendar(coordinator=coordinator)
+        ]
+    )
+
+
+class PostNLCalendar(CoordinatorEntity[PostNLCoordinator], CalendarEntity):
+
+    def __init__(self, coordinator):
+        super().__init__(coordinator=coordinator)
+        self._attr_name = "PostNL"
+        self._events: list[CalendarEvent] = []
+
+    @callback
+    def _handle_coordinator_update(self) -> None:
+        _LOGGER.debug('Updating calendar %s', self.name)
+
+        self._events = []
+
+        for package in self.coordinator.data['receiver']:
+            calendar_event = CalendarEvent(
+                start=datetime.datetime.fromisoformat(package.planned_from),
+                end=datetime.datetime.fromisoformat(package.planned_to),
+                summary=package.name,
+                uid=package.key
+            )
+
+            self._events.append(calendar_event)
+
+        for package in self.coordinator.data['sender']:
+            calendar_event = CalendarEvent(
+                start=datetime.datetime.fromisoformat(package.planned_from),
+                end=datetime.datetime.fromisoformat(package.planned_to),
+                summary=package.name,
+                uid=package.key
+            )
+
+            self._events.append(calendar_event)
+
+        self._events.sort(key=lambda event: event.start)
+
+        super()._handle_coordinator_update()
+
+    @property
+    def event(self) -> CalendarEvent | None:
+        """Return the next upcoming event."""
+
+        filtered_events = filter(lambda event: event.start.timestamp() > datetime.datetime.now().timestamp(), self._events)
+        events = list(filtered_events)
+
+        return events[0] if len(events) > 0 else None
+
+    async def async_get_events(
+        self,
+        hass: HomeAssistant,
+        start_date: datetime.datetime,
+        end_date: datetime.datetime,
+    ) -> list[CalendarEvent]:
+        """Return calendar events within a datetime range."""
+        _LOGGER.debug("Getting events between %s and %s", start_date, end_date)
+
+        filtered_events = filter(lambda event: event.start >= start_date <= end_date, self._events)
+        events = list(filtered_events)
+
+        _LOGGER.debug("Events: %s", events)
+
+        return events
diff --git a/custom_components/postnl/const.py b/custom_components/postnl/const.py
index 53bdc2e..acad4ab 100644
--- a/custom_components/postnl/const.py
+++ b/custom_components/postnl/const.py
@@ -9,5 +9,6 @@
 
 
 PLATFORMS = [
-    Platform.SENSOR
-]
\ No newline at end of file
+    Platform.SENSOR,
+    Platform.CALENDAR
+]
diff --git a/custom_components/postnl/sensor.py b/custom_components/postnl/sensor.py
index f9865fc..ebe907b 100644
--- a/custom_components/postnl/sensor.py
+++ b/custom_components/postnl/sensor.py
@@ -6,6 +6,7 @@
 from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
+from .calendar import PostNLCalendar
 from .coordinator import PostNLCoordinator
 from .structs.package import Package
 
@@ -32,7 +33,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry, async_add_e
     ])
 
 
-class PostNLDelivery(CoordinatorEntity, Entity):
+class PostNLDelivery(CoordinatorEntity[PostNLCoordinator], Entity):
     def __init__(self, coordinator, name, receiver: bool = True):
         """Initialize the PostNL sensor."""
         super().__init__(coordinator, context=name)
@@ -81,7 +82,7 @@ def _handle_coordinator_update(self) -> None:
 
         self.handle_coordinator_data()
 
-        self.async_write_ha_state()
+        super()._handle_coordinator_update()
 
     def handle_coordinator_data(self):
         self._attributes['delivered'] = []