import pandas as pd
from pathlib import Path
from rules.train_rules import classify_train_number


class TimetableService:
    def __init__(self, data_dir="data/timetable_parquet"):
        self.data_dir = Path(data_dir)

    def get_train_meta(self, train_number: int):
        return classify_train_number(train_number)

    def get_train_timetable(self, train_number: int, diagram_type: str):
        df = pd.read_parquet(
            self.data_dir,
            filters=[
                ("diagram_type", "=", diagram_type),
                ("train_number", "=", train_number)
            ]
        )
        return df.sort_values("seq")

    def get_full_info(self, train_number: int, diagram_type: str):
        meta = self.get_train_meta(train_number)
        timetable = self.get_train_timetable(train_number, diagram_type)

        return {
            "meta": meta,
            "timetable": timetable.to_dict("records")
        }
