import difflib
from functools import reduce
from logging import getLogger
from operator import iadd
from typing import Annotated
from fastmcp import Context
from pydantic import BaseModel, Field
from nagoya_bus_mcp.client import Client
log = getLogger(__name__)
class StationNumberResponse(BaseModel):
success: bool
station_name: str | None = None
station_number: int | None = None
class TimeTable(BaseModel):
route: Annotated[str, Field(description="路線")]
direction: Annotated[str, Field(description="方面")]
pole: Annotated[str, Field(description="乗り場")]
stop_stations: Annotated[list[str], Field(description="停車バス停のリスト")]
timetable: Annotated[dict[str, list[str]], Field(description="曜日別の時刻表")]
url: str
class TimeTableResponse(BaseModel):
station_number: int
timetables: list[TimeTable]
url: str
_cached_station_names: dict[str, int] | None = None
_cached_station_numbers: dict[int, str] | None = None
async def _get_station_names(client: Client) -> dict[str, int]:
global _cached_station_names # noqa: PLW0603
if _cached_station_names is None:
_cached_station_names = (await client.get_station_names()).root
return _cached_station_names
async def _get_station_numbers(client: Client) -> dict[int, str]:
global _cached_station_numbers # noqa: PLW0603
if _cached_station_numbers is None:
_cached_station_numbers = {
num: name for name, num in (await _get_station_names(client)).items()
}
return _cached_station_numbers
async def get_station_number(
ctx: Context, station_name: str
) -> StationNumberResponse | None:
"""Get station number for a given station name."""
client = ctx.request_context.lifespan_context.bus_client
log.info("Getting station number for %s", station_name)
station_names = await _get_station_names(client)
# First try exact match
station_number = station_names.get(station_name)
if station_number is not None:
return StationNumberResponse(
success=True, station_name=station_name, station_number=station_number
)
# If no exact match, try fuzzy matching
log.info("No exact match found for %s, trying fuzzy matching", station_name)
closest_matches = difflib.get_close_matches(
station_name, station_names.keys(), n=1, cutoff=0.6
)
if closest_matches:
closest_station = closest_matches[0]
closest_station_number = station_names[closest_station]
log.info(
"Found closest match: %s (station number: %s)",
closest_station,
closest_station_number,
)
return StationNumberResponse(
success=True,
station_name=closest_station,
station_number=closest_station_number,
)
log.info("No fuzzy match found for %s", station_name)
return StationNumberResponse(success=False)
async def get_timetable(ctx: Context, station_number: int) -> TimeTableResponse | None:
"""Get timetable for a given station."""
client = ctx.request_context.lifespan_context.bus_client
station_name = (await _get_station_numbers(client)).get(station_number)
if station_name is None:
log.warning("Station number %s not found", station_number)
return None
log.info("Getting timetable for station number %s", station_number)
diagram_response = await client.get_station_diagram(station_number)
if not diagram_response.root:
log.error(
"Failed to get station diagram for %s (%s)", station_name, station_number
)
return None
timetables: list[TimeTable] = []
for line, railways in diagram_response.root.items():
for railway_index, railway in enumerate(railways):
diagram: dict[str, list[str]] = {}
for day, hour_minutes in railway.diagram.root.items():
diagram[day] = []
for hour, minutes in hour_minutes.items():
for minute in minutes:
diagram[day].append(f"{hour}:{minute:02}")
timetables.append(
TimeTable(
route=line,
direction="・".join(railway.railway),
stop_stations=reduce(iadd, railway.stations, []),
pole=railway.polename,
timetable=diagram,
url=f"{client.base_url}/jp/pc/bus/timetable_dtl.html?name={station_name}&keito={line}&lineindex={railway_index}",
)
)
return TimeTableResponse(
timetables=timetables,
station_number=station_number,
url=f"{client.base_url}/jp/pc/bus/timetable_list.html?name={station_name}&toname=",
)
