from __future__ import annotations

import sqlalchemy as sa
from uuid import UUID
from sqlalchemy.orm import Session

from backend.app.db.models import TrainSchedule
from backend.app.models import TrainScheduleCreate
from backend.app.repositories.base import BaseRepository


class TrainScheduleRepository(BaseRepository[TrainSchedule]):
    """Persistence operations for train timetables."""

    model = TrainSchedule

    def __init__(self, session: Session) -> None:
        super().__init__(session)

    @staticmethod
    def _ensure_uuid(value: UUID | str) -> UUID:
        if isinstance(value, UUID):
            return value
        return UUID(str(value))

    def list_for_train(self, train_id: UUID | str) -> list[TrainSchedule]:
        identifier = self._ensure_uuid(train_id)
        statement = (
            sa.select(self.model)
            .where(self.model.train_id == identifier)
            .order_by(self.model.sequence_index.asc())
        )
        return list(self.session.scalars(statement))

    def create(self, data: TrainScheduleCreate) -> TrainSchedule:
        schedule = TrainSchedule(
            train_id=self._ensure_uuid(data.train_id),
            station_id=self._ensure_uuid(data.station_id),
            sequence_index=data.sequence_index,
            scheduled_arrival=data.scheduled_arrival,
            scheduled_departure=data.scheduled_departure,
            dwell_seconds=data.dwell_seconds,
        )
        self.session.add(schedule)
        return schedule