#!/usr/bin/env python3
import requests
import json
import argparse
import sys

BASE_URL = "https://efa.rvv.de/efa/"
VERSION = "10.5.17.3"

def query_stop(name):
    params = {
        "name_sf": name,
        "outputFormat": "rapidJSON",
        "type_sf": "any",
        "locationInfoActive": "1",
        "version": VERSION
    }
    response = requests.get(f"{BASE_URL}XML_STOPFINDER_REQUEST", params=params)
    return response.json()

def query_departures(stop_id, count=10):
    params = {
        "name_dm": stop_id,
        "type_dm": "stop",
        "outputFormat": "rapidJSON",
        "itdDateTimeDepArr": "dep",
        "useRealtime": "1",
        "mode": "direct",
        "depSequence": count,
        "version": VERSION
    }
    response = requests.get(f"{BASE_URL}XML_DM_REQUEST", params=params)
    return response.json()

def query_trip(origin_id, destination_id, count=4):
    params = {
        "name_origin": origin_id,
        "type_origin": "stop",
        "name_destination": destination_id,
        "type_destination": "stop",
        "outputFormat": "rapidJSON",
        "itdTripDateTimeDepArr": "dep",
        "calcNumberOfTrips": count,
        "useRealtime": "1",
        "mode": "direct",
        "version": VERSION
    }
    response = requests.get(f"{BASE_URL}XML_TRIP_REQUEST2", params=params)
    return response.json()

def main():
    parser = argparse.ArgumentParser(description="Query RVV Regensburg API")
    subparsers = parser.add_subparsers(dest="command", help="Command to run")

    # Stop Finder
    stop_parser = subparsers.add_parser("stop", help="Find a stop ID")
    stop_parser.add_argument("name", help="Name of the stop")

    # Departure Monitor
    dep_parser = subparsers.add_parser("departures", help="Get departures for a stop")
    dep_parser.add_argument("stop_id", help="Stop ID (e.g., de:09362:11010)")
    dep_parser.add_argument("--count", type=int, default=10, help="Number of departures")

    # Trip Planning
    trip_parser = subparsers.add_parser("trip", help="Plan a trip between two stops")
    trip_parser.add_argument("origin_id", help="Origin Stop ID")
    trip_parser.add_argument("destination_id", help="Destination Stop ID")
    trip_parser.add_argument("--count", type=int, default=4, help="Number of trips")

    args = parser.parse_args()

    if args.command == "stop":
        result = query_stop(args.name)
        points = result.get("locations", [])
        output = []
        for p in points:
            if p.get("type") == "stop":
                output.append({
                    "id": p["id"],
                    "stopId": p.get("properties", {}).get("stopId") or p.get("id"),
                    "name": p["name"]
                })
        print(json.dumps(output, indent=2, ensure_ascii=False))
    elif args.command == "departures":
        result = query_departures(args.stop_id, args.count)
        # EFA rapidJSON uses 'stopEvents' for departures
        stop_events = result.get("stopEvents", [])
        output = []
        for event in stop_events:
            transport = event.get("transportation", {})
            departure_time = event.get("departureTimePlanned")
            departure_time_rt = event.get("departureTimeEstimated")
            output.append({
                "line": transport.get("number"),
                "direction": transport.get("destination", {}).get("name"),
                "base_time": departure_time[11:16] if departure_time else None,
                "rt_time": departure_time_rt[11:16] if departure_time_rt else None,
                "delay": event.get("delay", 0)
            })
        print(json.dumps(output, indent=2, ensure_ascii=False))
    elif args.command == "trip":
        result = query_trip(args.origin_id, args.destination_id, args.count)
        # EFA rapidJSON uses 'journeys' for trips
        journeys = result.get("journeys", [])
        output = []
        for j in journeys:
            legs = j.get("legs", [])
            trip_info = {
                "legs": []
            }
            for leg in legs:
                origin = leg.get("origin", {})
                dest = leg.get("destination", {})
                transport = leg.get("transportation", {})
                trip_info["legs"].append({
                    "from": origin.get("name"),
                    "to": dest.get("name"),
                    "line": transport.get("number"),
                    "dep_planned": origin.get("departureTimePlanned", "")[11:16],
                    "dep_est": origin.get("departureTimeEstimated", "")[11:16],
                    "arr_planned": dest.get("arrivalTimePlanned", "")[11:16],
                    "arr_est": dest.get("arrivalTimeEstimated", "")[11:16]
                })
            output.append(trip_info)
        print(json.dumps(output, indent=2, ensure_ascii=False))
    else:
        parser.print_help()

if __name__ == "__main__":
    main()