from contextlib import contextmanager
from functools import lru_cache
from typing import Generator
from api_controllers.postgres.config import postgres_configs

from sqlalchemy import create_engine
from sqlalchemy.orm import declarative_base, sessionmaker, scoped_session, Session


# Configure the database engine with proper pooling
engine = create_engine(
    postgres_configs.url,
    pool_pre_ping=True,
    pool_size=10,  # Reduced from 20 to better match your CPU cores
    max_overflow=5,  # Reduced from 10 to prevent too many connections
    pool_recycle=600,  # Keep as is
    pool_timeout=30,  # Keep as is
    echo=False,  # Consider setting to False in production
)


Base = declarative_base()


# Create a cached session factory
@lru_cache()
def get_session_factory() -> scoped_session:
    """Create a thread-safe session factory."""
    session_local = sessionmaker(
        bind=engine,
        autocommit=False,
        autoflush=False,
        expire_on_commit=True,  # Prevent expired object issues
    )
    return scoped_session(session_local)


# Get database session with proper connection management
@contextmanager
def get_db() -> Generator[Session, None, None]:
    """Get database session with proper connection management.

    This context manager ensures:
    - Proper connection pooling
    - Session cleanup
    - Connection return to pool
    - Thread safety

    Yields:
        Session: SQLAlchemy session object
    """

    session_factory = get_session_factory()
    session = session_factory()
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()
        session_factory.remove()  # Clean up the session from the registry
