From d4bec29e49562229564d98e2421292519a4d02d8 Mon Sep 17 00:00:00 2001
From: Fred Boniface <fred@fjla.uk>
Date: Tue, 29 Oct 2024 13:08:27 +0000
Subject: [PATCH] Add 'train' methods

---
 pyOwlBoard/client.py   | 62 +++++++++++++++++++++++++++++++++++++++++-
 pyOwlBoard/contants.py |  1 +
 pyOwlBoard/utils.py    | 20 +++++++++++++-
 3 files changed, 81 insertions(+), 2 deletions(-)

diff --git a/pyOwlBoard/client.py b/pyOwlBoard/client.py
index 0874d9c..a1221bb 100644
--- a/pyOwlBoard/client.py
+++ b/pyOwlBoard/client.py
@@ -16,10 +16,12 @@
 
 import requests
 import logging
+from datetime import datetime
 from typing import List, Dict, Tuple
 from urllib.parse import urljoin
 
 from .contants import ENDPOINTS, VERSION
+from .utils import format_url_datetime, url_multijoin
 
 logger = logging.getLogger('OwlBoardClient')
 logger.setLevel(logging.INFO)
@@ -67,12 +69,70 @@ class OwlBoardClient:
         print(response.text)
 
     ### Train Methods ###
+    def get_trains_by_headcode(self, headcode: str, date: datetime):
+        if not isinstance(headcode, str):
+            raise TypeError("headcode must be a string")
+        if not isinstance(date, datetime):
+            raise TypeError("date must be a datetime object")
+
+        # Generate URL
+        try:
+            url_path = url_multijoin(
+                self.base_url,
+                ENDPOINTS['TIMETABLE_TRAINS'],
+                format_url_datetime(date),
+                'headcode',
+                headcode,
+            )
+            logger.debug(f"Generated URL: {url_path}")
+        except Exception as e:
+            logger.error(f"Error generating URL: {e}")
+        
+        # Send request
+        try:
+            response = requests.get(url_path, headers=self.headers)
+            response.raise_for_status()
+            logger.info("Request completed")
+            print(response.text)
+        except requests.RequestException as e:
+            logger.error(f"Request failed: {e}")
+            raise
+
+    def get_trains_by_trainUid(self, train_uid: str, date: datetime):
+        if not isinstance(train_uid, str):
+            raise TypeError("train_uid must be a string")
+        if not isinstance(date, datetime):
+            raise TypeError("date must be a datetime object")
+        
+        # Generate URL
+        try:
+            url_path = url_multijoin(
+                self.base_url,
+                ENDPOINTS['TIMETABLE_TRAINS'],
+                format_url_datetime(date),
+                'byTrainUid',
+                train_uid,
+            )
+            logger.debug(f"Generated URL: {url_path}")
+        except Exception as e:
+            logger.error(f"Error generating URL: {e}")
+            raise
+        
+        # Send request
+        try:
+            response = requests.get(url_path, headers=self.headers)
+            response.raise_for_status()
+            logger.info("Request completed")
+            print(response.text)
+        except requests.RequestException as e:
+            logger.error(f"Request failed: {e}")
+            raise
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.DEBUG)
 
     try:
         client = OwlBoardClient(base_url='https://owlboard.info', api_key="x")
-        client.get_pis_by_start_end_crs('bri', 'wsb')
+        client.get_trains_by_headcode("1A99", datetime.now())
     except Exception as e:
         logger.error(f"Failed to create client: {e}")
\ No newline at end of file
diff --git a/pyOwlBoard/contants.py b/pyOwlBoard/contants.py
index 0c026b9..82c0564 100644
--- a/pyOwlBoard/contants.py
+++ b/pyOwlBoard/contants.py
@@ -4,4 +4,5 @@ ENDPOINTS = {
     'TEST': '/api/v1/auth/test/',
     'PIS_BY_CODE': '/api/v2/pis/byCode/',
     'PIS_BY_START_END_CRS': '/api/v2/pis/byStartEnd/',
+    'TIMETABLE_TRAINS': '/api/v2/timetable/train'
 }
\ No newline at end of file
diff --git a/pyOwlBoard/utils.py b/pyOwlBoard/utils.py
index 8019d5b..71d8f51 100644
--- a/pyOwlBoard/utils.py
+++ b/pyOwlBoard/utils.py
@@ -12,4 +12,22 @@
 #    GNU General Public License for more details.
 
 #    You should have received a copy of the GNU General Public License
-#    along with this program.  If not, see <https://www.gnu.org/licenses/>.
\ No newline at end of file
+#    along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+from urllib.parse import urljoin
+from datetime import datetime
+
+def url_multijoin(*parts: str) -> str:
+    if not parts:
+        return ""
+
+    base = parts[0]
+    for part in parts[1:]:
+        base = urljoin(base.rstrip('/') + '/', part.lstrip('/'))
+
+    return base
+
+def format_url_datetime(input: datetime) -> str:
+    if not isinstance(input, datetime):
+        raise TypeError("Expected a datetime object")
+    return input.strftime("%Y-%m-%d")
\ No newline at end of file