From bf30e323f15b5b4d3cedddd649586c622dd7c990 Mon Sep 17 00:00:00 2001
From: Kumi <git@kumi.email>
Date: Tue, 2 Jul 2024 11:02:33 +0200
Subject: [PATCH] feat(core): Add shipment management and notifier handlers

Introduced methods to add and remove shipments in the Core class
and enhanced the notifier system with handlers for various shipment
events. This improvement facilitates better shipment tracking and
notification capabilities, allowing dynamic shipment management
through the notifier interface. Minor code cleanup in `__main__.py`.
---
 src/trackbert/__main__.py         |  6 -----
 src/trackbert/classes/core.py     | 16 +++++++++++
 src/trackbert/classes/notifier.py | 44 ++++++++++++++++++++++++++++---
 src/trackbert/notifiers/matrix.py |  1 -
 4 files changed, 57 insertions(+), 10 deletions(-)

diff --git a/src/trackbert/__main__.py b/src/trackbert/__main__.py
index 7c3b9f9..7d0d1a5 100644
--- a/src/trackbert/__main__.py
+++ b/src/trackbert/__main__.py
@@ -4,7 +4,6 @@ from tabulate import tabulate
 import argparse
 import asyncio
 
-
 from .classes.core import Core
 
 
@@ -12,7 +11,6 @@ def main():
     parser = argparse.ArgumentParser()
 
     # Arguments related to the tracker
-
     parser.add_argument(
         "--tracking-number",
         "-n",
@@ -58,7 +56,6 @@ def main():
     )
 
     # Arguments related to the config file
-
     parser.add_argument(
         "--generate-config",
         action="store_true",
@@ -76,7 +73,6 @@ def main():
     args = parser.parse_args()
 
     # Generate config file if requested
-
     config_file = Path(args.config_file or "config.ini")
 
     if args.generate_config:
@@ -91,7 +87,6 @@ def main():
             exit(0)
 
     # Load config file
-
     if args.config_file and not config_file.exists():
         print(f"Config file {config_file} does not exist. Use -g to generate it.")
         exit(1)
@@ -99,7 +94,6 @@ def main():
     tracker = Core(config_file)
 
     # List carriers if requested
-
     if args.list_carriers:
         print("Supported carriers:\n")
 
diff --git a/src/trackbert/classes/core.py b/src/trackbert/classes/core.py
index a9cff89..ea1ac93 100644
--- a/src/trackbert/classes/core.py
+++ b/src/trackbert/classes/core.py
@@ -261,6 +261,18 @@ class Core:
                 self.db.write_event(event)
                 self.notify_event(shipment, event, event == events[-1])
 
+    def add_shipment(self, tracking_number: str, carrier: str, description: Optional[str] = None):
+        logging.info(f"Adding shipment {tracking_number} with carrier {carrier}")
+        self.db.create_shipment(tracking_number, carrier, description)
+        for notifier in self.notifiers:
+            notifier.on_shipment_added(tracking_number, carrier, description)
+
+    def remove_shipment(self, tracking_number: str):
+        logging.info(f"Removing shipment {tracking_number}")
+        self.db.disable_shipment(tracking_number)
+        for notifier in self.notifiers:
+            notifier.on_shipment_removed(tracking_number)
+
     def start_loop(self) -> Never:
         logging.debug("Starting loop")
 
@@ -333,8 +345,12 @@ class Core:
 
     def start(self, config: Optional[PathLike] = None):
         self.notify("Trackbert", "Starting up")
+        for notifier in self.notifiers:
+            notifier.set_core(self)
         self.start_loop()
 
     async def start_async(self, config: Optional[PathLike] = None):
         self.notify("Trackbert", "Starting up")
+        for notifier in self.notifiers:
+            notifier.set_core(self)
         await self.start_loop_async()
diff --git a/src/trackbert/classes/notifier.py b/src/trackbert/classes/notifier.py
index e224f6b..96aa5cc 100644
--- a/src/trackbert/classes/notifier.py
+++ b/src/trackbert/classes/notifier.py
@@ -1,10 +1,48 @@
+from typing import Optional
+
+import logging
+
 class BaseNotifier:
     def __init__(self, **kwargs):
-        pass
+        self.core = None
 
-    def notify(self, title: str, message: str, urgent: bool = False) -> None:
-        raise NotImplementedError
+    def set_core(self, core):
+        self.core = core
 
     @property
     def enabled(self) -> bool:
         raise NotImplementedError
+
+    def notify(self, title: str, message: str, urgent: bool = False) -> None:
+        raise NotImplementedError
+
+    # Handlers for shipment events
+
+    def on_shipment_added(self, tracking_number: str, carrier: str, description: Optional[str] = None):
+        pass
+
+    def on_shipment_removed(self, tracking_number: str):
+        pass
+
+    def on_shipment_updated(self, tracking_number: str, carrier: str, description: Optional[str] = None):
+        pass
+
+    def on_shipment_status(self, tracking_number: str, status: str):
+        pass
+
+    # Methods to interact with the core
+
+    def add_shipment(
+        self, tracking_number: str, carrier: str, description: Optional[str] = None
+    ):
+        if self.core:
+            self.core.add_shipment(tracking_number, carrier, description)
+        else:
+            logging.error("Core instance is not set in notifier")
+
+    def remove_shipment(self, tracking_number: str):
+        if self.core:
+            self.core.remove_shipment(tracking_number)
+        else:
+            logging.error("Core instance is not set in notifier")
+
diff --git a/src/trackbert/notifiers/matrix.py b/src/trackbert/notifiers/matrix.py
index 1bc8c26..9e67b66 100644
--- a/src/trackbert/notifiers/matrix.py
+++ b/src/trackbert/notifiers/matrix.py
@@ -1,4 +1,3 @@
-from typing import Dict, Any
 from urllib.request import Request, urlopen
 
 import json