from typing import List, Optional

from fastapi import Depends, FastAPI, HTTPException, Query
from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select


class TeamBase(SQLModel):
    name: str = Field(index=True)
    headquarters: str


class Team(TeamBase, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)

    heroes: List["Hero"] = Relationship(back_populates="team")


class TeamCreate(TeamBase):
    pass


class TeamRead(TeamBase):
    id: int


class TeamUpdate(SQLModel):
    id: Optional[int] = None
    name: Optional[str] = None
    headquarters: Optional[str] = None


class HeroBase(SQLModel):
    name: str = Field(index=True)
    secret_name: str
    age: Optional[int] = Field(default=None, index=True)

    team_id: Optional[int] = Field(default=None, foreign_key="team.id")


class Hero(HeroBase, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)

    team: Optional[Team] = Relationship(back_populates="heroes")


class HeroRead(HeroBase):
    id: int


class HeroCreate(HeroBase):
    pass


class HeroUpdate(SQLModel):
    name: Optional[str] = None
    secret_name: Optional[str] = None
    age: Optional[int] = None
    team_id: Optional[int] = None


sqlite_file_name = "database.db"
sqlite_url = f"sqlite:///{sqlite_file_name}"

connect_args = {"check_same_thread": False}
engine = create_engine(sqlite_url, echo=True, connect_args=connect_args)


def create_db_and_tables():
    SQLModel.metadata.create_all(engine)


def get_session():
    with Session(engine) as session:
        yield session


app = FastAPI()


@app.on_event("startup")
def on_startup():
    create_db_and_tables()


@app.post("/heroes/", response_model=HeroRead)
def create_hero(*, session: Session = Depends(get_session), hero: HeroCreate):
    db_hero = Hero.from_orm(hero)
    session.add(db_hero)
    session.commit()
    session.refresh(db_hero)
    return db_hero


@app.get("/heroes/", response_model=List[HeroRead])
def read_heroes(
    *,
    session: Session = Depends(get_session),
    offset: int = 0,
    limit: int = Query(default=100, lte=100),
):
    heroes = session.exec(select(Hero).offset(offset).limit(limit)).all()
    return heroes


@app.get("/heroes/{hero_id}", response_model=HeroRead)
def read_hero(*, session: Session = Depends(get_session), hero_id: int):
    hero = session.get(Hero, hero_id)
    if not hero:
        raise HTTPException(status_code=404, detail="Hero not found")
    return hero


@app.patch("/heroes/{hero_id}", response_model=HeroRead)
def update_hero(
    *, session: Session = Depends(get_session), hero_id: int, hero: HeroUpdate
):
    db_hero = session.get(Hero, hero_id)
    if not db_hero:
        raise HTTPException(status_code=404, detail="Hero not found")
    hero_data = hero.dict(exclude_unset=True)
    for key, value in hero_data.items():
        setattr(db_hero, key, value)
    session.add(db_hero)
    session.commit()
    session.refresh(db_hero)
    return db_hero


@app.delete("/heroes/{hero_id}")
def delete_hero(*, session: Session = Depends(get_session), hero_id: int):

    hero = session.get(Hero, hero_id)
    if not hero:
        raise HTTPException(status_code=404, detail="Hero not found")
    session.delete(hero)
    session.commit()
    return {"ok": True}


@app.post("/teams/", response_model=TeamRead)
def create_team(*, session: Session = Depends(get_session), team: TeamCreate):
    db_team = Team.from_orm(team)
    session.add(db_team)
    session.commit()
    session.refresh(db_team)
    return db_team


@app.get("/teams/", response_model=List[TeamRead])
def read_teams(
    *,
    session: Session = Depends(get_session),
    offset: int = 0,
    limit: int = Query(default=100, lte=100),
):
    teams = session.exec(select(Team).offset(offset).limit(limit)).all()
    return teams


@app.get("/teams/{team_id}", response_model=TeamRead)
def read_team(*, team_id: int, session: Session = Depends(get_session)):
    team = session.get(Team, team_id)
    if not team:
        raise HTTPException(status_code=404, detail="Team not found")
    return team


