Update app.py

app.py

@@ -10,16 +10,10 @@ CLIENT_ID = os.environ.get('client_id')
 API_KEY = os.environ.get('api_key')
 
 def get_station_info(pattern, client_id, api_key):
-    """Ermittelt EVA-Nummer und offiziellen Namen."""
     url = f"https://apis.deutschebahn.com/db-api-marketplace/apis/timetables/v1/station/{pattern}"
-    headers = {
-        'DB-Client-Id': client_id,
-        'DB-Api-Key': api_key,
-        'accept': 'application/xml'
-    }
+    headers = {'DB-Client-Id': client_id, 'DB-Api-Key': api_key, 'accept': 'application/xml'}
     try:
         response = requests.get(url, headers=headers, timeout=10)
-        response.raise_for_status()
         root = ET.fromstring(response.content)
         station = root.find('station')
         if station is not None:
@@ -29,13 +23,8 @@ def get_station_info(pattern, client_id, api_key):
     return None
 
 def fetch_timetable_hour(eva, date_str, hour_str, client_id, api_key):
-    """Ruft den Fahrplan für eine bestimmte Stunde ab."""
     url = f"https://apis.deutschebahn.com/db-api-marketplace/apis/timetables/v1/plan/{eva}/{date_str}/{hour_str}"
-    headers = {
-        'DB-Client-Id': client_id,
-        'DB-Api-Key': api_key,
-        'accept': 'application/xml'
-    }
+    headers = {'DB-Client-Id': client_id, 'DB-Api-Key': api_key, 'accept': 'application/xml'}
     try:
         response = requests.get(url, headers=headers, timeout=10)
         if response.status_code == 200:
@@ -45,12 +34,9 @@ def fetch_timetable_hour(eva, date_str, hour_str, client_id, api_key):
     return None
 
 def search_next_3_connections(departure, destination, client_id, api_key):
-    """Hauptfunktion: Findet die nächsten 3 Verbindungen ab 'jetzt'."""
-    
     if not all([departure, destination, client_id, api_key]):
         return {"error": "Bitte alle Felder ausfüllen."}
 
-    # 1. IBNRs ermitteln
     dep_info = get_station_info(departure, client_id, api_key)
     dest_info = get_station_info(destination, client_id, api_key)
     
@@ -60,8 +46,7 @@ def search_next_3_connections(departure, destination, client_id, api_key):
     now = datetime.now()
     found_connections = []
     
-    # 2. Suche in der aktuellen und der nächsten Stunde (um 3 Verbindungen zu garantieren)
-    for hour_offset in range(3): # Wir schauen bis zu 2 Stunden voraus
+    for hour_offset in range(3): 
         search_time = now + timedelta(hours=hour_offset)
         datum = search_time.strftime("%y%m%d")
         stunde = search_time.strftime("%H")
@@ -71,61 +56,78 @@ def search_next_3_connections(departure, destination, client_id, api_key):
             
         for train in root.findall('.//s'):
             dp = train.find('dp')
-            if dp is None: continue
+            tl = train.find('tl')
+            if dp is None or tl is None: continue
             
-            # Prüfen ob Ziel im geplanten Pfad (ppth)
+            # Check if destination is in the path
             path = dp.get('ppth', '')
             if dest_info['name'].lower() in path.lower() or destination.lower() in path.lower():
-                
-                pt = dp.get('pt', '') # Geplante Zeit: YYMMDDHHMM
+                pt = dp.get('pt', '')
                 if not pt: continue
                 
                 dep_dt = datetime.strptime(pt, "%y%m%d%H%M")
                 
-                # Nur Verbindungen in der Zukunft (ab Zeitstempel jetzt)
                 if dep_dt >= now:
-                    tl = train.find('tl')
+                    # To get arrival time, we must check the destination station for the same train ID
+                    # For performance, we assume arrival is within +2 hours of departure
+                    arr_time_str = "N/A"
+                    duration_str = "N/A"
+                    
+                    # Search destination timetable for the arrival
+                    # We check the hour of departure and the next hour for arrival
+                    arrival_root = fetch_timetable_hour(dest_info['eva'], dep_dt.strftime("%y%m%d"), dep_dt.strftime("%H"), client_id, api_key)
+                    
+                    if arrival_root is not None:
+                        # Find the train by number (n) and category (c)
+                        for arr_train in arrival_root.findall('.//s'):
+                            arr_tl = arr_train.find('tl')
+                            arr_ar = arr_train.find('ar')
+                            if arr_tl is not None and arr_ar is not None:
+                                if arr_tl.get('n') == tl.get('n'):
+                                    arr_pt = arr_ar.get('pt', '')
+                                    arr_dt = datetime.strptime(arr_pt, "%y%m%d%H%M")
+                                    arr_time_str = arr_dt.strftime("%H:%M")
+                                    
+                                    # Calculate duration
+                                    diff = arr_dt - dep_dt
+                                    duration_str = f"{int(diff.total_seconds() // 60)} min"
+                                    break
+
                     found_connections.append({
                         "startort": dep_info['name'],
                         "zielort": dest_info['name'],
                         "abfahrtszeit": dep_dt.strftime("%H:%M"),
+                        "ankunftszeitszeit": arr_time_str,
                         "gleis": dp.get('pp', 'n/a'),
-                        "zug": f"{tl.get('c', '')} {tl.get('n', '')}",
-                        "_dt": dep_dt # Hilfsfeld zum Sortieren
+                        "duration": duration_str,
+                        "_dt": dep_dt 
                     })
 
-    # 3. Sortieren und Top 3 auswählen
     found_connections.sort(key=lambda x: x['_dt'])
     
-    # Bereinigen der Ergebnisse für die JSON-Ausgabe
-    final_results = []
-    for c in found_connections[:3]:
-        # Entferne das interne Datetime-Objekt vor der JSON-Konvertierung
-        res = {k: v for k, v in c.items() if k != "_dt"}
-        final_results.append(res)
-
-    return final_results
+    # Return list of first 3 matches formatted without internal sort key
+    return [{k: v for k, v in c.items() if k != "_dt"} for c in found_connections[:3]]
 
 # --- Gradio UI ---
 
 def ui_wrapper(dep, dest):
     if not CLIENT_ID or not API_KEY:
-        return {"error": "API credentials not configured. Please set client_id and api_key as environment variables."}
+        return {"error": "API credentials not configured."}
     
     results = search_next_3_connections(dep, dest, CLIENT_ID, API_KEY)
-    return json.dumps(results, indent=4, ensure_ascii=False)
+    return results # Return list directly for gr.JSON
 
 with gr.Blocks(title="DB JSON Fahrplan", theme=gr.themes.Soft()) as demo:
     gr.Markdown("# │ DB │ Deutsche Bahn train connections")
     
     with gr.Row():
-        dep_input = gr.Textbox(label="Abfahrtsort", placeholder="z.B. Berlin")
-        dest_input = gr.Textbox(label="Zielort", placeholder="z.B. Hamburg")
+        dep_input = gr.Textbox(label="Abfahrtsort", placeholder="z.B. Schweinfurt Hbf")
+        dest_input = gr.Textbox(label="Zielort", placeholder="z.B. Oerlenbach")
         
     btn = gr.Button("Suchen", variant="primary")
-    output = gr.JSON(label="JSON Ergebnis") # Nutzt Gradio JSON Komponente für schöneres Format
+    output = gr.JSON(label="JSON Ergebnis")
     
     btn.click(fn=ui_wrapper, inputs=[dep_input, dest_input], outputs=output)
 
 if __name__ == "__main__":
-    demo.launch(mcp_server=True)
+    demo.launch()