from typing import Iterable, Self

from sqlalchemy import (
    BigInteger,
    Column,
    Enum,
    Float,
    ForeignKey,
    select,
    String,
    Table,
)
from sqlalchemy.orm import Mapped, relationship, selectinload, with_polymorphic
from sqlalchemy.schema import Index

from ..db import Base, db
from ..idfm_interface.idfm_types import TransportMode, IdfmState, StopAreaType

stop_area_stop_association_table = Table(
    "stop_area_stop_association_table",
    Base.metadata,
    Column("stop_id", ForeignKey("_stops.id")),
    Column("stop_area_id", ForeignKey("stop_areas.id")),
)


class _Stop(Base):

    db = db

    id = Column(BigInteger, primary_key=True)
    kind = Column(String)

    name = Column(String, nullable=False, index=True)
    town_name = Column(String, nullable=False)
    postal_region = Column(String, nullable=False)
    xepsg2154 = Column(BigInteger, nullable=False)
    yepsg2154 = Column(BigInteger, nullable=False)
    version = Column(String, nullable=False)
    created_ts = Column(BigInteger)
    changed_ts = Column(BigInteger, nullable=False)
    lines: Mapped[list["Line"]] = relationship(
        "Line",
        secondary="line_stop_association_table",
        back_populates="stops",
        # lazy="joined",
        lazy="selectin",
    )
    areas: Mapped[list["StopArea"]] = relationship(
        "StopArea", secondary=stop_area_stop_association_table, back_populates="stops"
    )

    __tablename__ = "_stops"
    __mapper_args__ = {"polymorphic_identity": "_stops", "polymorphic_on": kind}
    __table_args__ = (
        # To optimize the ilike requests
        Index(
            "name_idx_gin",
            name,
            postgresql_ops={"name": "gin_trgm_ops"},
            postgresql_using="gin",
        ),
    )

    # TODO: Test https://www.cybertec-postgresql.com/en/postgresql-more-performance-for-like-and-ilike-statements/
    # TODO: Should be able to remove with_polymorphic ?
    @classmethod
    async def get_by_name(cls, name: str) -> list[Self]:
        stop_stop_area = with_polymorphic(_Stop, [Stop, StopArea])
        stmt = (
            select(stop_stop_area)
            .where(stop_stop_area.name.ilike(f"%{name}%"))
            .options(
                selectinload(stop_stop_area.areas),
                selectinload(stop_stop_area.lines),
            )
        )
        res = await cls.db.session.execute(stmt)
        return res.scalars()


class Stop(_Stop):

    id = Column(BigInteger, ForeignKey("_stops.id"), primary_key=True)

    latitude = Column(Float, nullable=False)
    longitude = Column(Float, nullable=False)
    transport_mode = Column(Enum(TransportMode), nullable=False)
    accessibility = Column(Enum(IdfmState), nullable=False)
    visual_signs_available = Column(Enum(IdfmState), nullable=False)
    audible_signs_available = Column(Enum(IdfmState), nullable=False)
    record_id = Column(String, nullable=False)
    record_ts = Column(BigInteger, nullable=False)

    __tablename__ = "stops"
    __mapper_args__ = {"polymorphic_identity": "stops", "polymorphic_load": "inline"}


class StopArea(_Stop):

    id = Column(BigInteger, ForeignKey("_stops.id"), primary_key=True)

    type = Column(Enum(StopAreaType), nullable=False)
    stops: Mapped[list[_Stop]] = relationship(
        _Stop,
        secondary=stop_area_stop_association_table,
        back_populates="areas",
        lazy="selectin",
        # lazy="joined",
    )

    __tablename__ = "stop_areas"
    __mapper_args__ = {"polymorphic_identity": "stop_areas", "polymorphic_load": "inline"}

    @classmethod
    async def add_stops(cls, stop_area_to_stop_ids: Iterable[tuple[str, str]]) -> int:
        stop_area_ids, stop_ids = set(), set()
        for stop_area_id, stop_id in stop_area_to_stop_ids:
            stop_area_ids.add(stop_area_id)
            stop_ids.add(stop_id)

        res = await cls.db.session.execute(
            select(StopArea)
            .where(StopArea.id.in_(stop_area_ids))
            .options(selectinload(StopArea.stops))
        )
        stop_areas = {stop_area.id: stop_area for stop_area in res.scalars()}

        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 stop_area_id, stop_id in stop_area_to_stop_ids:
            if (stop_area := stop_areas.get(stop_area_id)) is not None:
                if (stop := stops.get(stop_id)) is not None:
                    stop_area.stops.append(stop)
                    found += 1
                else:
                    print(f"No stop found for {stop_id} id")
            else:
                print(f"No stop area found for {stop_area_id}")

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