from typing import List, Optional

from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select

class HeroTeamLink(SQLModel, table=True):

team_id: Optional[int] = Field(

default=None, foreign_key="team.id", primary_key=True

)

hero_id: Optional[int] = Field(

default=None, foreign_key="hero.id", primary_key=True

)

is_training: bool = False

team: "Team" = Relationship(back_populates="hero_links")

hero: "Hero" = Relationship(back_populates="team_links")

class Team(SQLModel, table=True):

id: Optional[int] = Field(default=None, primary_key=True)

name: str = Field(index=True)

headquarters: str

hero_links: List[HeroTeamLink] = Relationship(back_populates="team")

class Hero(SQLModel, table=True):

id: Optional[int] = Field(default=None, primary_key=True)

name: str = Field(index=True)

secret_name: str

age: Optional[int] = Field(default=None, index=True)

team_links: List[HeroTeamLink] = Relationship(back_populates="hero")

sqlite_file_name = "database.db"

sqlite_url = f"sqlite:///{sqlite_file_name}"

engine = create_engine(sqlite_url, echo=True)

def create_db_and_tables():

SQLModel.metadata.create_all(engine)

def create_heroes():

with Session(engine) as session:

team_preventers = Team(name="Preventers", headquarters="Sharp Tower")

team_z_force = Team(name="Z-Force", headquarters="Sister Margaret’s Bar")

hero_deadpond = Hero(

name="Deadpond",

secret_name="Dive Wilson",

)

hero_rusty_man = Hero(

name="Rusty-Man",

secret_name="Tommy Sharp",

age=48,

)

hero_spider_boy = Hero(

name="Spider-Boy",

secret_name="Pedro Parqueador",

)

deadpond_team_z_link = HeroTeamLink(team=team_z_force, hero=hero_deadpond)

deadpond_preventers_link = HeroTeamLink(

team=team_preventers, hero=hero_deadpond, is_training=True

)

spider_boy_preventers_link = HeroTeamLink(

team=team_preventers, hero=hero_spider_boy, is_training=True

)

rusty_man_preventers_link = HeroTeamLink(

team=team_preventers, hero=hero_rusty_man

)

session.add(deadpond_team_z_link)

session.add(deadpond_preventers_link)

session.add(spider_boy_preventers_link)

session.add(rusty_man_preventers_link)

session.commit()

for link in team_z_force.hero_links:

print("Z-Force hero:", link.hero, "is training:", link.is_training)

for link in team_preventers.hero_links:

print("Preventers hero:", link.hero, "is training:", link.is_training)

def update_heroes():

with Session(engine) as session:

hero_spider_boy = session.exec(

select(Hero).where(Hero.name == "Spider-Boy")

).one()

team_z_force = session.exec(select(Team).where(Team.name == "Z-Force")).one()

spider_boy_z_force_link = HeroTeamLink(

team=team_z_force, hero=hero_spider_boy, is_training=True

)

team_z_force.hero_links.append(spider_boy_z_force_link)

session.add(team_z_force)

session.commit()

print("Updated Spider-Boy's Teams:", hero_spider_boy.team_links)

print("Z-Force heroes:", team_z_force.hero_links)

for link in hero_spider_boy.team_links:

if link.team.name == "Preventers":

link.is_training = False

session.add(hero_spider_boy)

session.commit()

for link in hero_spider_boy.team_links:

print("Spider-Boy team:", link.team, "is training:", link.is_training)

def main():

create_db_and_tables()

create_heroes()

update_heroes()

if __name__ == "__main__":

main()