from collections.abc import Generator

from sqlalchemy import create_engine, inspect, text
from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker

from app.config import get_settings


class Base(DeclarativeBase):
    pass


settings = get_settings()
connect_args = {"check_same_thread": False} if settings.database_url.startswith("sqlite") else {}
engine = create_engine(settings.database_url, pool_pre_ping=True, connect_args=connect_args)
SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False)


def get_db() -> Generator[Session, None, None]:
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


_LIGHTWEIGHT_ALTERS: list[tuple[str, str, str]] = [
    ("attachments", "remote_url", "VARCHAR(1000)"),
    ("users", "unit_id", "INTEGER"),
    ("units", "parent_id", "INTEGER"),
    ("units", "category", "VARCHAR(20) DEFAULT 'company'"),
]


def _apply_lightweight_alters() -> None:
    inspector = inspect(engine)
    for table, column, ddl_type in _LIGHTWEIGHT_ALTERS:
        if not inspector.has_table(table):
            continue
        existing = {c["name"] for c in inspector.get_columns(table)}
        if column in existing:
            continue
        with engine.begin() as conn:
            conn.execute(text(f'ALTER TABLE {table} ADD COLUMN {column} {ddl_type}'))


def init_db() -> None:
    from app.models import all_models  # noqa: F401

    Base.metadata.create_all(bind=engine)
    _apply_lightweight_alters()