"""Demonstrate that shared-cache URI mode fixes the concurrent-session
data-loss issue described in
https://github.com/sqlalchemy/sqlalchemy/discussions/13428

The script runs the reproducer from that discussion under two
configurations:

1. StaticPool (the old default for :memory:) — concurrent sessions
   share one DBAPI connection and a ROLLBACK in one session silently
   discards uncommitted work from another.

2. Shared-cache URI ("file::memory:?cache=shared&uri=true") — each
   session gets its own DBAPI connection with independent transaction
   state, backed by the same in-memory database.  Concurrent sessions
   work correctly.
"""

import asyncio

from sqlalchemy import Column, Integer, Text, create_engine, select
from sqlalchemy.ext.asyncio import (
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)
from sqlalchemy.orm import DeclarativeBase, Session
from sqlalchemy.pool import StaticPool


class Base(DeclarativeBase):
    pass


class Row(Base):
    __tablename__ = "rows"
    id = Column(Integer, primary_key=True)
    v = Column(Text)


def test_sync_staticpool_loses_data():
    """Variant A from the discussion: sync, StaticPool, single thread.

    Two sessions share one DBAPI connection.  s2's rollback wipes
    s1's flushed INSERT; s1.commit() then commits nothing.
    """
    eng = create_engine(
        "sqlite://",
        poolclass=StaticPool,
        connect_args={"check_same_thread": False},
    )
    Base.metadata.create_all(eng)

    s1 = Session(eng)
    s1.add(Row(v="s1"))
    s1.flush()

    s2 = Session(eng)
    s2.execute(select(Row))
    s2.rollback()

    s1.commit()

    with Session(eng) as s3:
        rows = s3.execute(select(Row)).scalars().all()

    s1.close()
    s2.close()
    eng.dispose()

    if len(rows) == 0:
        print(
            "SYNC StaticPool: CONFIRMED data loss — "
            "0 rows after commit (expected 1)"
        )
        return True
    else:
        print(f"SYNC StaticPool: got {len(rows)} rows (unexpected)")
        return False


def test_sync_shared_cache_preserves_data():
    """Sync variant using shared-cache URI.

    Each session gets its own DBAPI connection with independent
    transaction state.  s1 flushes and commits an INSERT; s2 opens
    its own session, does work, and rolls back — s1's committed
    row is unaffected.
    """
    eng = create_engine(
        "sqlite://github.com/file::memory:?cache=shared&uri=true"
    )
    Base.metadata.create_all(eng)

    s1 = Session(eng)
    s1.add(Row(v="s1"))
    s1.commit()

    s2 = Session(eng)
    s2.add(Row(v="s2"))
    s2.flush()
    s2.rollback()

    with Session(eng) as s3:
        rows = s3.execute(select(Row)).scalars().all()

    s1.close()
    s2.close()
    eng.dispose()

    if len(rows) == 1:
        print(
            "SYNC shared-cache: OK — "
            "1 row preserved, s2 rollback did not affect s1"
        )
        return True
    else:
        print(
            f"SYNC shared-cache: UNEXPECTED — "
            f"got {len(rows)} rows (expected 1)"
        )
        return False


async def test_async_staticpool_loses_data():
    """Variant B from the discussion: async, default :memory: engine.

    Two concurrent tasks each get a session; the reader's implicit
    rollback on exit wipes the writer's flushed INSERT.
    """
    eng = create_async_engine("sqlite+aiosqlite://")
    async with eng.begin() as c:
        await c.run_sync(Base.metadata.create_all)

    sf = async_sessionmaker(eng, class_=AsyncSession)
    barrier = asyncio.Barrier(2)

    async def writer():
        async with sf() as s:
            s.add(Row(v="w"))
            await s.flush()
            await barrier.wait()
            await asyncio.sleep(0.05)
            await s.commit()

    async def reader():
        async with sf() as s:
            await s.execute(select(Row))
            await barrier.wait()

    await asyncio.gather(writer(), reader())

    async with sf() as s:
        rows = (await s.execute(select(Row))).scalars().all()

    await eng.dispose()

    if len(rows) == 0:
        print(
            "ASYNC StaticPool: CONFIRMED data loss — "
            "0 rows after commit (expected 1)"
        )
        return True
    else:
        print(f"ASYNC StaticPool: got {len(rows)} rows (unexpected)")
        return False


async def test_async_shared_cache_preserves_data():
    """Async variant using shared-cache URI with create_async_engine.

    Each coroutine gets its own DBAPI connection with independent
    transaction state.  The writer commits, then the reader opens
    a session and rolls back — the committed row is unaffected.
    """
    eng = create_async_engine(
        "sqlite+aiosqlite://github.com/file::memory:?cache=shared&uri=true"
    )
    async with eng.begin() as c:
        await c.run_sync(Base.metadata.create_all)

    sf = async_sessionmaker(eng, class_=AsyncSession)
    barrier = asyncio.Barrier(2)

    async def writer():
        async with sf() as s:
            s.add(Row(v="w"))
            await s.commit()
            await barrier.wait()

    async def reader():
        async with sf() as s:
            await barrier.wait()
            s.add(Row(v="r"))
            await s.flush()
            await s.rollback()

    await asyncio.gather(writer(), reader())

    async with sf() as s:
        rows = (await s.execute(select(Row))).scalars().all()

    await eng.dispose()

    if len(rows) == 1:
        print(
            "ASYNC shared-cache: OK — "
            "1 row preserved, reader rollback did not affect writer"
        )
        return True
    else:
        print(
            f"ASYNC shared-cache: UNEXPECTED — "
            f"got {len(rows)} rows (expected 1)"
        )
        return False


if __name__ == "__main__":
    results = []

    print("=" * 60)
    print("StaticPool data-loss reproducer + shared-cache fix")
    print("=" * 60)
    print()

    print("--- Sync tests ---")
    results.append(test_sync_staticpool_loses_data())
    results.append(test_sync_shared_cache_preserves_data())

    print()
    print("--- Async tests ---")
    results.append(asyncio.run(test_async_staticpool_loses_data()))
    results.append(asyncio.run(test_async_shared_cache_preserves_data()))

    print()
    if all(results):
        print("All 4 tests behaved as expected.")
    else:
        print("Some tests did not behave as expected.")