carrramba-encore-rate/backend/idfm_matrix_backend/models/line.py

from asyncio import gather as asyncio_gather
from collections import defaultdict
from typing import Iterable, Self

from sqlalchemy import (
    BigInteger,
    Boolean,
    Column,
    Enum,
    ForeignKey,
    Integer,
    select,
    String,
    Table,
)
from sqlalchemy.orm import Mapped, relationship, selectinload
from sqlalchemy.orm.attributes import set_committed_value
from sqlalchemy.sql.expression import tuple_

from ..db import Base, db
from ..idfm_interface.idfm_types import (
    IdfmState,
    IdfmLineState,
    TransportMode,
    TransportSubMode,
)
from .stop import _Stop

line_stop_association_table = Table(
    "line_stop_association_table",
    Base.metadata,
    Column("line_id", ForeignKey("lines.id")),
    Column("stop_id", ForeignKey("_stops.id")),
)


class LinePicto(Base):

    db = db

    id = Column(String, primary_key=True)
    mime_type = Column(String, nullable=False)
    height_px = Column(Integer, nullable=False)
    width_px = Column(Integer, nullable=False)
    filename = Column(String, nullable=False)
    url = Column(String, nullable=False)
    thumbnail = Column(Boolean, nullable=False)
    format = Column(String, nullable=False)

    __tablename__ = "line_pictos"


class Line(Base):

    db = db

    id = Column(String, primary_key=True)

    short_name = Column(String)
    name = Column(String, nullable=False)
    status = Column(Enum(IdfmLineState), nullable=False)
    transport_mode = Column(Enum(TransportMode), nullable=False)
    transport_submode = Column(Enum(TransportSubMode), nullable=False)

    network_name = Column(String)
    group_of_lines_id = Column(String)
    group_of_lines_shortname = Column(String)

    colour_web_hexa = Column(String, nullable=False)
    text_colour_hexa = Column(String, nullable=False)

    operator_id = Column(String)
    operator_name = Column(String)

    accessibility = Column(Enum(IdfmState), nullable=False)
    visual_signs_available = Column(Enum(IdfmState), nullable=False)
    audible_signs_available = Column(Enum(IdfmState), nullable=False)

    picto_id = Column(String, ForeignKey("line_pictos.id"))
    picto: Mapped[LinePicto] = relationship(LinePicto, lazy="selectin")

    record_id = Column(String, nullable=False)
    record_ts = Column(BigInteger, nullable=False)

    stops: Mapped[list["_Stop"]] = relationship(
        "_Stop",
        secondary=line_stop_association_table,
        back_populates="lines",
        lazy="selectin",
    )

    __tablename__ = "lines"

    @classmethod
    async def get_by_name(
        cls, name: str, operator_name: None | str = None
    ) -> list[Self]:
        filters = {"name": name}
        if operator_name is not None:
            filters["operator_name"] = operator_name

        lines = None
        stmt = (
            select(Line)
            .filter_by(**filters)
            .options(selectinload(Line.stops), selectinload(Line.picto))
        )
        res = await cls.db.session.execute(stmt)
        lines = res.scalars().all()
        return lines

    @classmethod
    async def _add_picto_to_line(cls, line: str | Self, picto: LinePicto) -> None:
        if isinstance(line, str):
            if (lines := await cls.get_by_name(line)) is not None:
                if len(lines) == 1:
                    line = lines[0]
                else:
                    for candidate_line in lines:
                        if candidate_line.operator_name == "RATP":
                            line = candidate_line
                            break

        if isinstance(line, Line) and line.picto is None:
            line.picto = picto
            line.picto_id = picto.id

    @classmethod
    async def add_pictos(cls, line_to_pictos: dict[str | Self, LinePicto]) -> None:
        await asyncio_gather(
            *[
                cls._add_picto_to_line(line, picto)
                for line, picto in line_to_pictos.items()
            ]
        )

        await cls.db.session.commit()

    @classmethod
    async def add_stops(cls, line_to_stop_ids: Iterable[tuple[str, str, str]]) -> int:
        line_names_ops, stop_ids = set(), set()
        for line_name, operator_name, stop_id in line_to_stop_ids:
            line_names_ops.add((line_name, operator_name))
            stop_ids.add(stop_id)

        res = await cls.db.session.execute(
            select(Line).where(
                tuple_(Line.name, Line.operator_name).in_(line_names_ops)
            )
        )

        lines = defaultdict(list)
        for line in res.scalars():
            lines[(line.name, line.operator_name)].append(line)

        res = await cls.db.session.execute(select(_Stop).where(_Stop.id.in_(stop_ids)))
        stops = {stop.id: stop for stop in res.scalars()}

        found = 0
        for line_name, operator_name, stop_id in line_to_stop_ids:
            if (stop := stops.get(stop_id)) is not None:
                if (stop_lines := lines.get((line_name, operator_name))) is not None:
                    if len(stop_lines) > 1:
                        print(stop_lines)
                    for stop_line in stop_lines:
                        stop_line.stops.append(stop)
                    found += 1
                else:
                    print(f"No line found for {line_name}/{operator_name}")
            else:
                print(
                    f"No stop found for {stop_id} id (used by {line_name}/{operator_name})"
                )

        await cls.db.session.commit()
        return found