from collections.abc import AsyncIterator
from contextlib import asynccontextmanager

import uvicorn
from fastapi import FastAPI
from sqlalchemy.ext.asyncio import create_async_engine
from sqlmodel import Field, Relationship, SQLModel, select
from sqlmodel.ext.asyncio.session import AsyncSession


class Parent(SQLModel, table=True):
    __tablename__ = "parent"

    id: int | None = Field(default=None, primary_key=True)
    name: str
    intermediates: list["Intermediate"] = Relationship(sa_relationship_kwargs={"lazy": "selectin"})


class ParentResponseSchema(SQLModel):
    id: int
    name: str
    intermediates: list["IntermediateResponseSchema"]


class ParentRequestSchema(SQLModel):
    name: str
    intermediates: list[str]


class Intermediate(SQLModel, table=True):
    __tablename__ = "intermediate"

    id: int | None = Field(default=None, primary_key=True)
    name: str
    parent_id: int = Field(foreign_key="parent.id")
    children: list["Child"] = Relationship(sa_relationship_kwargs={"lazy": "selectin"})


class IntermediateResponseSchema(SQLModel):
    id: int
    name: str
    parent_id: int
    children: list["ChildResponseSchema"]


class Child(SQLModel, table=True):
    __tablename__ = "child"

    id: int | None = Field(default=None, primary_key=True)
    name: str
    intermediate_id: int = Field(foreign_key="intermediate.id")


class ChildResponseSchema(SQLModel):
    id: int
    name: str


engine = create_async_engine("postgresql+asyncpg://postgres:admin123@localhost:5432/postgres")


@asynccontextmanager
async def lifespan(_: FastAPI) -> AsyncIterator[None]:
    async with engine.begin() as conn:
        await conn.run_sync(SQLModel.metadata.create_all)
    yield
    # async with engine.begin() as conn:
    #     await conn.run_sync(SQLModel.metadata.drop_all)


app = FastAPI(lifespan=lifespan)


@app.post("/parent", response_model=ParentResponseSchema)
async def create(body: ParentRequestSchema):
    async with AsyncSession(engine) as session:
        query = select(Parent).where(Parent.name == body.name)
        result = await session.exec(query)
        parent = result.first()
        if parent is None:
            parent = Parent(name=body.name)
            session.add(parent)
            await session.commit()
            await session.refresh(parent)

        for intermediate_name in body.intermediates:
            query = select(Intermediate).where(Intermediate.name == intermediate_name)
            result = await session.exec(query)
            intermediate = result.first()
            if intermediate is None:
                intermediate = Intermediate(name=intermediate_name, parent_id=parent.id)
                session.add(intermediate)
                await session.commit()
                await session.refresh(intermediate)

            parent.intermediates.append(intermediate)
    return parent


@app.get("/parent", response_model=ParentResponseSchema)
async def alive():
    async with AsyncSession(engine) as session:
        parent = await session.get(Parent, 1)
    return parent


if __name__ == "__main__":
    uvicorn.run("main:app", host="localhost", port=8080, reload=True)