"""
Demonstrate relationships and joined reads using SQLModel with generic kit.

Run: uv run python examples/03_relationships.py
"""

from typing import List, Optional
from sqlalchemy.orm import selectinload
from sqlmodel import SQLModel, select, Field, Relationship

from sqlmodel_pg_kit import create_all
from sqlmodel_pg_kit.db import get_session


class Team(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    name: str = Field(index=True)
    heroes: List["Hero"] = Relationship(back_populates="team")


class Hero(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    name: str = Field(index=True)
    age: Optional[int] = None
    team_id: Optional[int] = Field(default=None, foreign_key="team.id")
    team: Optional[Team] = Relationship(back_populates="heroes")


def main():
    create_all()

    with get_session() as s:
        # Clean
        s.execute(Hero.__table__.delete())
        s.execute(Team.__table__.delete())
        s.commit()

        # Create Team and Heroes
        t = Team(name="Avengers")
        s.add(t)
        s.commit()
        s.refresh(t)

        h1 = Hero(name="Thor", age=1500, team_id=t.id)
        h2 = Hero(name="Hulk", age=49, team_id=t.id)
        s.add(h1)
        s.add(h2)
        s.commit()

        # Query heroes and eager-load team via selectinload
        stmt = select(Hero).options(selectinload(Hero.team)).order_by(Hero.id.asc())
        heroes: List[Hero] = s.exec(stmt).all()
        print([(h.name, h.team.name if h.team else None) for h in heroes])


if __name__ == "__main__":
    main()