from __future__ import annotations
import os
from contextlib import asynccontextmanager
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession
from quart import Quart, g

DATABASE_URL = (
    os.getenv("DATABASE_URL_ASYNC")
    or os.getenv("DATABASE_URL")
    or "postgresql+asyncpg://localhost/coop"
)

_engine = create_async_engine(
    DATABASE_URL,
    future=True,
    echo=False,
    pool_pre_ping=True,
    pool_size=-1  # ned to look at this!!!
)

_Session = async_sessionmaker(
    bind=_engine,
    class_=AsyncSession,
    expire_on_commit=False,
)

@asynccontextmanager
async def get_session():
    """Always create a fresh AsyncSession for this block."""
    sess = _Session()
    try:
        yield sess
    finally:
        await sess.close()



def register_db(app: Quart):

  @app.before_request
  async def open_session():
      g.s = _Session()
      g.tx = await g.s.begin()
      g.had_error = False

  @app.after_request
  async def maybe_commit(response):
      # Runs BEFORE bytes are sent.
      if not g.had_error and 200 <= response.status_code < 400:
          try:
              if hasattr(g, "tx"):
                  await g.tx.commit()
          except Exception as e:
              print(f'commit failed {e}')
              if hasattr(g, "tx"):
                  await g.tx.rollback()
              from quart import make_response
              return await make_response("Commit failed", 500)
      return response

  @app.teardown_request
  async def finish(exc):
      try:
          # If an exception occurred OR we didn't commit (still in txn), roll back.
          if hasattr(g, "s"):
              if exc is not None or g.s.in_transaction():
                  if hasattr(g, "tx"):
                      await g.tx.rollback()
      finally:
          if hasattr(g, "s"):
              await g.s.close()

  @app.errorhandler(Exception)
  async def mark_error(e):
      g.had_error = True
      raise