from contextlib import contextmanager
from pathlib import Path

from sqlalchemy import create_engine
from sqlalchemy.orm import DeclarativeBase, sessionmaker

from app.config import get_settings


class Base(DeclarativeBase):
    pass


_engine = None
_SessionLocal = None


def init_db(url: str | None = None) -> None:
    """Создаёт engine и таблицы. Вызывается при старте приложения (и в тестах)."""
    global _engine, _SessionLocal
    url = url or get_settings().database_url
    if url.startswith("sqlite:///") and ":memory:" not in url:
        Path(url.removeprefix("sqlite:///")).parent.mkdir(parents=True, exist_ok=True)
    _engine = create_engine(url, connect_args={"check_same_thread": False})
    _SessionLocal = sessionmaker(_engine, expire_on_commit=False)
    from app import models  # noqa: F401  (регистрация моделей)

    Base.metadata.create_all(_engine)
    _migrate()


def _migrate() -> None:
    """create_all не добавляет колонки в существующие таблицы — докидываем вручную."""
    with _engine.begin() as conn:
        cols = {r[1] for r in conn.exec_driver_sql("PRAGMA table_info(devices)")}
        if "use_tls" not in cols:
            conn.exec_driver_sql("ALTER TABLE devices ADD COLUMN use_tls BOOLEAN NOT NULL DEFAULT 1")
        if "group_id" not in cols:
            conn.exec_driver_sql("ALTER TABLE devices ADD COLUMN group_id INTEGER")
        if "note" not in cols:
            conn.exec_driver_sql("ALTER TABLE devices ADD COLUMN note TEXT")


@contextmanager
def session_scope():
    """Короткая транзакция: commit при успехе, rollback при ошибке."""
    session = _SessionLocal()
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()