from collections.abc import AsyncGenerator, Generator
from datetime import datetime, timezone
import uuid

import pytest
import pytest_asyncio
from sqlalchemy import DateTime, ForeignKey, Integer, String, text
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
from testcontainers.postgres import PostgresContainer


class Base(DeclarativeBase):
    pass


class User(Base):
    __tablename__ = "users"

    id: Mapped[int] = mapped_column(Integer, primary_key=True)
    name: Mapped[str] = mapped_column(String, nullable=False)
    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)


class Post(Base):
    __tablename__ = "posts"

    id: Mapped[int] = mapped_column(Integer, primary_key=True)
    user_id: Mapped[int] = mapped_column(Integer, ForeignKey("users.id"), nullable=False)
    title: Mapped[str] = mapped_column(String, nullable=False)


@pytest.fixture(scope="session")
def pg_container() -> Generator[PostgresContainer]:
    with PostgresContainer("postgres:16-alpine") as pg:
        yield pg


@pytest_asyncio.fixture(params=["sqlite", "postgres"])
async def async_engine(request: pytest.FixtureRequest, pg_container: PostgresContainer) -> AsyncGenerator[AsyncEngine]:
    if request.param == "sqlite":
        engine = create_async_engine("sqlite+aiosqlite:///:memory:")
        async with engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)
        yield engine
        await engine.dispose()
    else:
        pg_url = (
            f"postgresql+asyncpg://{pg_container.username}:{pg_container.password}"
            f"@{pg_container.get_container_host_ip()}"
            f":{pg_container.get_exposed_port(5432)}/{pg_container.dbname}"
        )

        schema = f"t{uuid.uuid4().hex}"
        engine = create_async_engine(
            pg_url,
            connect_args={"server_settings": {"search_path": schema}},
        )
        async with engine.execution_options(isolation_level="AUTOCOMMIT").connect() as conn:
            await conn.execute(text(f'CREATE SCHEMA "{schema}"'))
            await conn.run_sync(Base.metadata.create_all)
        yield engine
        async with engine.execution_options(isolation_level="AUTOCOMMIT").connect() as conn:
            await conn.execute(text(f'DROP SCHEMA "{schema}" CASCADE'))


@pytest_asyncio.fixture
async def async_session(async_engine: AsyncEngine) -> AsyncGenerator[AsyncSession]:
    factory = async_sessionmaker(async_engine, expire_on_commit=False)
    async with factory() as session:
        yield session
        await session.rollback()


@pytest_asyncio.fixture
async def seeded_users(async_engine: AsyncEngine) -> list[User]:
    start = datetime(2024, 1, 1, 0, 0, 0, tzinfo=timezone.utc)
    users = [
        User(
            id=i,
            name=f"User {i:02d}",
            created_at=start.replace(hour=i % 24, minute=0, second=0),
        )
        for i in range(1, 21)
    ]
    factory = async_sessionmaker(async_engine, expire_on_commit=False)
    async with factory() as session:
        session.add_all(users)
        await session.commit()
    return users


@pytest_asyncio.fixture
async def seeded_posts(async_engine: AsyncEngine, seeded_users: list[User]) -> list[Post]:
    # 3 posts per user; user N gets post IDs (3N-2), (3N-1), 3N
    posts = [
        Post(
            id=((user_id - 1) * 3) + offset,
            user_id=user_id,
            title=f"Post {((user_id - 1) * 3) + offset:03d}",
        )
        for user_id in range(1, 21)
        for offset in range(1, 4)
    ]
    factory = async_sessionmaker(async_engine, expire_on_commit=False)
    async with factory() as session:
        session.add_all(posts)
        await session.commit()
    return posts
