from collections.abc import Generator
from contextlib import contextmanager

from sqlalchemy import create_engine
from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session, declarative_base, sessionmaker

from app.core.config import get_settings


Base = declarative_base()

_engine_cache: dict[str, Engine] = {}
_session_factory_cache: dict[str, sessionmaker[Session]] = {}
_prepared_urls: set[str] = set()


def _connect_args(database_url: str) -> dict[str, object]:
    if database_url.startswith("sqlite"):
        return {"check_same_thread": False}
    return {}


def get_engine(database_url: str | None = None) -> Engine:
    target_url = database_url or get_settings().database_url
    engine = _engine_cache.get(target_url)
    if engine is None:
        engine = create_engine(
            target_url,
            future=True,
            pool_pre_ping=not target_url.startswith("sqlite"),
            connect_args=_connect_args(target_url),
        )
        _engine_cache[target_url] = engine

    return engine


def prepare_database(database_url: str | None = None) -> Engine:
    target_url = database_url or get_settings().database_url
    engine = get_engine(target_url)
    if target_url not in _prepared_urls:
        Base.metadata.create_all(bind=engine)
        _prepared_urls.add(target_url)
    return engine


def _get_session_factory(database_url: str | None = None) -> sessionmaker[Session]:
    target_url = database_url or get_settings().database_url
    session_factory = _session_factory_cache.get(target_url)
    if session_factory is None:
        session_factory = sessionmaker(bind=prepare_database(target_url), autoflush=False, autocommit=False, future=True)
        _session_factory_cache[target_url] = session_factory
    return session_factory


def open_db_session(database_url: str | None = None) -> Generator[Session, None, None]:
    session_factory = _get_session_factory(database_url)
    db = session_factory()
    try:
        yield db
    finally:
        db.close()


@contextmanager
def db_session_context(database_url: str | None = None) -> Generator[Session, None, None]:
    session_factory = _get_session_factory(database_url)
    db = session_factory()
    try:
        yield db
    finally:
        db.close()