from __future__ import annotations

from uuid import UUID

import sqlalchemy as sa
from sqlalchemy.orm import Session

from backend.app.db.models import Train
from backend.app.models import TrainCreate
from backend.app.repositories.base import BaseRepository


class TrainRepository(BaseRepository[Train]):
    model = Train

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

    def list_all(self) -> list[Train]:
        statement = sa.select(self.model)
        return list(self.session.scalars(statement))

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

    def create(self, data: TrainCreate) -> Train:
        train = Train(
            designation=data.designation,
            operator_id=self._optional_uuid(data.operator_id),
            home_station_id=self._optional_uuid(data.home_station_id),
            capacity=data.capacity,
            max_speed_kph=data.max_speed_kph,
            consist=data.consist,
        )
        self.session.add(train)
        return train