1 files changed, 85 insertions, 48 deletions
diff --git a/prediction_controller.py b/prediction_controller.py
index 4d5c1ed..0188e2e 100644
--- a/prediction_controller.py
+++ b/prediction_controller.py
@@ -1,48 +1,85 @@
-# prediction_controller.py
-from dateutil.relativedelta import relativedelta
-from datetime import datetime, timedelta
-import math
-
-class PredictionController:
-    def __init__(self, calendar_manager, date_calculator, keyword_list):
-        self.calendar_manager = calendar_manager
-        self.date_calculator = date_calculator
-        self.launch_date = None
-        self.duration = None
-        self.prediction = None
-        self.keyword_list = keyword_list
-        for keyword in keyword_list:
-            self.keyword = []
-        
-    def set_parameters(self, launch_date, duration_years):
-        self.launch_date = datetime.fromisoformat(launch_date)
-        self.duration = relativedelta(years=duration_years)
-    
-    def make_prediction(self, launch_date, duration_years):
-        self.set_parameters(launch_date, duration_years)
-
-        prediction = self.launch_date + self.duration - timedelta(days=1)
-
-        keyword_args = {}
-
-        for entry in self.calendar_manager.entries:
-            for keyword in self.keyword_list:
-                if entry.keyword == keyword:
-                    if keyword not in keyword_args:
-                        keyword_args[keyword] = []
-                    keyword_args[keyword].append((entry.start_date, entry.end_date, entry.id))
-                    break
-
-        prediction, corrected_events = self.date_calculator.calculate_prediction(self.launch_date, self.duration, **keyword_args)
-        self.prediction = prediction
-        self.calendar_manager.correct_dates(corrected_events)
-        
-    def get_launch_date(self):
-        return self.launch_date
-        
-    def get_duration(self):
-        return self.duration
-        
-    def get_prediction(self):
-        return self.prediction
-
+# prediction_controller.py
+from dateutil.relativedelta import relativedelta
+from datetime import datetime, timedelta
+import math
+from prediction_storage import PredictionStorage
+
+class PredictionController:
+    def __init__(self, calendar_manager, date_calculator, keyword_list, prediction_storage=None):
+        self.calendar_manager = calendar_manager
+        self.date_calculator = date_calculator
+        self.launch_date = None
+        self.duration = None
+        self.prediction = None
+        self.keyword_list = keyword_list
+        self.prediction_storage = prediction_storage
+        for keyword in keyword_list:
+            self.keyword = []
+        
+    def set_parameters(self, launch_date, duration_years):
+        self.launch_date = datetime.fromisoformat(launch_date) if isinstance(launch_date, str) else launch_date
+        self.duration = relativedelta(years=duration_years)
+    
+    def make_prediction(self, launch_date, duration_years, description=None):
+        self.set_parameters(launch_date, duration_years)
+
+        prediction = self.launch_date + self.duration - timedelta(days=1)
+
+        keyword_args = {}
+
+        for entry in self.calendar_manager.entries:
+            for keyword in self.keyword_list:
+                if entry.keyword == keyword:
+                    if keyword not in keyword_args:
+                        keyword_args[keyword] = []
+                    keyword_args[keyword].append((entry.start_date, entry.end_date, entry.id))
+                    break
+
+        prediction, corrected_events = self.date_calculator.calculate_prediction(self.launch_date, self.duration, **keyword_args)
+        self.prediction = prediction
+        self.calendar_manager.correct_dates(corrected_events)
+        
+        # Store the prediction if we have a storage object
+        if self.prediction_storage:
+            self.prediction_storage.add_prediction(
+                launch_date=self.launch_date,
+                duration_years=duration_years,
+                predicted_date=prediction,
+                keyword_args=keyword_args,
+                description=description
+            )
+        
+        return prediction
+        
+    def get_launch_date(self):
+        return self.launch_date
+        
+    def get_duration(self):
+        return self.duration
+        
+    def get_prediction(self):
+        return self.prediction
+        
+    def get_all_predictions(self):
+        """Get all stored predictions."""
+        if self.prediction_storage:
+            return self.prediction_storage.list_predictions()
+        return []
+    
+    def get_prediction_by_id(self, prediction_id):
+        """Get a specific prediction by ID."""
+        if self.prediction_storage:
+            return self.prediction_storage.get_prediction_by_id(prediction_id)
+        return None
+    
+    def search_predictions(self, start_date=None, end_date=None, keyword=None):
+        """Search predictions by date range or keyword."""
+        if self.prediction_storage:
+            return self.prediction_storage.search_predictions(start_date, end_date, keyword)
+        return []
+    
+    def update_prediction_description(self, prediction_id, description):
+        """Update a prediction's description."""
+        if self.prediction_storage:
+            return self.prediction_storage.update_prediction_description(prediction_id, description)
+        return None
+\ No newline at end of file