@app.patch("/teams/{team_id}", response_model=TeamRead)
def update_team(
    *,
    session: Session = Depends(get_session),
    team_id: int,
    team: TeamUpdate,
):
    db_team = session.get(Team, team_id)
    if not db_team:
        raise HTTPException(status_code=404, detail="Team not found")
    team_data = team.dict(exclude_unset=True)
    for key, value in team_data.items():
        setattr(db_team, key, value)
    session.add(db_team)
    session.commit()
    session.refresh(db_team)
    return db_team


@app.delete("/teams/{team_id}")
def delete_team(*, session: Session = Depends(get_session), team_id: int):
    team = session.get(Team, team_id)
    if not team:
        raise HTTPException(status_code=404, detail="Team not found")
    session.delete(team)
    session.commit()
    return {"ok": True}

sqlmodel.Relationship sqlmodel.create_engine sqlmodel.Session sqlmodel.SQLModel.metadata.create_all sqlmodel.Field fastapi.HTTPException sqlmodel.select fastapi.Query fastapi.Depends fastapi.FastAPI
[((1384, 1447), 'sqlmodel.create_engine', 'create_engine', (['sqlite_url'], {'echo': '(True)', 'connect_args': 'connect_args'}), '(sqlite_url, echo=True, connect_args=connect_args)\n', (1397, 1447), False, 'from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select\n'), ((1607, 1616), 'fastapi.FastAPI', 'FastAPI', ([], {}), '()\n', (1614, 1616), False, 'from fastapi import Depends, FastAPI, HTTPException, Query\n'), ((221, 238), 'sqlmodel.Field', 'Field', ([], {'index': '(True)'}), '(index=True)\n', (226, 238), False, 'from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select\n'), ((321, 358), 'sqlmodel.Field', 'Field', ([], {'default': 'None', 'primary_key': '(True)'}), '(default=None, primary_key=True)\n', (326, 358), False, 'from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select\n'), ((387, 422), 'sqlmodel.Relationship', 'Relationship', ([], {'back_populates': '"""team"""'}), "(back_populates='team')\n", (399, 422), False, 'from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select\n'), ((675, 692), 'sqlmodel.Field', 'Field', ([], {'index': '(True)'}), '(index=True)\n', (680, 692), False, 'from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select\n'), ((739, 770), 'sqlmodel.Field', 'Field', ([], {'default': 'None', 'index': '(True)'}), '(default=None, index=True)\n', (744, 770), False, 'from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select\n'), ((801, 843), 'sqlmodel.Field', 'Field', ([], {'default': 'None', 'foreign_key': '"""team.id"""'}), "(default=None, foreign_key='team.id')\n", (806, 843), False, 'from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select\n'), ((904, 941), 'sqlmodel.Field', 'Field', ([], {'default': 'None', 'primary_key': '(True)'}), '(default=None, primary_key=True)\n', (909, 941), False, 'from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select\n'), ((970, 1007), 'sqlmodel.Relationship', 'Relationship', ([], {'back_populates': '"""heroes"""'}), "(back_populates='heroes')\n", (982, 1007), False, 'from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select\n'), ((1482, 1518), 'sqlmodel.SQLModel.metadata.create_all', 'SQLModel.metadata.create_all', (['engine'], {}), '(engine)\n', (1510, 1518), False, 'from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select\n'), ((1776, 1796), 'fastapi.Depends', 'Depends', (['get_session'], {}), '(get_session)\n', (1783, 1796), False, 'from fastapi import Depends, FastAPI, HTTPException, Query\n'), ((2046, 2066), 'fastapi.Depends', 'Depends', (['get_session'], {}), '(get_session)\n', (2053, 2066), False, 'from fastapi import Depends, FastAPI, HTTPException, Query\n'), ((2106, 2133), 'fastapi.Query', 'Query', ([], {'default': '(100)', 'lte': '(100)'}), '(default=100, lte=100)\n', (2111, 2133), False, 'from fastapi import Depends, FastAPI, HTTPException, Query\n'), ((2323, 2343), 'fastapi.Depends', 'Depends', (['get_session'], {}), '(get_session)\n', (2330, 2343), False, 'from fastapi import Depends, FastAPI, HTTPException, Query\n'), ((2603, 2623), 'fastapi.Depends', 'Depends', (['get_session'], {}), '(get_session)\n', (2610, 2623), False, 'from fastapi import Depends, FastAPI, HTTPException, Query\n'), ((3081, 3101), 'fastapi.Depends', 'Depends', (['get_session'], {}), '(get_session)\n', (3088, 3101), False, 'from fastapi import Depends, FastAPI, HTTPException, Query\n'), ((3400, 3420), 'fastapi.Depends', 'Depends', (['get_session'], {}), '(get_session)\n', (3407, 3420), False, 'from fastapi import Depends, FastAPI, HTTPException, Query\n'), ((3668, 3688), 'fastapi.Depends', 'Depends', (['get_session'], {}), '(get_session)\n', (3675, 3688), False, 'from fastapi import Depends, FastAPI, HTTPException, Query\n'), ((3728, 3755), 'fastapi.Query', 'Query', ([], {'default': '(100)', 'lte': '(100)'}), '(default=100, lte=100)\n', (3733, 3755), False, 'from fastapi import Depends, FastAPI, HTTPException, Query\n'), ((3956, 3976), 'fastapi.Depends', 'Depends', (['get_session'], {}), '(get_session)\n', (3963, 3976), False, 'from fastapi import Depends, FastAPI, HTTPException, Query\n'), ((4225, 4245), 'fastapi.Depends', 'Depends', (['get_session'], {}), '(get_session)\n', (4232, 4245), False, 'from fastapi import Depends, FastAPI, HTTPException, Query\n'), ((4711, 4731), 'fastapi.Depends', 'Depends', (['get_session'], {}), '(get_session)\n', (4718, 4731), False, 'from fastapi import Depends, FastAPI, HTTPException, Query\n'), ((1549, 1564), 'sqlmodel.Session', 'Session', (['engine'], {}), '(engine)\n', (1556, 1564), False, 'from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select\n'), ((2429, 2484), 'fastapi.HTTPException', 'HTTPException', ([], {'status_code': '(404)', 'detail': '"""Hero not found"""'}), "(status_code=404, detail='Hero not found')\n", (2442, 2484), False, 'from fastapi import Depends, FastAPI, HTTPException, Query\n'), ((2734, 2789), 'fastapi.HTTPException', 'HTTPException', ([], {'status_code': '(404)', 'detail': '"""Hero not found"""'}), "(status_code=404, detail='Hero not found')\n", (2747, 2789), False, 'from fastapi import Depends, FastAPI, HTTPException, Query\n'), ((3188, 3243), 'fastapi.HTTPException', 'HTTPException', ([], {'status_code': '(404)', 'detail': '"""Hero not found"""'}), "(status_code=404, detail='Hero not found')\n", (3201, 3243), False, 'from fastapi import Depends, FastAPI, HTTPException, Query\n'), ((4048, 4103), 'fastapi.HTTPException', 'HTTPException', ([], {'status_code': '(404)', 'detail': '"""Team not found"""'}), "(status_code=404, detail='Team not found')\n", (4061, 4103), False, 'from fastapi import Depends, FastAPI, HTTPException, Query\n'), ((4365, 4420), 'fastapi.HTTPException', 'HTTPException', ([], {'status_code': '(404)', 'detail': '"""Team not found"""'}), "(status_code=404, detail='Team not found')\n", (4378, 4420), False, 'from fastapi import Depends, FastAPI, HTTPException, Query\n'), ((4817, 4872), 'fastapi.HTTPException', 'HTTPException', ([], {'status_code': '(404)', 'detail': '"""Team not found"""'}), "(status_code=404, detail='Team not found')\n", (4830, 4872), False, 'from fastapi import Depends, FastAPI, HTTPException, Query\n'), ((2164, 2176), 'sqlmodel.select', 'select', (['Hero'], {}), '(Hero)\n', (2170, 2176), False, 'from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select\n'), ((3785, 3797), 'sqlmodel.select', 'select', (['Team'], {}), '(Team)\n', (3791, 3797), False, 'from sqlmodel import Field, Relationship, Session, SQLModel, create_engine, select\n')]