from typing import AsyncIterator
from fastapi import Request
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
from sqlalchemy import String


class Base(DeclarativeBase):
    pass


# User Thread Mapping Table
class UserThreadMapping(Base):
    """Mapping of users to their chat threads.

    Used to keep track of which user owns which chat thread.
    Also stores meta data like thread name.
    """

    __tablename__ = "userThreads"
    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    userId: Mapped[int] = mapped_column(nullable=False)
    threadId: Mapped[str] = mapped_column(String(255), unique=True, nullable=False)
    threadName: Mapped[str] = mapped_column(String(255), nullable=False)


# Dependency that pulls the sessionmaker off app.state
#  This is set in app.py on startup in @asynccontextmanager
async def get_session(request: Request) -> AsyncIterator[AsyncSession]:
    SessionLocal: async_sessionmaker[AsyncSession] = (
        request.app.state.checkpoint_sessionmaker
    )
    async with SessionLocal() as session:
        yield session


# Optional helper to init tables at startup (demo only)
async def init_models(engine) -> None:
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)