import os
from collections.abc import Generator
from contextlib import contextmanager

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

DATABASE_URL_DEFAULT = "sqlite:///finanz.db"

# Optional: allow an external engine (e.g. for testing with in-memory SQLite)
_engine: Engine | None = None


def set_engine(engine: Engine) -> None:
    global _engine
    _engine = engine


def get_set_engine() -> Engine | None:
    return _engine


def get_engine() -> Engine:
    url = os.environ.get("DATABASE_URL", DATABASE_URL_DEFAULT)
    return create_engine(url)


def get_sessionmaker(engine: Engine | None = None) -> sessionmaker[Session]:
    engine = engine or get_engine()
    return sessionmaker(bind=engine)


@contextmanager
def get_session(engine: Engine | None = None) -> Generator[Session]:
    session_factory = get_sessionmaker(engine)
    session = session_factory()
    try:
        yield session
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()